06_pytorch_transfer_learning

[adi018]

Hello,
I am working on an object detection application using a pre-trained FasterRCNN network on COCO dataset. For my use case, I want my network to detect only four classes i.e. A, B, C and D.

I have downloaded the model and frozed the weights of the backbone and RPN of the model. Now when I initialize the ROI head should I use num_classes = 4 or num_classes = 5 . According to https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html or https://discuss.pytorch.org/t/finetuning-fasterrcnn-for-multi-class-examples/92707/3, I should set num_classes = 5 . Could you please give an explanation? Thanks!

Please find the model summary below:

=======================================================================================================================================
Layer (type (var_name))                                 Input Shape          Output Shape         Param #              Trainable
=======================================================================================================================================
FasterRCNN (FasterRCNN)                                 [1, 3, 512, 512]     --                   --                   Partial
├─GeneralizedRCNNTransform (transform)                  [1, 3, 512, 512]     --                   --                   --
├─BackboneWithFPN (backbone)                            [1, 3, 800, 800]     [1, 256, 13, 13]     --                   False
│    └─IntermediateLayerGetter (body)                   [1, 3, 800, 800]     [1, 2048, 25, 25]    --                   False
│    │    └─Conv2d (conv1)                              [1, 3, 800, 800]     [1, 64, 400, 400]    (9,408)              False
│    │    └─FrozenBatchNorm2d (bn1)                     [1, 64, 400, 400]    [1, 64, 400, 400]    --                   --
│    │    └─ReLU (relu)                                 [1, 64, 400, 400]    [1, 64, 400, 400]    --                   --
│    │    └─MaxPool2d (maxpool)                         [1, 64, 400, 400]    [1, 64, 200, 200]    --                   --
│    │    └─Sequential (layer1)                         [1, 64, 200, 200]    [1, 256, 200, 200]   (212,992)            False
│    │    └─Sequential (layer2)                         [1, 256, 200, 200]   [1, 512, 100, 100]   (1,212,416)          False
│    │    └─Sequential (layer3)                         [1, 512, 100, 100]   [1, 1024, 50, 50]    (7,077,888)          False
│    │    └─Sequential (layer4)                         [1, 1024, 50, 50]    [1, 2048, 25, 25]    (14,942,208)         False
│    └─FeaturePyramidNetwork (fpn)                      --                   [1, 256, 13, 13]     --                   False
│    │    └─ModuleList (inner_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (layer_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (inner_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (layer_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (inner_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (layer_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (inner_blocks)                   --                   --                   (recursive)          False
│    │    └─ModuleList (layer_blocks)                   --                   --                   (recursive)          False
│    │    └─LastLevelMaxPool (extra_blocks)             [1, 256, 200, 200]   [1, 256, 200, 200]   --                   --
├─RegionProposalNetwork (rpn)                           --                   [1000, 4]            --                   False
│    └─RPNHead (head)                                   [1, 256, 200, 200]   [1, 3, 200, 200]     --                   False
│    │    └─Sequential (conv)                           [1, 256, 200, 200]   [1, 256, 200, 200]   (590,080)            False
│    │    └─Conv2d (cls_logits)                         [1, 256, 200, 200]   [1, 3, 200, 200]     (771)                False
│    │    └─Conv2d (bbox_pred)                          [1, 256, 200, 200]   [1, 12, 200, 200]    (3,084)              False
│    │    └─Sequential (conv)                           [1, 256, 100, 100]   [1, 256, 100, 100]   (recursive)          False
│    │    └─Conv2d (cls_logits)                         [1, 256, 100, 100]   [1, 3, 100, 100]     (recursive)          False
│    │    └─Conv2d (bbox_pred)                          [1, 256, 100, 100]   [1, 12, 100, 100]    (recursive)          False
│    │    └─Sequential (conv)                           [1, 256, 50, 50]     [1, 256, 50, 50]     (recursive)          False
│    │    └─Conv2d (cls_logits)                         [1, 256, 50, 50]     [1, 3, 50, 50]       (recursive)          False
│    │    └─Conv2d (bbox_pred)                          [1, 256, 50, 50]     [1, 12, 50, 50]      (recursive)          False
│    │    └─Sequential (conv)                           [1, 256, 25, 25]     [1, 256, 25, 25]     (recursive)          False
│    │    └─Conv2d (cls_logits)                         [1, 256, 25, 25]     [1, 3, 25, 25]       (recursive)          False
│    │    └─Conv2d (bbox_pred)                          [1, 256, 25, 25]     [1, 12, 25, 25]      (recursive)          False
│    │    └─Sequential (conv)                           [1, 256, 13, 13]     [1, 256, 13, 13]     (recursive)          False
│    │    └─Conv2d (cls_logits)                         [1, 256, 13, 13]     [1, 3, 13, 13]       (recursive)          False
│    │    └─Conv2d (bbox_pred)                          [1, 256, 13, 13]     [1, 12, 13, 13]      (recursive)          False
│    └─AnchorGenerator (anchor_generator)               --                   [159882, 4]          --                   --
├─RoIHeads (roi_heads)                                  --                   --                   --                   True
│    └─MultiScaleRoIAlign (box_roi_pool)                --                   [1000, 256, 7, 7]    --                   --
│    └─TwoMLPHead (box_head)                            [1000, 256, 7, 7]    [1000, 1024]         --                   True
│    │    └─Linear (fc6)                                [1000, 12544]        [1000, 1024]         12,846,080           True
│    │    └─Linear (fc7)                                [1000, 1024]         [1000, 1024]         1,049,600            True
│    └─FastRCNNPredictor (box_predictor)                [1000, 1024]         [1000, 5]            --                   True
│    │    └─Linear (cls_score)                          [1000, 1024]         [1000, 5]            5,125                True
│    │    └─Linear (bbox_pred)                          [1000, 1024]         [1000, 20]           20,500               True
=======================================================================================================================================
Total params: 41,314,536
Trainable params: 13,921,305
Non-trainable params: 27,393,231
Total mult-adds (G): 133.98
=======================================================================================================================================
Input size (MB): 3.15
Forward/backward pass size (MB): 1454.98
Params size (MB): 165.26
Estimated Total Size (MB): 1623.39
=======================================================================================================================================

[mrdbourke]

Hi @adi018,

I don't have much experience with the PyTorch pre-trained FasterRCNN but my guess is you'd want the heads to be the same number of classes as you've got (e.g. 4 for A, B, C, D).

Have you tried both combinations?

One model with 4 classes and one model with 5 classes?

That would be the first thing I try.