Totally _freeze_stages in MobileNetV2

mmaction/models/backbones/mobilenet_v2.py

@@ -129,8 +129,9 @@ class MobileNetV2(nn.Module):
             channels in each layer by this amount. Default: 1.0.
         out_indices (None or Sequence[int]): Output from which stages.
             Default: (7, ).
-        frozen_stages (int): Stages to be frozen (all param fixed).
-            Default: -1, which means not freezing any parameters.
+        frozen_stages (int): Stages to be frozen (all param fixed). Note that
+            the last stage in ``MobileNetV2`` is ``conv2``. Default: -1,
+            which means not freezing any parameters.
         conv_cfg (dict): Config dict for convolution layer.
             Default: None, which means using conv2d.
         norm_cfg (dict): Config dict for normalization layer.
@@ -169,8 +170,8 @@ def __init__(self,
                 raise ValueError('the item in out_indices must in '
                                  f'range(0, 8). But received {index}')
 
-        if frozen_stages not in range(-1, 8):
-            raise ValueError('frozen_stages must be in range(-1, 8). '
+        if frozen_stages not in range(-1, 9):
+            raise ValueError('frozen_stages must be in range(-1, 9). '
                              f'But received {frozen_stages}')
         self.out_indices = out_indices
         self.frozen_stages = frozen_stages
@@ -281,10 +282,12 @@ def forward(self, x):
 
     def _freeze_stages(self):
         if self.frozen_stages >= 0:
+            self.conv1.eval()
             for param in self.conv1.parameters():
                 param.requires_grad = False
         for i in range(1, self.frozen_stages + 1):
-            layer = getattr(self, f'layer{i}')
+            layer_name = self.layers[i - 1]
+            layer = getattr(self, layer_name)
             layer.eval()
             for param in layer.parameters():
                 param.requires_grad = False

tests/test_models/test_common_modules/test_mobilenet_v2.py

@@ -33,8 +33,8 @@ def is_block(modules):
         model.init_weights()
 
     with pytest.raises(ValueError):
-        # frozen_stages must in range(1, 8)
-        MobileNetV2(frozen_stages=8)
+        # frozen_stages must in range(1, 9)
+        MobileNetV2(frozen_stages=9)
 
     with pytest.raises(ValueError):
         # tout_indices in range(-1, 8)
@@ -60,6 +60,18 @@ def is_block(modules):
         for param in layer.parameters():
             assert param.requires_grad is False
 
+    # Test MobileNetV2 with all stages frozen
+    frozen_stages = 8
+    model = MobileNetV2(frozen_stages=frozen_stages)
+    model.init_weights()
+    model.train()
+
+    for mod in model.modules():
+        if not isinstance(mod, MobileNetV2):
+            assert mod.training is False
+        for param in mod.parameters():
+            assert param.requires_grad is False
+
     # Test MobileNetV2 with norm_eval=True
     model = MobileNetV2(norm_eval=True)
     model.init_weights()