making _validate_resnet_trainable_layers generic and using it in mobilenet

Author: datumbox

torchvision/models/detection/backbone_utils.py

@@ -108,19 +108,20 @@ def resnet_fpn_backbone(
     return BackboneWithFPN(backbone, return_layers, in_channels_list, out_channels, extra_blocks=extra_blocks)
 
 
-def _validate_resnet_trainable_layers(pretrained, trainable_backbone_layers):
+def _validate_trainable_layers(pretrained, trainable_backbone_layers, max_value, default_value):
     # dont freeze any layers if pretrained model or backbone is not used
     if not pretrained:
         if trainable_backbone_layers is not None:
             warnings.warn(
                 "Changing trainable_backbone_layers has not effect if "
                 "neither pretrained nor pretrained_backbone have been set to True, "
-                "falling back to trainable_backbone_layers=5 so that all layers are trainable")
-        trainable_backbone_layers = 5
-    # by default, freeze first 2 blocks following Faster R-CNN
+                "falling back to trainable_backbone_layers={} so that all layers are trainable".format(max_value))
+        trainable_backbone_layers = max_value
+
+    # by default freeze first blocks
     if trainable_backbone_layers is None:
-        trainable_backbone_layers = 3
-    assert trainable_backbone_layers <= 5 and trainable_backbone_layers >= 0
+        trainable_backbone_layers = default_value
+    assert 0 <= trainable_backbone_layers <= max_value
     return trainable_backbone_layers
 
 

torchvision/models/detection/faster_rcnn.py

@@ -15,7 +15,7 @@
 from .rpn import RPNHead, RegionProposalNetwork
 from .roi_heads import RoIHeads
 from .transform import GeneralizedRCNNTransform
-from .backbone_utils import resnet_fpn_backbone, _validate_resnet_trainable_layers
+from .backbone_utils import resnet_fpn_backbone, _validate_trainable_layers
 
 
 __all__ = [
@@ -350,8 +350,8 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
             Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
     """
     # check default parameters and by default set it to 3 if possible
-    trainable_backbone_layers = _validate_resnet_trainable_layers(
-        pretrained or pretrained_backbone, trainable_backbone_layers)
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 5, 3)
 
     if pretrained:
         # no need to download the backbone if pretrained is set

torchvision/models/detection/keypoint_rcnn.py

@@ -7,7 +7,7 @@
 from ..utils import load_state_dict_from_url
 
 from .faster_rcnn import FasterRCNN
-from .backbone_utils import resnet_fpn_backbone, _validate_resnet_trainable_layers
+from .backbone_utils import resnet_fpn_backbone, _validate_trainable_layers
 
 
 __all__ = [
@@ -319,8 +319,8 @@ def keypointrcnn_resnet50_fpn(pretrained=False, progress=True,
             Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
     """
     # check default parameters and by default set it to 3 if possible
-    trainable_backbone_layers = _validate_resnet_trainable_layers(
-        pretrained or pretrained_backbone, trainable_backbone_layers)
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 5, 3)
 
     if pretrained:
         # no need to download the backbone if pretrained is set

torchvision/models/detection/mask_rcnn.py

@@ -8,7 +8,7 @@
 from ..utils import load_state_dict_from_url
 
 from .faster_rcnn import FasterRCNN
-from .backbone_utils import resnet_fpn_backbone, _validate_resnet_trainable_layers
+from .backbone_utils import resnet_fpn_backbone, _validate_trainable_layers
 
 __all__ = [
     "MaskRCNN", "maskrcnn_resnet50_fpn",
@@ -314,8 +314,8 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True,
             Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
     """
     # check default parameters and by default set it to 3 if possible
-    trainable_backbone_layers = _validate_resnet_trainable_layers(
-        pretrained or pretrained_backbone, trainable_backbone_layers)
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 5, 3)
 
     if pretrained:
         # no need to download the backbone if pretrained is set

torchvision/models/detection/retinanet.py

@@ -12,7 +12,7 @@
 from . import _utils as det_utils
 from .anchor_utils import AnchorGenerator
 from .transform import GeneralizedRCNNTransform
-from .backbone_utils import resnet_fpn_backbone, _validate_resnet_trainable_layers, mobilenet_fpn_backbone
+from .backbone_utils import resnet_fpn_backbone, _validate_trainable_layers, mobilenet_fpn_backbone
 from ...ops.feature_pyramid_network import LastLevelP6P7
 from ...ops import sigmoid_focal_loss
 from ...ops import boxes as box_ops
@@ -556,6 +556,7 @@ def forward(self, images, targets=None):
             return losses, detections
         return self.eager_outputs(losses, detections)
 
+
 # TODO: replace with pytorch links
 model_urls = {
     'retinanet_mobilenet_v3_large_fpn_coco':
@@ -608,8 +609,8 @@ def retinanet_resnet50_fpn(pretrained=False, progress=True,
             Valid values are between 0 and 5, with 5 meaning all backbone layers are trainable.
     """
     # check default parameters and by default set it to 3 if possible
-    trainable_backbone_layers = _validate_resnet_trainable_layers(
-        pretrained or pretrained_backbone, trainable_backbone_layers)
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 5, 3)
 
     if pretrained:
         # no need to download the backbone if pretrained is set
@@ -626,8 +627,8 @@ def retinanet_resnet50_fpn(pretrained=False, progress=True,
     return model
 
 
-def retinanet_mobilenet_v3_large_fpn(pretrained=False, progress=True,
-                                     num_classes=91, pretrained_backbone=True, **kwargs):
+def retinanet_mobilenet_v3_large_fpn(pretrained=False, progress=True, num_classes=91, pretrained_backbone=True,
+                                     trainable_backbone_layers=None, **kwargs):
     """
     Constructs a RetinaNet model with a MobileNetV3-Large-FPN backbone. It works similarly
     to RetinaNet with ResNet-50-FPN backbone. See `retinanet_resnet50_fpn` for more details.
@@ -642,10 +643,19 @@ def retinanet_mobilenet_v3_large_fpn(pretrained=False, progress=True,
     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 5, with 5 meaning all backbone layers are trainable.
     """
+    # check default parameters and by default set it to 3 if possible
+    trainable_backbone_layers = _validate_trainable_layers(
+        pretrained or pretrained_backbone, trainable_backbone_layers, 6, 2)
+
     if pretrained:
         pretrained_backbone = False
-    backbone = mobilenet_fpn_backbone("mobilenet_v3_large", pretrained_backbone, returned_layers=[4, 5])
+    backbone = mobilenet_fpn_backbone("mobilenet_v3_large", pretrained_backbone, returned_layers=[4, 5],
+                                      trainable_layers=trainable_backbone_layers)
 
     anchor_sizes = ((128,), (256,), (512,))
     aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes)