Deprecating the mobilenetv3.SqueezeExcitation layer.

Author: datumbox

torchvision/models/mobilenetv3.py

@@ -1,13 +1,13 @@
+import warnings
 import torch
 
 from functools import partial
 from torch import nn, Tensor
-from torch.nn import functional as F
-from typing import Any, Callable, Dict, List, Optional, Sequence
+from typing import Any, Callable, List, Optional, Sequence
 
 from .._internally_replaced_utils import load_state_dict_from_url
-from torchvision.models.efficientnet import SqueezeExcitation as SElayer
-from torchvision.models.mobilenetv2 import _make_divisible, ConvBNActivation
+from .efficientnet import SqueezeExcitation as SElayer
+from .mobilenetv2 import _make_divisible, ConvBNActivation
 
 
 __all__ = ["MobileNetV3", "mobilenet_v3_large", "mobilenet_v3_small"]
@@ -20,10 +20,14 @@
 
 
 class SqueezeExcitation(SElayer):
+    """DEPRECATED
+    """
     def __init__(self, input_channels: int, squeeze_factor: int = 4):
         squeeze_channels = _make_divisible(input_channels // squeeze_factor, 8)
-        super().__init__(input_channels, squeeze_channels, activation=nn.ReLU, scale_activation=nn.Hardsigmoid)
+        super().__init__(input_channels, squeeze_channels, scale_activation=nn.Hardsigmoid)
         self.relu = self.activation
+        warnings.warn(
+            "This SqueezeExcitation class is deprecated and will be removed in future versions.", FutureWarning)
 
 
 class InvertedResidualConfig:
@@ -47,7 +51,7 @@ def adjust_channels(channels: int, width_mult: float):
 class InvertedResidual(nn.Module):
     # Implemented as described at section 5 of MobileNetV3 paper
     def __init__(self, cnf: InvertedResidualConfig, norm_layer: Callable[..., nn.Module],
-                 se_layer: Callable[..., nn.Module] = SqueezeExcitation):
+                 se_layer: Callable[..., nn.Module] = partial(SElayer, scale_activation=nn.Hardsigmoid)):
         super().__init__()
         if not (1 <= cnf.stride <= 2):
             raise ValueError('illegal stride value')
@@ -68,7 +72,8 @@ def __init__(self, cnf: InvertedResidualConfig, norm_layer: Callable[..., nn.Mod
                                        stride=stride, dilation=cnf.dilation, groups=cnf.expanded_channels,
                                        norm_layer=norm_layer, activation_layer=activation_layer))
         if cnf.use_se:
-            layers.append(se_layer(cnf.expanded_channels))
+            squeeze_channels = _make_divisible(cnf.expanded_channels // 4, 8)
+            layers.append(se_layer(cnf.expanded_channels, squeeze_channels))
 
         # project
         layers.append(ConvBNActivation(cnf.expanded_channels, cnf.out_channels, kernel_size=1, norm_layer=norm_layer,

torchvision/models/quantization/mobilenetv3.py

@@ -1,8 +1,9 @@
 import torch
 from torch import nn, Tensor
 from ..._internally_replaced_utils import load_state_dict_from_url
-from torchvision.models.mobilenetv3 import InvertedResidual, InvertedResidualConfig, ConvBNActivation, MobileNetV3,\
-    SqueezeExcitation, model_urls, _mobilenet_v3_conf
+from ..efficientnet import SqueezeExcitation as SElayer
+from ..mobilenetv3 import InvertedResidual, InvertedResidualConfig, ConvBNActivation, MobileNetV3,\
+    model_urls, _mobilenet_v3_conf
 from torch.quantization import QuantStub, DeQuantStub, fuse_modules
 from typing import Any, List, Optional
 from .utils import _replace_relu
@@ -16,7 +17,7 @@
 }
 
 
-class QuantizableSqueezeExcitation(SqueezeExcitation):
+class QuantizableSqueezeExcitation(SElayer):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self.skip_mul = nn.quantized.FloatFunctional()
@@ -25,7 +26,7 @@ def forward(self, input: Tensor) -> Tensor:
         return self.skip_mul.mul(self._scale(input), input)
 
     def fuse_model(self) -> None:
-        fuse_modules(self, ['fc1', 'relu'], inplace=True)
+        fuse_modules(self, ['fc1', 'activation'], inplace=True)
 
 
 class QuantizableInvertedResidual(InvertedResidual):