Batching Loss — StreamPETR Head

This page documents the batched loss computation used by StreamPETRHeadBatched. It builds on the same approach as the focal head but extends it to multiple decoder layers, keeping the computation batched over both samples and layers while returning per-layer losses.

The reference head is already partially batched over samples: within each decoder layer, it concatenates the per-sample targets and invokes each classification and bounding-box loss function once. The optimized head does not newly batch those loss-function calls over samples. Instead, it batches the surrounding per-sample target preparation, keeps variable-length ground truth in RaggedBatch objects, computes assignment cost matrices in batches, and computes bounding-box loss only for matched prediction/target pairs. The SciPy Hungarian solver still handles each layer/sample entry separately. The optimized head additionally invokes the classification and bounding-box loss functions across all decoder layers at once, including the DN losses; the reference processes one layer at a time.

You can find the implementation in the repository here:

• StreamPETRHeadBatched: demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

• StreamPETRHead (reference): demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head.py

• Helper classes:

  • HungarianAssigner3DBatched: demos/streampetr/projects/mmdet3d_plugin/core/bbox/assigners/hungarian_assigner_3d_batched.py

  • HungarianAssigner3D (non-batched reference): demos/streampetr/projects/mmdet3d_plugin/core/bbox/assigners/hungarian_assigner_3d.py

Note

The used HungarianAssigner3DBatched follows the same approach as the HungarianAssigner2DBatched, with some changes in the cost matrix computation (but still using the same approach to compute the matrices). For more details, please also see the source code of the HungarianAssigner3DBatched and the page on the matching implementation.

Overview and Goals

The more fully batched loss computation is based on the same approach as used in focal_head_batched.py. Please, see the notes in focal_head_batched.py first for a detailed explanation of the approach.

The original StreamPETR head already batches the final classification and bounding-box loss-function calls over samples within each decoder layer: it concatenates the per-sample targets and flattens the sample and query dimensions before invoking the loss functions. The changes here instead batch the surrounding target preparation and cost-matrix computation over samples and decoder layers, preserve the structured batch dimensions until the loss-function boundary, and compute bounding-box loss only for matched prediction/target pairs. The SciPy Hungarian solver itself still runs separately for each layer/sample entry. The changes also add batching of the loss-function calls across decoder layers, including the DN losses, which is not present in the original implementation.

Here, we extend the approach to the case of multiple decoder layers, where the loss is computed for each decoder layer (i.e. batched, but not reduced over the layers). The explanation here is focused on this extension. The extension is done as follows:

1. The GT data is duplicated for each decoder layer to match the shape of the predictions. This is similar to the original implementation. However, lists are used there. Here, we use RaggedBatch instances. After converting the GT data to RaggedBatch objects, an additional batch dimension is added, and the GT data is duplicated for each decoder layer, so that the additional batch dimension iterates over the decoder layers.

2. The method _compute_per_layer_loss() is used as a wrapper around the actual loss functions and allows for reduction of the loss over the samples while retaining the per-decoder-layer separation. Note that this is done to work around the fact that the loss functions can only reduce over all objects in the batch at once, but cannot return per-sample losses, i.e. they cannot reduce over the objects only, keeping the batch dimension(s) intact.

Also note that the loss computation includes the DN loss computation (dn_loss_batched()), which is not present for the focal head. The inputs for the DN loss computation are obtained from prepare_for_loss(), which returns the data with the first dimension being the decoder layers, so that the inputs can be directly used for batched loss computation.

The implementation of the loss computation is done in the loss() function. Please see the notes & doc-strings there and in the related functions called from there for more details.

Below, we discuss the implementation in detail, focusing on the parts of the implementation that differ from or extend the focal head case.

Top-level Loss Computation

Pre-processing

Initial validation and GT preparation (combining samples as ragged) mirrors the focal head.

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Pre-processing (`loss()`) =====
        nvtx_wrp.range_push("StreamPETRHeadBatched.loss")
        stopwatch.start_meas("StreamPETRHeadBatched.loss")
        assert gt_bboxes_ignore is None, \
            f'{self.__class__.__name__} only supports ' \
            f'for gt_bboxes_ignore setting to None.'
        
        loss_dict = dict()

        all_cls_scores = preds_dicts['all_cls_scores']
        all_bbox_preds = preds_dicts['all_bbox_preds']

        num_dec_layers = len(all_cls_scores)
        device = gt_labels_list[0].device
        #@NOTE
        # Note that `gt_bboxes_list` is on the CPU, so we need to transfer it to the GPU.
        # However, this is done when calling `combine_data()` below. Typically, it is more efficient to
        # do the combining of samples onto batched on the CPU, so we do not not transfer the data to the GPU
        # before calling `combine_data()`.
        gt_bboxes_list = [
            torch.cat((gt_bboxes.gravity_center, gt_bboxes.tensor[:, 3:]), dim=1)
            for gt_bboxes in gt_bboxes_list
        ]
        gt_bboxes = bhl.combine_data(gt_bboxes_list, device=device, flatten_batch_dims=True)

        #@NOTE
        # Using `other_with_same_sample_sizes=gt_bboxes` to share the memory for the 
        # `sample_sizes` and `mask` tensors of the two RaggedBatch instances.
        gt_labels = bhl.combine_data(gt_labels_list, other_with_same_sample_sizes=gt_bboxes, flatten_batch_dims=True)

Prepare Per-layer GT

A key difference comes next where GT is duplicated across decoder layers to match prediction shapes. Introduce a new batch dimension for decoder layers and repeat GT along it to match prediction layout:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Prepare per-layer GT =====
        #@NOTE
        # Duplicating the GT for each layer in order to match the shape of the predictions, which
        # are also per-layer (but in contrast to the GT, the predictions differ for each layer). 
        # This is done by introducing a new batch dimension (iterating over decoder layers) and repeating the GT data 
        # for each decoder layer along that dimension.
        gt_bboxes = gt_bboxes.unsqueeze_batch_dim(0)
        gt_bboxes = gt_bboxes.repeat_samples(num_dec_layers, batch_dim=0)
        gt_labels = gt_labels.unsqueeze_batch_dim(0)
        gt_labels = gt_labels.repeat_samples(num_dec_layers, batch_dim=0)
        

Compute Batched Loss over Decoder Layers

Delegate to loss_batched to compute per-layer losses in a batched manner:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Compute batched loss over decoder layers =====
        #@NOTE 
        # Compute the loss for each decoder layer. Note that in the original implementation, the loss is 
        # computed for each decoder layer separately. Here, we batch the computation over all decoder layers. 
        # 
        # Batching over the individual samples in the batch is done both in the original and our 
        # implementations (although the approaches differ).
        losses_cls, losses_bbox = self.loss_batched(all_cls_scores, all_bbox_preds, gt_bboxes, gt_labels)

        # loss from the last decoder layer

Assemble Results

Collect losses into the output dictionary and finalize timers:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Assemble results =====
        loss_dict['loss_cls'] = losses_cls[-1]
        loss_dict['loss_bbox'] = losses_bbox[-1]

        # loss from other decoder layers
        for ndl, (loss_cls_i, loss_bbox_i) in enumerate(zip(losses_cls[:-1], losses_bbox[:-1])):
            loss_dict[f'd{ndl}.loss_cls'] = loss_cls_i
            loss_dict[f'd{ndl}.loss_bbox'] = loss_bbox_i

DN Loss (Batched over Layers) & Build Loss Dictionary

When DN is enabled, DN losses are also computed batched over decoder layers and merged into the dictionary. Note that in the original implementation, the dn_loss_single() function is called for each decoder layer separately. Here, we implement dn_loss_batched() to compute the DN loss batched over all decoder layers.

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== DN Loss (batched over layers) =====
        #@NOTE: Compute the DN loss. 
        if preds_dicts['dn_mask_dict'] is not None:
            known_labels, known_bboxs, output_known_class, output_known_coord, num_tgt = \
                self.prepare_for_loss(preds_dicts['dn_mask_dict'])
            #@NOTE
            # Duplicating the GT for each decoder layer to match the shape of the predictions, which
            # are also per-layer (but differ for each layer in contrast to the GT). In the original 
            # implementation, lists are used instead of `repeat()`, as the GT is not batched over the decoder 
            # layers and instead, the loss is computed for each layer individually. Here, the computation is 
            # batched, so we need the GT data to be repeated for each decoder layer.
            #   
            # Also, note that here, we use tensors and not `RaggedBatch` instances for the GT data, as the 
            # GT is of uniform size.
            all_known_bboxs = known_bboxs.unsqueeze(0).repeat((num_dec_layers, 1, 1))
            all_known_labels = known_labels.unsqueeze(0).repeat((num_dec_layers, 1))
            all_num_tgt = torch.full(
                (num_dec_layers,), num_tgt, dtype=torch.float32, device=all_known_labels.device
            )

            #@NOTE: Computing the DN loss batched over the decoder layers.
            dn_losses_cls, dn_losses_bbox = self.dn_loss_batched(
                output_known_class, output_known_coord, all_known_bboxs, all_known_labels, all_num_tgt
            )

            loss_dict['dn_loss_cls'] = dn_losses_cls[-1]
            loss_dict['dn_loss_bbox'] = dn_losses_bbox[-1]
            #@NOTE
            # Here, we split the loss for each decoder layer and store them in the loss dictionary
            # This results in a dictionary with the same structure as in the original implementation.
            for ndl, (dn_l_cls_i, dn_l_bbox_i) in enumerate(zip(dn_losses_cls[:-1], dn_losses_bbox[:-1])):
                loss_dict[f'd{ndl}.dn_loss_cls'] = dn_l_cls_i
                loss_dict[f'd{ndl}.dn_loss_bbox'] = dn_l_bbox_i
        elif self.with_dn:
            #@NOTE
            # In this case, the DN loss computation is the same as the loss computation, so we
            # re-use the results from the loss computation. Note that this is different from the original 
            # implementation, where the loss is re-computed. 
            # 
            # This branch is not used in the configurations used for the runtime experiments and
            # is only added for feature completeness w.r.t. the original implementation.
            loss_dict['dn_loss_cls'] = loss_dict['loss_cls'].detach().clone()
            loss_dict['dn_loss_bbox'] = loss_dict['loss_bbox'].detach().clone()
            for ndl in range(num_dec_layers - 1):
                loss_dict[f'd{ndl}.dn_loss_cls'] = loss_dict[f'd{ndl}.loss_cls'].detach().clone()
                loss_dict[f'd{ndl}.dn_loss_bbox'] = loss_dict[f'd{ndl}.loss_bbox'].detach().clone()

Finalize

The final step is to return the loss dictionary.

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Finalize (`loss()`) =====
        nvtx_wrp.range_pop("StreamPETRHeadBatched.loss")
        stopwatch.end_meas("StreamPETRHeadBatched.loss")

        return loss_dict

Inside loss_batched()

Pre-processing

Shape extraction and entry-point for the per-layer batched computation:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Pre-processing (`loss_batched()`) =====
        nvtx_wrp.range_push("StreamPETRHeadBatched.loss_batched")
        num_dec_layers = cls_scores.size(0)
        num_samples_in_batch = cls_scores.size(1)
        num_preds = cls_scores.size(2)
        

Perform the Matching

Call the 3D batched assigner to obtain matched indices. Note that this is analogous to the focal head case.

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Perform the matching =====
        #@NOTE
        # The following is the same as the case of `FocalHeadBatched.loss_batched()`, i.e. we
        # call the assign() method directly and obtain the matched indices.
        # 
        # This is the same as in the focal head case. However, the assigner itself if a different class
        # (i.e. `HungarianAssigner3DBatched` instead of `HungarianAssigner2DBatched`). Also, the data
        # has multiple batch dimensions (which are handled automatically by the assigner implementation).
        matched_pred_indices, matched_gt_indices = self.assigner.assign(bbox_preds, cls_scores, gt_bboxes,
                                                                        gt_labels, None, self.match_costs, 
                                                                        self.match_with_velo)
        

Per-decoder-layer Counts and Average Factors

Compute per-layer number of positives/negatives and average factors. In contrast to the focal head case, the number of positives/negatives is computed for each decoder layer separately.

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Per-layer counts and avg factors =====
        #@NOTE: `per_sample_nums_pos` will have the same shape as the batch dimensions.
        per_sample_nums_pos = matched_pred_indices.sample_sizes
        #@NOTE
        # Here, we sum over the samples (the second batch dimension, i.e. `dim==1`) to obtain the total 
        # number of positive samples for each decoder layer. Note that the first batch dimension, i.e. 
        # `dim==0`, is over the decoder layers, and we keep that dimension.
        #   
        # For comparison: In the original implementation, the number of positive samples is 
        # computed for each decoder layer separately (in the `get_targets()` function).
        per_layer_nums_pos = per_sample_nums_pos.sum(dim=1)
        per_layer_nums_neg = num_preds * num_samples_in_batch - per_layer_nums_pos

        #@NOTE
        # The following is the batched version of the per-layer normalization factor computation (for the 
        # classification loss). This is equivalent to the (per-sample) computation in the original 
        # implementation.
        per_layer_cls_avg_factors = per_layer_nums_pos * 1.0 + \
            per_layer_nums_neg * self.bg_cls_weight
        if self.sync_cls_avg_factor:
            per_layer_cls_avg_factors = reduce_mean(per_layer_cls_avg_factors)
        #@NOTE 
        # Make sure we don't have any division by zero (similar to original implementation).
        # This is equivalent to the original implementation, but done here for tensor elements
        # instead of for a scalar value.
        per_layer_cls_avg_factors[per_layer_cls_avg_factors < 1.0] = 1.0

Align Labels and Align Matched Boxes

Map GT labels to predictions (for classification) and align matched GT/prediction boxes (for bounding box regression loss):

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Align labels and align matched boxes =====
        nvtx_wrp.range_push("Indexing [StreamPETRHeadBatched.loss_batched]")
        #@NOTE:
        # Align/map back the GT labels to the predicted labels as is done for losses which 
        # are computed for all predictions, not just the matched ones. This is analogous to the focal head 
        # classification loss. Please have a look at the notes in `focal_head_batched.py` for a more detailed 
        # description. The main difference is that here, we have multiple batch dimensions. However, this is 
        # handled automatically by `bhl.batched_index_mapping()`.
        gt_labels_mapped_to_pred = torch.full(
            (num_dec_layers, num_samples_in_batch, num_preds), 
            fill_value=self.cls_out_channels, 
            dtype=gt_labels.dtype, device=gt_labels.device
        )
        gt_labels_mapped_to_pred = bhl.batched_index_mapping(
            gt_labels, matched_gt_indices, matched_pred_indices, gt_labels_mapped_to_pred
        )

        #@NOTE:
        # Get the aligned (i.e. indexed using indices obtained from the matching) versions of the 
        # bounding boxes. Again, this is analogous to the focal head case (except for the additional batch 
        # dimensions). Please have a look at the notes in `focal_head_batched.py` for a more detailed 
        # description.
        gt_bboxes_matched = bhl.batched_indexing_access(gt_bboxes, matched_gt_indices)
        pred_bboxes_matched = bhl.batched_indexing_access(bbox_preds, matched_pred_indices)
        nvtx_wrp.range_pop()

        label_weights = torch.ones(
            (num_dec_layers, num_samples_in_batch, num_preds), 
            dtype=torch.float32, device = gt_labels.device
        )

Losses: Classification

Compute classification loss per layer using the wrapper to reduce over samples while keeping layers separate.

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Losses: Classification (per-layer wrapper) =====
        nvtx_wrp.range_push("loss_cls computation[StreamPETRHeadBatched.loss_batched]")
        #@NOTE
        # The method `_compute_per_layer_loss()` is a wrapper to the actual loss computation and enables
        # batching over both samples and decoder layers, with the possibility to compute the loss for each
        # layer separately, while combining (i.e. summing or computing the average) the loss over the samples.
        # This is necessary as the loss function does not differentiate between the entries in a batch. 
        # Using the loss function directly, it would be not possible to compute the per-layer loss separately.
        # Please refer to the documentation of `_compute_per_layer_loss()` and the comments
        # in that function for more details.
        # 
        # Furthermore, note that the classification loss is computed for all predictions, not just the matched 
        # ones (also see the discussion of the classification loss in `focal_head_batched.py`).
        loss_cls = self._compute_per_layer_loss(
            self.loss_cls, gt_labels_mapped_to_pred, cls_scores, label_weights, per_layer_cls_avg_factors, 
            num_batch_dims=2
        )

        nvtx_wrp.range_pop()
        loss_cls = torch.nan_to_num(loss_cls)

Synchronize Positives

Synchronize positive counts across devices:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Synchronize positives =====
        #@NOTE: reduce_mean is used to synchronize over all devices, as in the original implementation.
        per_layer_nums_pos = per_layer_nums_pos.type(loss_cls.dtype)
        per_layer_nums_pos = reduce_mean(per_layer_nums_pos)
        per_layer_nums_pos[per_layer_nums_pos < 1.0] = 1.0

Prepare Bounding Box Targets & Masks

Next, we prepare the bounding box targets and masks:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Prepare bbox targets & masks =====
        gt_bboxes_matched_normalized = gt_bboxes_matched.apply(
            lambda bboxs: normalize_bbox(bboxs, self.pc_range)
        )

        isnotnan = torch.isfinite(gt_bboxes_matched_normalized.tensor).all(dim=-1, keepdim=True)
        
        #@NOTE
        # Get the existence weights. In the original implementation, the weights are returned 
        # by the `get_targets()` function (also see notes in `FocalHeadBatched.loss_batched()`). Note that
        # the weights here do not exactly correspond to the weights in the original implementation, as the 
        # existence weights here are computed for the data tensor (i.e. `tensor`) of the `RaggedBatch` 
        # instance (containing filler values to ensure uniform shape), while in the original implementation, 
        # the weights are computed for a tensor with the same shape as the predictions.
        bbox_existence_weights_matched = gt_bboxes_matched_normalized.get_existence_weights()

        #@NOTE 
        # The code weights are per-element weights for the individual elements which describe a single
        # bounding box. Therefore, they can be directly multiplied with the existence weights without any
        # adaptations regarding the batch shape, as long as the size of the inner-most dimension of the 
        # existence weights corresponds to the bounding box descriptors.
        bbox_weights_matched = bbox_existence_weights_matched * self.code_weights

        #@NOTE
        # The handling of non-finite values differs from the original implementation. In the original code,
        # boolean indexing with `isnotnan` as a mask is used to filter out non-finite values. However, 
        # this would lead to flattening of the batch dimensions and the dimension iterating over individual 
        # objects inside a sample. Here, we use `apply_mask_to_tensor()` to zero out non-finite values while 
        # preserving the current shape of the data.
        # This approach:
        #   1. Maintains the batch dimensions for efficient batched processing
        #   2. Properly handles `NaN` values (unlike simple multiplication, where `NaN * 0 = NaN`)
        # 
        # Note that our approach changes elements containing `NaN` values to filler (0) values, while the 
        # original implementation removes the `NaN` elements instead (using boolean indexing).
        gt_bboxes_matched_normalized = gt_bboxes_matched_normalized.apply(
            lambda tensor: bhl.apply_mask_to_tensor(tensor, isnotnan)
        )

        #@NOTE 
        # Here, we can multiply instead of using `apply_mask_to_tensor()` as the predictions and weights do 
        # not contain `NaN` values.
        pred_bboxes_matched = pred_bboxes_matched.apply(lambda bboxs: bboxs * isnotnan)
        bbox_weights_matched = bbox_weights_matched * isnotnan

Losses: Bounding Box Overlap

Compute bbox regression loss per layer using the wrapper, operating only on matched predictions:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Losses: BBox (per-layer wrapper) =====
        nvtx_wrp.range_push("loss_bbox computation[StreamPETRHeadBatched.loss_batched]")
        #@NOTE
        # The wrapper method `_compute_per_layer_loss()` is used again here.
        # 
        # In contrast to the classification loss, the bounding box loss is computed only for matched 
        # predictions. See notes in `focal_head_batched.py` for a discussion of how the input data differs 
        # between the two cases.
        loss_bbox = self._compute_per_layer_loss(
            self.loss_bbox, gt_bboxes_matched_normalized.tensor, pred_bboxes_matched.tensor, 
            bbox_weights_matched, per_layer_nums_pos, num_batch_dims=2
        )

        loss_bbox = torch.nan_to_num(loss_bbox)
        nvtx_wrp.range_pop()

Finalize

End of per-layer loss computation:

demos/streampetr/projects/mmdet3d_plugin/models/dense_heads/streampetr_head_batched.py

        # ===== Finalize (`loss_batched()`) =====
        nvtx_wrp.range_pop("StreamPETRHeadBatched.loss_batched")
        return loss_cls, loss_bbox

Other Notable Implementations Changes

In this section, we discussed the changes needed to implement the batched loss computation for the StreamPETR head.

In addition to the changes discussed above, there are some other notable changes to the implementation. Please refer to the source code (streampetr_head_batched.py, see beginning of this page) for more details. Changes not discussed above in detail are:

• The dn_loss_batched() method is used to compute the DN loss batched over all decoder layers. As the GT data is of uniform size, there is not need to use the RaggedBatch data format. However, the implementation is similar to loss_batched() in that the GT data is repeated for the individual decoder layers, and the loss is computed in a batched manner.

• The _compute_per_layer_loss() method is used to compute the loss for the individual decoder layers. This is achieved by disabling the reduction in the called loss function, and performing the reduction after the call in _compute_per_layer_loss() instead. In this way, we can decide over which dimensions to reduce, allowing us to keep the batch dimension corresponding to the decoder layers intact, while reducing over the individual samples. Note that the type of reduction can still be configured for the loss function, and the configuration will be read and applied accordingly in the _compute_per_layer_loss() method.