Repository: pytorch/vision

Issue:
Replace SqueezeExcitation of MobileNetV3 with the one of EfficientNet
### 🚀 The feature

Investigate if it's possible to replace the `SqueezeExcitation` of MobileNetV3 in a BC manner and use the one introduced on `EfficientNet` at #pull.

### Motivation, pitch

This will reduce the amount of different SE implementations that we have and open the possibility of moving the layer on a separate area where it can be reused.

### Alternatives

_No response_

### Additional context

_No response_

cc @datumbox
[RFC] API For Common Layers In Torchvision
Picks up from discussion in https://github.com/pytorch/vision/pull/4293#discussion_r696471536

## 🚀 Feature

API for Commonly used Layers in building models.

## Motivation

A huge code duplication is involved in building very basic blocks for large neural networks. Some of these blocks can be standardized and re-used. Also these could be offered to end user as an API so that downstream libraries can build models easily.

E.g. for duplication are SqueezeExcitation, ConvBNRelu, etc.

## Pitch

Create an API called `torchvision.nn` or `torchvision.layers`
Our implementations need to be generic but not locked to certain activation functions or channels, etc.
These can be simply classes based on `nn.Module`.

An example

```
class ConvNormAct(nn.Module):
    def __init__(
        self,
        in_planes: int,
        out_planes: int,
        kernel_size: int = 3,
        stride: int = 1,
        groups: int = 1,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None,
        activation_layer: Optional[Callable[..., nn.Module]] = None,
    ) -> None: 

        super().__init__()
        #  Thsi could be even a nn.Sequential Block instead.
        padding = (kernel_size - 1) // 2 * dilation
        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, dilation, groups, bias=False)
        self.norm_layer = norm_layer
        self.activation = activation_layer

    def forward(x: Tensor):
         x = self.conv1(x)
         if self.norm_layer is not None:
              x = self.norm_layer(x)
         x = self.activation(x)
        return x

```

User can use this as

```
from torchvision import layers
l1 = layers.ConvNormAct(3, 10, norm_layer=nn.BatchNorm2d(), activation_layer=nn.Relu(inplace=True))
dummy_input = torch.randn(3, 224, 224)
out = l1(dummy_input)
```

Also layers can be mixed into new custom models.

E.g.

```
class SimpleModel(nn.Module):
    def __init__():
        self.l1 = layers.ConvNormAct(3, 10, activation_layer=nn.Relu(inplace=True))
        self.l2  = layers.ConvNormAct(10, 3, activation_layer=nn.Relu(inplace=True))

    def forward(x: Tensor):
        x = self.l1(x)
        x = self.l2(x)
        return x
```

## Points to Consider

### We have torchvision.ops then why layers? 

Ops are transforms that do manipulations with pre-processing and post processing of structures such as Boxes, Masks, Anchors, etc. These are not used in "model building" but are optional steps for specific models.
Also these are 

E.g. NMS, IoU, RoI, etc.

One doesn't need ops for every model.
E.g. You don't need to do RoI align, for DeTR. Or you don't computer IoU for segmentation masks.

With separate API can be clear distinction in what are `layer` for models and operators for `tasks` such as detection, segmentation.

## Should `torchvision.nn` contain losses?

This is tricky, and for now I see no clear winner.
PyTorch does not differentiate the API for losses or layers.
E.g. we do `nn.Conv2d` which builds a convolutional layer. Also we do `nn.CrossEntropy` or `nn.MSE` which builds a loss function.

I'm not sure whether layers should be `torchvision.layers` or `torchvision.nn` (if implemented of course)

Users don't need to worry about colliding namespaces. They can do.

```
from torch import nn
from torchvision import nn as tnn
```

Note that `nn` seems to be the convention adopted by torchtext.

### Other points to consider.

1. Portability: - 
Currently most of the torchvision models are easily copy pastable. E.g. We can easily copy paste mobilenetv2.py file and edit it on the go to customize models.
By Adding such API we can reduce the internal code duplication but these files would no longer be single standalone files for models.

2. Layer Customization : -
Layer Customization has far too many options to consider. 
E.g. there are several implementations possible for BasicBlock of ResNet or some slight modifications of inverted residual layer.
One can't create an implementation that will suit all the needs for everyone. If one tries to, then the API would be significantly complicated.

4. TorchScript: -
We shouldn't be hampering torchscript compatibility of any model while implementing above API. 

## Additional context

Some candidates for layers
1. [ConvNormAct](https://github.com/pytorch/vision/blob/5ef75d7bc666c9af3dd05b62b075df4c23ca467c/torchvision/models/mobilenetv2.py#L32)  
2. LinearDropoutAct
3. [SqueezeExcitation](https://github.com/pytorch/vision/blob/5ef75d7bc666c9af3dd05b62b075df4c23ca467c/torchvision/models/efficientnet.py#L34)
4. [StochasticDepth](https://github.com/pytorch/vision/blob/main/torchvision/ops/stochastic_depth.py)
5. [BasicBlock](https://github.com/pytorch/vision/blob/5ef75d7bc666c9af3dd05b62b075df4c23ca467c/torchvision/models/resnet.py#L37)
6. [MLP](https://github.com/facebookresearch/detr/blob/eb9f7e03ed8e2ed2cd55528989fe7df890bc3fc0/models/detr.py#L289) Simple Multi Layer Perceptron, often duplicated in downstream library codebases E.g. Detr
7. [FrozenBatchNorm2d](https://github.com/pytorch/vision/blob/5ef75d7bc666c9af3dd05b62b075df4c23ca467c/torchvision/ops/misc.py#L45)  Also used in [DETR](https://github.com/facebookresearch/detr/blob/eb9f7e03ed8e2ed2cd55528989fe7df890bc3fc0/models/backbone.py#L19).
Not sure why it is under ops, it doesn't it there

Also Quantizable versions of these !
1. QuantizableConvNormAct
2. QuantizableLinearDropoutAct

Quantizable versions will allow users to directly fuse up the models.

Additionally I will recommend not hurrying this feature, we could create `torchvision.experimental.nn` and start working out things.(or probably I can try in a fork)
Linking plans #3911 #4187 

P.S. I'm a junior developer and all my thoughts are probably step too far. So please forgive me if I'm wrong.

cc @datumbox @pmeier @NicolasHug  @fmassa 


Discussion:

> I'm a junior developer and all my thoughts are probably step too far. So please forgive me if I'm wrong.

Don't worry @oke-aditya , that is what issues are for. You've put in quite some work to compile this, so even if we totally disagree (not saying that we are) you don't need to be sorry at all. Thanks a lot for your work!

---

I'm all for reducing duplications. For my own projects I have written `ConvBNReLU` quite a few times, so having this in a central location is a good thing to have. I can't say much about the other blocks that you have proposed, but even if downstream libraries don't need them, I consider reducing duplications in our own code a win.

> I'm not sure whether layers should be `torchvision.layers` or `torchvision.nn`

For the naming issue I would go for `torchvision.nn` to follow `torch.nn` and `torchtext.nn`. In our implementations, I would promote `import torchvision.nn as tvnn` as canonical import to avoid possible conflicts with the import of the other domain libraries.

> Should `torchvision.nn` contain losses?

I don't see why not. Not sure if we want to push it as feature, i.e. making it a priority to add something there, but if someone wants to contribute a loss I think we can accept this. So I would only put the general model blocks in there in the first place and see what the community wants.

> Portability:-

Since `torchvision` is a proper library (can be installed and imported from), I wouldn't worry about the needs of someone copy-pasting our code. Plus, in the `MobileNetV2` example you gave, we import internal stuff

https://github.com/pytorch/vision/blob/5ef75d7bc666c9af3dd05b62b075df4c23ca467c/torchvision/models/mobilenetv2.py#L4

so copy-pasting the code is already not possible.

> Layer Customization : -

Yes, there is always a trade-off between customizability and ease-of-use. Given that we implement quite a few important models, IMO we can use this as "limit" for customizability. That means, all our stuff should be covered without some hacky workarounds or special casing. Everything beyond that probably  warrants a custom implementation of the user. We can always make it more general in the future if there is some need.

> Additionally I will recommend not hurrying this feature, we could create `torchvision.experimental.nn` and start working out things.(or probably I can try in a fork)

The standard way is to "mark" (in the documentation as well as in the release notes) a module in a prototype state. The aim of this is twofold:

1. We actively encourage feedback on these modules to better fit the user need. Of course we put in some thoughts in the design before we go public, but like you said: there are so many users with different needs, that it is impossible to get it right in the first try. 
2. For prototype (and somewhat weaker also for beta) functionality, we are not bound to BC. If we not completely botch our initial thoughts we probably won't change the whole UI after the first round of iterations, but smaller UI changes are ok.

Thanks @oke-aditya for bringing up this topic!
Just sharing my 2 cents. I think there are two kinds of APIs: developer APIs and public APIs. Developers APIs are those helper function/utils that we use within the codebase to improve dev efficiency and avoid code duplication. Public APIs are the products we offer to the community and commit on backward compatibility. 

For example, I would treat layers like ConvNormAct as developer API since it's just an shortcut for composing three operators together. Can we put this kind of API under `torchvision.ops.misc`?  and only used within torchvision?

For layers like SqueezeExcitation that are back by papers, referenced in many model architectures. We could make them public to the community. Putting them in `torchvision.nn` or `torchvision.ops` both sound good to me.

Hi @kazhang I agree to your thoughts. We should be exposing only those layers that are referenced in model architectures.

It would be nice to reduce internal code duplication by keeping them under `nn.misc.py` or perhaps naming these classes with `_` or keeping in `_.py` file in layers / nn folder.

I would prefer `nn` as to separate out model building logic from model pre / post processing / helper logic.




Resolving patch:
diff --git a/torchvision/models/mobilenetv3.py b/torchvision/models/mobilenetv3.py
--- a/torchvision/models/mobilenetv3.py
+++ b/torchvision/models/mobilenetv3.py
@@ -1,12 +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.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"]
@@ -18,25 +19,16 @@
 }
 
 
-class SqueezeExcitation(nn.Module):
-    # Implemented as described at Figure 4 of the MobileNetV3 paper
+class SqueezeExcitation(SElayer):
+    """DEPRECATED
+    """
     def __init__(self, input_channels: int, squeeze_factor: int = 4):
-        super().__init__()
         squeeze_channels = _make_divisible(input_channels // squeeze_factor, 8)
-        self.fc1 = nn.Conv2d(input_channels, squeeze_channels, 1)
-        self.relu = nn.ReLU(inplace=True)
-        self.fc2 = nn.Conv2d(squeeze_channels, input_channels, 1)
-
-    def _scale(self, input: Tensor, inplace: bool) -> 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)
-
-    def forward(self, input: Tensor) -> Tensor:
-        scale = self._scale(input, True)
-        return scale * input
+        super().__init__(input_channels, squeeze_channels, scale_activation=nn.Hardsigmoid)
+        self.relu = self.activation
+        delattr(self, 'activation')
+        warnings.warn(
+            "This SqueezeExcitation class is deprecated and will be removed in future versions.", FutureWarning)
 
 
 class InvertedResidualConfig:
@@ -60,7 +52,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')
@@ -81,7 +73,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,
diff --git a/torchvision/models/quantization/mobilenetv3.py b/torchvision/models/quantization/mobilenetv3.py
--- a/torchvision/models/quantization/mobilenetv3.py
+++ b/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,16 +17,53 @@
 }
 
 
-class QuantizableSqueezeExcitation(SqueezeExcitation):
+class QuantizableSqueezeExcitation(SElayer):
+    _version = 2
+
     def __init__(self, *args: Any, **kwargs: Any) -> None:
+        kwargs["scale_activation"] = nn.Hardsigmoid
         super().__init__(*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) -> None:
-        fuse_modules(self, ['fc1', 'relu'], inplace=True)
+        fuse_modules(self, ['fc1', 'activation'], inplace=True)
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        version = local_metadata.get("version", None)
+
+        if version is None or version < 2:
+            default_state_dict = {
+                "scale_activation.activation_post_process.scale": torch.tensor([1.]),
+                "scale_activation.activation_post_process.zero_point": torch.tensor([0], dtype=torch.int32),
+                "scale_activation.activation_post_process.fake_quant_enabled": torch.tensor([1]),
+                "scale_activation.activation_post_process.observer_enabled": torch.tensor([1]),
+            }
+            for k, v in default_state_dict.items():
+                full_key = prefix + k
+                if full_key not in state_dict:
+                    state_dict[full_key] = v
+
+        super()._load_from_state_dict(
+            state_dict,
+            prefix,
+            local_metadata,
+            strict,
+            missing_keys,
+            unexpected_keys,
+            error_msgs,
+        )
 
 
 class QuantizableInvertedResidual(InvertedResidual):
@@ -78,7 +116,7 @@ def _load_weights(
     arch: str,
     model: QuantizableMobileNetV3,
     model_url: Optional[str],
-    progress: bool,
+    progress: bool
 ) -> None:
     if model_url is None:
         raise ValueError("No checkpoint is available for {}".format(arch))


Test patch: