How to implement a custom loss without rewriting box_predictor

[alexander-soare]

My goal is to apply label smoothing to the classification loss. So I write a new ROIHeads class:

@ROI_HEADS_REGISTRY.register()
class CustomROIHeads(StandardROIHeads):
    def _forward_box(self, features: Dict[str, torch.Tensor], proposals: List[Instances]):
        """
        Forward logic of the box prediction branch. If `self.train_on_pred_boxes is True`,
            the function puts predicted boxes in the `proposal_boxes` field of `proposals` argument.
        Args:
            features (dict[str, Tensor]): mapping from feature map names to tensor.
                Same as in :meth:`ROIHeads.forward`.
            proposals (list[Instances]): the per-image object proposals with
                their matching ground truth.
                Each has fields "proposal_boxes", and "objectness_logits",
                "gt_classes", "gt_boxes".
        Returns:
            In training, a dict of losses.
            In inference, a list of `Instances`, the predicted instances.
        """
        features = [features[f] for f in self.box_in_features]
        box_features = self.box_pooler(features, [x.proposal_boxes for x in proposals])
        box_features = self.box_head(box_features)
        predictions = self.box_predictor(box_features)
        del box_features

        if self.training:
            losses = self.box_predictor.losses(predictions, proposals)
            # proposals is modified in-place below, so losses must be computed first.
            if self.train_on_pred_boxes:
                with torch.no_grad():
                    pred_boxes = self.box_predictor.predict_boxes_for_gt_classes(
                        predictions, proposals
                    )
                    for proposals_per_image, pred_boxes_per_image in zip(proposals, pred_boxes):
                        proposals_per_image.proposal_boxes = Boxes(pred_boxes_per_image)
            return losses
        else:
            pred_instances, _ = self.box_predictor.inference(predictions, proposals)
            return pred_instances

Actually I haven't done anything new there, but I'm trying to figure out how to replace self.box_predictor. In the docs, this example is given:

@ROI_HEADS_REGISTRY.register()
class MyStandardROIHeads(StandardROIHeads):
  def __init__(self, cfg, input_shape):
    super().__init__(cfg, input_shape,
                     box_predictor=MyRCNNOutput(...))

But it's not clear to me how I can make use of the existing MyFastRCNNOutputLayers class. Because I'd like to replace the loss method there, something like this:

class MyFastRCNNOutputLayers(FastRCNNOutputLayers):
    def losses(self, predictions, proposals):
        """
        Args:
            predictions: return values of :meth:`forward()`.
            proposals (list[Instances]): proposals that match the features that were used
                to compute predictions. The fields ``proposal_boxes``, ``gt_boxes``,
                ``gt_classes`` are expected.
        Returns:
            Dict[str, Tensor]: dict of losses
        """

        scores, proposal_deltas = predictions

        # parse classification outputs
        gt_classes = (
            cat([p.gt_classes for p in proposals], dim=0) if len(proposals) else torch.empty(0)
        )
        _log_classification_stats(scores, gt_classes)

        # parse box regression outputs
        if len(proposals):
            proposal_boxes = cat([p.proposal_boxes.tensor for p in proposals], dim=0)  # Nx4
            assert not proposal_boxes.requires_grad, "Proposals should not require gradients!"
            # If "gt_boxes" does not exist, the proposals must be all negative and
            # should not be included in regression loss computation.
            # Here we just use proposal_boxes as an arbitrary placeholder because its
            # value won't be used in self.box_reg_loss().
            gt_boxes = cat(
                [(p.gt_boxes if p.has("gt_boxes") else p.proposal_boxes).tensor for p in proposals],
                dim=0,
            )
        else:
            proposal_boxes = gt_boxes = torch.empty((0, 4), device=proposal_deltas.device)

        losses = {
            # CUSTOM MODIFICATION, APPLY LABEL SMOOTHING TO GT_CLASSES
            "loss_cls": cross_entropy(scores, gt_classes, reduction="mean"),
            "loss_box_reg": self.box_reg_loss(
                proposal_boxes, gt_boxes, proposal_deltas, gt_classes
            ),
        }
        return {k: v * self.loss_weight.get(k, 1.0) for k, v in losses.items()}

The problem is, I have no idea how and where to init this. StandardROIHeads does not init the box_predictor but rather, it's passed in as an argument. And there are multiple instances of StandardROIHeads in an FPN framework, which only complicates things more, because it means I probably need to do the initialisation somewhere upstream.

I realised that maybe if I knew what was supposed to go in the ellipsis from the example in the docs:

@ROI_HEADS_REGISTRY.register()
class MyStandardROIHeads(StandardROIHeads):
  def __init__(self, cfg, input_shape):
    super().__init__(cfg, input_shape,
                     box_predictor=MyRCNNOutput(...))

... I would be able to solve it. But the init signature of FastRCNNOutputLayers requires arguments which I don't know where to obtain.

[ppwwyyxx]

I realised that maybe if I knew what was supposed to go in the ellipsis from the example in the docs:

detectron2/detectron2/modeling/roi_heads/fast_rcnn.py

Lines 398 to 413 in 6d25060

	@classmethod
	def from_config(cls, cfg, input_shape):
	return {
	"input_shape": input_shape,
	"box2box_transform": Box2BoxTransform(weights=cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_WEIGHTS),
	# fmt: off
	"num_classes" : cfg.MODEL.ROI_HEADS.NUM_CLASSES,
	"cls_agnostic_bbox_reg" : cfg.MODEL.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG,
	"smooth_l1_beta" : cfg.MODEL.ROI_BOX_HEAD.SMOOTH_L1_BETA,
	"test_score_thresh" : cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST,
	"test_nms_thresh" : cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST,
	"test_topk_per_image" : cfg.TEST.DETECTIONS_PER_IMAGE,
	"box_reg_loss_type" : cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_TYPE,
	"loss_weight" : {"loss_box_reg": cfg.MODEL.ROI_BOX_HEAD.BBOX_REG_LOSS_WEIGHT},
	# fmt: on
	}

shows what the arguments could look like if initialized from a config file. Though eventually it's up to you to provide appropriate arguments for the new use case.

[alexander-soare]

Thanks so much for picking this up. The thing is, all I need to change is the way the loss is calculated within the losses method of FastRCNNOutputLayers, so the arguments should be exactly the same. The challenge I'm facing though, is that box_predictor (which is the instance of FastRCNNOutputLayers) is instantiated somewhere other than StandardROIHeads. It's actually fed into StandardROIHeads as an argument to its constructor. Even if I ignore that fact, and try to instantiate a new version of box_predictor in the constructor of my subclass CustomROIHeads, I'll have to track down what all the arguments to the original box_predictor were in order to match the framework (and there are several because I'm dealing with FPN, so I need to know which layer I'm in).

Perhaps there's a better way to do this, but:

• As far as I'm aware, there is no way to "register" the equivalent of FastRCNNOutputLayers.
• I'm not able to track down where FastRCNNOutputLayers is instantiated to tackle the problem there.

[ppwwyyxx]

Here

detectron2/detectron2/modeling/roi_heads/roi_heads.py

Line 636 in 6d25060

box_predictor = FastRCNNOutputLayers(cfg, box_head.output_shape)

for StandardROIHeads

[alexander-soare]

Oh... I looked everywhere but here. Thank you.

On a side note while I have your attention, thanks to you and team for putting this awesome framework together. I know it's tough to build build something which is so easy to use and also customisable. Great job!

[deepaksinghcv]

Hey @alexander-soare, @ppwwyyxx
I have a similar requirement where I want to add a new custom head with its own loss function along with the losses of FastRCNNOutputLayers.
I have added the custom head, the custom loss; I can even notice the updated model definition when the model definition is being logged,the losses are being calculated only for the first iteration.
After the first iteration I'm facing the following error:

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by (1) passing the keyword argument `find_unused_parameters=True` to `torch.nn.parallel.DistributedDataParallel`; (2) making sure all `forward` function outputs participate in calculating loss. If you already have done the above two steps, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's `forward` function. Please include the loss function and the structure of the return value of `forward` of your module when reporting this issue (e.g. list, dict, iterable).

Any idea what's causing this issue.