Extend SqueezeExcitation to support custom min_value and activation.

Author: datumbox

torchvision/models/mobilenetv3.py

@@ -20,22 +20,24 @@
 
 class SqueezeExcitation(nn.Module):
     # Implemented as described at Figure 4 of the MobileNetV3 paper
-    def __init__(self, input_channels: int, squeeze_factor: int = 4):
+    def __init__(self, input_channels: int, squeeze_factor: int = 4, min_value: Optional[int] = None,
+                 activation_fn: Callable[..., Tensor] = F.hardsigmoid):
         super().__init__()
-        squeeze_channels = _make_divisible(input_channels // squeeze_factor, 8)
+        squeeze_channels = _make_divisible(input_channels // squeeze_factor, 8, min_value)
         self.fc1 = nn.Conv2d(input_channels, squeeze_channels, 1)
         self.relu = nn.ReLU(inplace=True)
         self.fc2 = nn.Conv2d(squeeze_channels, input_channels, 1)
+        self.activation_fn = activation_fn
 
-    def _scale(self, input: Tensor, inplace: bool) -> Tensor:
+    def _scale(self, input: Tensor) -> Tensor:
         scale = F.adaptive_avg_pool2d(input, 1)
         scale = self.fc1(scale)
         scale = self.relu(scale)
         scale = self.fc2(scale)
-        return F.hardsigmoid(scale, inplace=inplace)
+        return self.activation_fn(scale)
 
     def forward(self, input: Tensor) -> Tensor:
-        scale = self._scale(input, True)
+        scale = self._scale(input)
         return scale * input
 
 

torchvision/models/quantization/mobilenetv3.py

@@ -22,7 +22,7 @@ def __init__(self, *args, **kwargs):
         self.skip_mul = nn.quantized.FloatFunctional()
 
     def forward(self, input: Tensor) -> Tensor:
-        return self.skip_mul.mul(self._scale(input, False), input)
+        return self.skip_mul.mul(self._scale(input), input)
 
     def fuse_model(self):
         fuse_modules(self, ['fc1', 'relu'], inplace=True)