When adding dropout layers with class inheritance, they are always inserted at the end of the architecture (while that was not programmed!)

[pieterblok]

Hi guys, sorry for this noob question (and taking so much place, due to the code blocks). I promise it reads fast!

I followed the tutorial https://detectron2.readthedocs.io/en/latest/tutorials/write-models.html to rewrite a part of Faster R-CNN.

Specially I want to add two dropout layers before the linear layers in FastRCNNConvFCHead (./detectron2/modeling/roi_heads/box_head.py).

With this code its works perfect (it's just a copy-paste of the standard FastRCNNConvFCHead with two additional dropout-layers before the linear layers):

@ROI_BOX_HEAD_REGISTRY.register()
class FastRCNNConvFCHeadDropout(nn.Sequential):
    """
    A head with several 3x3 conv layers (each followed by norm & relu) and then
    several fc layers (each followed by relu).
    """

    @configurable
    def __init__(
        self, input_shape: ShapeSpec, *, conv_dims: List[int], fc_dims: List[int], conv_norm=""
    ):
        """
        NOTE: this interface is experimental.

        Args:
            input_shape (ShapeSpec): shape of the input feature.
            conv_dims (list[int]): the output dimensions of the conv layers
            fc_dims (list[int]): the output dimensions of the fc layers
            conv_norm (str or callable): normalization for the conv layers.
                See :func:`detectron2.layers.get_norm` for supported types.
        """
        super().__init__()
        assert len(conv_dims) + len(fc_dims) > 0

        self._output_size = (input_shape.channels, input_shape.height, input_shape.width)

        self.conv_norm_relus = []
        for k, conv_dim in enumerate(conv_dims):
            conv = Conv2d(
                self._output_size[0],
                conv_dim,
                kernel_size=3,
                padding=1,
                bias=not conv_norm,
                norm=get_norm(conv_norm, conv_dim),
                activation=nn.ReLU(),
            )
            self.add_module("conv{}".format(k + 1), conv)
            self.conv_norm_relus.append(conv)
            self._output_size = (conv_dim, self._output_size[1], self._output_size[2])

        self.fcs = []
        for k, fc_dim in enumerate(fc_dims):
            if k == 0:
                self.add_module("flatten", nn.Flatten())
            dropout = nn.Dropout(p=0.5)
            self.add_module("dropout{}".format(k + 1), dropout)
            fc = nn.Linear(int(np.prod(self._output_size)), fc_dim)
            self.add_module("fc{}".format(k + 1), fc)
            self.add_module("fc_relu{}".format(k + 1), nn.ReLU())
            self.fcs.append(fc)
            self._output_size = fc_dim

        for layer in self.conv_norm_relus:
            weight_init.c2_msra_fill(layer)
        for layer in self.fcs:
            weight_init.c2_xavier_fill(layer)

    @classmethod
    def from_config(cls, cfg, input_shape):
        num_conv = cfg.MODEL.ROI_BOX_HEAD.NUM_CONV
        conv_dim = cfg.MODEL.ROI_BOX_HEAD.CONV_DIM
        num_fc = cfg.MODEL.ROI_BOX_HEAD.NUM_FC
        fc_dim = cfg.MODEL.ROI_BOX_HEAD.FC_DIM
        return {
            "input_shape": input_shape,
            "conv_dims": [conv_dim] * num_conv,
            "fc_dims": [fc_dim] * num_fc,
            "conv_norm": cfg.MODEL.ROI_BOX_HEAD.NORM,
        }

    def forward(self, x):
        for layer in self:
            x = layer(x)
        return x

    @property
    @torch.jit.unused
    def output_shape(self):
        """
        Returns:
            ShapeSpec: the output feature shape
        """
        o = self._output_size
        if isinstance(o, int):
            return ShapeSpec(channels=o)
        else:
            return ShapeSpec(channels=o[0], height=o[1], width=o[2])

When I print the specific architecture, everything is perfect:

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x_dropout.yaml"))
model = build_model(cfg)
print(model.roi_heads.box_head)

FastRCNNConvFCHeadDropout(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (dropout1): Dropout(p=0.5, inplace=False)
  (fc1): Linear(in_features=12544, out_features=1024, bias=True)
  (fc_relu1): ReLU()
  (dropout2): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=1024, out_features=1024, bias=True)
  (fc_relu2): ReLU()
)

As you can understand, the code-block above is rather big for 2 additional lines of code, so I was thinking to add the dropout layers by class inheritance:

@ROI_BOX_HEAD_REGISTRY.register()
class FastRCNNConvFCHeadDropout(FastRCNNConvFCHead):

    def __init__(self, cfg, input_shape):
        super().__init__(cfg, input_shape)

        self.fcs = []
        num_fc = cfg.MODEL.ROI_BOX_HEAD.NUM_FC
        fc_dim = cfg.MODEL.ROI_BOX_HEAD.FC_DIM
        fc_dims = [fc_dim] * num_fc

        for k, fc_dim in enumerate(fc_dims):
            if k == 0:
                self.add_module("flatten", nn.Flatten())
            dropout = nn.Dropout(p=0.5)
            self.add_module("dropout{}".format(k + 1), dropout)
            fc = nn.Linear(int(np.prod(self._output_size)), fc_dim)
            self.add_module("fc{}".format(k + 1), fc)
            self.add_module("fc_relu{}".format(k + 1), nn.ReLU())
            self.fcs.append(fc)
            self._output_size = fc_dim

        for layer in self.conv_norm_relus:
            weight_init.c2_msra_fill(layer)
        for layer in self.fcs:
            weight_init.c2_xavier_fill(layer)
        

    def forward(self, x):
        for layer in self:
            x = layer(x)
        return x

When I print the specific architecture:

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x_dropout.yaml"))
model = build_model(cfg)
print(model.roi_heads.box_head)

FastRCNNConvFCHeadDropout(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (fc1): Linear(in_features=1024, out_features=1024, bias=True)
  (fc_relu1): ReLU()
  (fc2): Linear(in_features=1024, out_features=1024, bias=True)
  (fc_relu2): ReLU()
  (dropout1): Dropout(p=0.5, inplace=False)
  (dropout2): Dropout(p=0.5, inplace=False)

Unfortunately, the dropout layers are inserted at the end of the architecture, while this was not programmed. Maybe I missed something with the class inheritance. What am I doing wrong?

Is this maybe caused by the nn.Sequential (class inheritance of FastRCNNConvFCHead)? How can I properly insert the dropout layers before the linear layers by using class inheritance?

Thanks in advance, Pieter

[ppwwyyxx]

This is how Sequential works. see also pytorch/pytorch#43876

you might need to remove all the modules after the insertion location, and re-add them.

[pieterblok]

@ppwwyyxx thank you for your reply. I will first do it with a different function! The change in the nn.sequential is a bit too complex for me...