Adding a FasterRCNN + MobileNetV3 with & w/o FPN models.

Author: datumbox

test/expect/ModelTester.test_fasterrcnn_mobilenet_v3_large_expect.pkl

@@ -37,6 +37,8 @@ def get_available_video_models():
     'googlenet': lambda x: x.logits,
     'inception_v3': lambda x: x.logits,
     "fasterrcnn_resnet50_fpn": lambda x: x[1],
+    "fasterrcnn_mobilenet_v3_large": lambda x: x[1],
+    "fasterrcnn_mobilenet_v3_large_fpn": lambda x: x[1],
     "maskrcnn_resnet50_fpn": lambda x: x[1],
     "keypointrcnn_resnet50_fpn": lambda x: x[1],
     "retinanet_resnet50_fpn": lambda x: x[1],
@@ -105,6 +107,8 @@ def _test_detection_model(self, name, dev):
         if "retinanet" in name:
             # Reduce the default threshold to ensure the returned boxes are not empty.
             kwargs["score_thresh"] = 0.01
+        elif "fasterrcnn_mobilenet" in name:
+            kwargs["box_score_thresh"] = 0.02076
         model = models.detection.__dict__[name](num_classes=50, pretrained_backbone=False, **kwargs)
         model.eval().to(device=dev)
         input_shape = (3, 300, 300)

test/expect/ModelTester.test_fasterrcnn_mobilenet_v3_large_fpn_expect.pkl

@@ -97,14 +97,15 @@ def test_assign_targets_to_proposals(self):
         self.assertEqual(labels[0].dtype, torch.int64)
 
     def test_forward_negative_sample_frcnn(self):
-        model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
-            num_classes=2, min_size=100, max_size=100)
+        for name in ["fasterrcnn_resnet50_fpn", "fasterrcnn_mobilenet_v3_large", "fasterrcnn_mobilenet_v3_large_fpn"]:
+            model = torchvision.models.detection.__dict__[name](
+                num_classes=2, min_size=100, max_size=100)
 
-        images, targets = self._make_empty_sample()
-        loss_dict = model(images, targets)
+            images, targets = self._make_empty_sample()
+            loss_dict = model(images, targets)
 
-        self.assertEqual(loss_dict["loss_box_reg"], torch.tensor(0.))
-        self.assertEqual(loss_dict["loss_rpn_box_reg"], torch.tensor(0.))
+            self.assertEqual(loss_dict["loss_box_reg"], torch.tensor(0.))
+            self.assertEqual(loss_dict["loss_rpn_box_reg"], torch.tensor(0.))
 
     def test_forward_negative_sample_mrcnn(self):
         model = torchvision.models.detection.maskrcnn_resnet50_fpn(
@@ -130,7 +131,7 @@ def test_forward_negative_sample_krcnn(self):
 
     def test_forward_negative_sample_retinanet(self):
         model = torchvision.models.detection.retinanet_resnet50_fpn(
-            num_classes=2, min_size=100, max_size=100)
+            num_classes=2, min_size=100, max_size=100, pretrained_backbone=False)
 
         images, targets = self._make_empty_sample()
         loss_dict = model(images, targets)

test/test_models.py

@@ -4,6 +4,7 @@
 
 from torchvision.ops import misc as misc_nn_ops
 from .._utils import IntermediateLayerGetter
+from .. import mobilenet
 from .. import resnet
 
 
@@ -122,3 +123,50 @@ def _validate_trainable_layers(pretrained, trainable_backbone_layers, max_value,
         trainable_backbone_layers = default_value
     assert 0 <= trainable_backbone_layers <= max_value
     return trainable_backbone_layers
+
+
+def mobilenet_backbone(
+    backbone_name,
+    pretrained,
+    fpn,
+    norm_layer=misc_nn_ops.FrozenBatchNorm2d,
+    trainable_layers=2,
+    returned_layers=None,
+    extra_blocks=None
+):
+    backbone = mobilenet.__dict__[backbone_name](pretrained=pretrained, norm_layer=norm_layer).features
+
+    # Gather the indeces of blocks which are strided. These are the locations of C1, ..., Cn-1 blocks.
+    # The first and last blocks are always included because they are the C0 (conv1) and Cn.
+    stage_indeces = [0] + [i for i, b in enumerate(backbone) if getattr(b, "is_strided", False)] + [len(backbone) - 1]
+    num_stages = len(stage_indeces)
+
+    # find the index of the layer from which we wont freeze
+    assert 0 <= trainable_layers <= num_stages
+    freeze_before = num_stages if trainable_layers == 0 else stage_indeces[num_stages - trainable_layers]
+
+    # freeze layers only if pretrained backbone is used
+    for b in backbone[:freeze_before]:
+        for parameter in b.parameters():
+            parameter.requires_grad_(False)
+
+    out_channels = 256
+    if fpn:
+        if extra_blocks is None:
+            extra_blocks = LastLevelMaxPool()
+
+        if returned_layers is None:
+            returned_layers = [num_stages - 2, num_stages - 1]
+        assert min(returned_layers) >= 0 and max(returned_layers) < num_stages
+        return_layers = {f'{stage_indeces[k]}': str(v) for v, k in enumerate(returned_layers)}
+
+        in_channels_list = [backbone[stage_indeces[i]].out_channels for i in returned_layers]
+        return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks)
+    else:
+        m = nn.Sequential(
+            backbone,
+            # depthwise linear combination of channels to reduce their size
+            nn.Conv2d(backbone[-1].out_channels, out_channels, 1),
+        )
+        m.out_channels = out_channels
+        return m

test/test_models_detection_negative_samples.py

@@ -15,11 +15,11 @@
 from .rpn import RPNHead, RegionProposalNetwork
 from .roi_heads import RoIHeads
 from .transform import GeneralizedRCNNTransform
-from .backbone_utils import resnet_fpn_backbone, _validate_trainable_layers
+from .backbone_utils import resnet_fpn_backbone, _validate_trainable_layers, mobilenet_backbone
 
 
 __all__ = [
-    "FasterRCNN", "fasterrcnn_resnet50_fpn",
+    "FasterRCNN", "fasterrcnn_resnet50_fpn", "fasterrcnn_mobilenet_v3_large", "fasterrcnn_mobilenet_v3_large_fpn"
 ]
 
 
@@ -291,6 +291,8 @@ def forward(self, x):
 model_urls = {
     'fasterrcnn_resnet50_fpn_coco':
         'https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth',
+    'fasterrcnn_mobilenet_v3_large_coco': None,
+    'fasterrcnn_mobilenet_v3_large_fpn_coco': None,
 }
 
 
@@ -367,3 +369,83 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
         model.load_state_dict(state_dict)
         overwrite_eps(model, 0.0)
     return model
+
+
+def fasterrcnn_mobilenet_v3_large(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True,
+                                  trainable_backbone_layers=None, **kwargs):
+    """
+    Constructs a Faster R-CNN model with a MobileNetV3-Large backbone. It works similarly
+    to Faster R-CNN with ResNet-50 FPN backbone. See `fasterrcnn_resnet50_fpn` for more details.
+
+    Example::
+
+        >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large(pretrained=True)
+        >>> model.eval()
+        >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
+        >>> predictions = model(x)
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on COCO train2017
+        progress (bool): If True, displays a progress bar of the download to stderr
+        num_classes (int): number of output classes of the model (including the background)
+        pretrained_backbone (bool): If True, returns a model with backbone pre-trained on Imagenet
+        trainable_backbone_layers (int): number of trainable (not frozen) resnet layers starting from final block.
+            Valid values are between 0 and 6, with 6 meaning all backbone layers are trainable.
+    """
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 6, 3)
+
+    if pretrained:
+        pretrained_backbone = False
+    backbone = mobilenet_backbone("mobilenet_v3_large", pretrained_backbone, False,
+                                  trainable_layers=trainable_backbone_layers)
+
+    anchor_sizes = ((32, 64, 128, 256, 512), )
+    aspect_ratios = ((0.5, 1.0, 2.0), )
+
+    model = FasterRCNN(backbone, num_classes, rpn_anchor_generator=AnchorGenerator(anchor_sizes, aspect_ratios),
+                       **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls['fasterrcnn_mobilenet_v3_large_coco'], progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def fasterrcnn_mobilenet_v3_large_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True,
+                                      trainable_backbone_layers=None, **kwargs):
+    """
+    Constructs a Faster R-CNN model with a MobileNetV3-Large FPN backbone. It works similarly
+    to Faster R-CNN with ResNet-50 FPN backbone. See `fasterrcnn_resnet50_fpn` for more details.
+
+    Example::
+
+        >>> model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(pretrained=True)
+        >>> model.eval()
+        >>> x = [torch.rand(3, 300, 400), torch.rand(3, 500, 400)]
+        >>> predictions = model(x)
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on COCO train2017
+        progress (bool): If True, displays a progress bar of the download to stderr
+        num_classes (int): number of output classes of the model (including the background)
+        pretrained_backbone (bool): If True, returns a model with backbone pre-trained on Imagenet
+        trainable_backbone_layers (int): number of trainable (not frozen) resnet layers starting from final block.
+            Valid values are between 0 and 6, with 6 meaning all backbone layers are trainable.
+    """
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 6, 3)
+
+    if pretrained:
+        pretrained_backbone = False
+    backbone = mobilenet_backbone("mobilenet_v3_large", pretrained_backbone, True,
+                                  trainable_layers=trainable_backbone_layers)
+
+    anchor_sizes = ((32, 64, 128, 256, 512, ), ) * 3
+    aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)
+
+    model = FasterRCNN(backbone, num_classes, rpn_anchor_generator=AnchorGenerator(anchor_sizes, aspect_ratios),
+                       **kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls['fasterrcnn_mobilenet_v3_large_fpn_coco'], progress=progress)
+        model.load_state_dict(state_dict)
+    return model

torchvision/models/detection/backbone_utils.py

@@ -19,7 +19,7 @@
 
 
 __all__ = [
-    "RetinaNet", "retinanet_resnet50_fpn",
+    "RetinaNet", "retinanet_resnet50_fpn"
 ]