Add Squeeze and Excitation to ResNeXt models (#426)

Summary:
Pull Request resolved: #426

Added a `SqueezeAndExcitation` layer to a new sub-package, `models.common` (open to other suggestions, I didn't want to have a `generic.py` or `util.py` as that is too vague and broad).

Plugged in the layer to `ResNeXt` models.

Differential Revision: D20283172

fbshipit-source-id: 21d5183a61d7aa13fca094afe95ecb0aa18f1632

Author: mannatsingh

classy_vision/models/__init__.py

@@ -93,6 +93,7 @@ def build_model(config):
 from .resnet import ResNet  # isort:skip
 from .resnext import ResNeXt  # isort:skip
 from .resnext3d import ResNeXt3D  # isort:skip
+from .squeeze_and_excitation_layer import SqueezeAndExcitationLayer  # isort:skip
 
 
 __all__ = [
@@ -107,4 +108,5 @@ def build_model(config):
     "ResNet",
     "ResNeXt",
     "ResNeXt3D",
+    "SqueezeAndExcitationLayer",
 ]

classy_vision/models/resnext.py

@@ -8,6 +8,7 @@
 Implementation of ResNeXt (https://arxiv.org/pdf/1611.05431.pdf)
 """
 
+import copy
 import math
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -16,6 +17,7 @@
 
 from . import register_model
 from .classy_model import ClassyModel
+from .squeeze_and_excitation_layer import SqueezeAndExcitationLayer
 
 
 # global setting for in-place ReLU:
@@ -55,6 +57,8 @@ def __init__(
         mid_planes_and_cardinality=None,
         reduction=4,
         final_bn_relu=True,
+        use_se=False,
+        se_reduction_ratio=16,
     ):
 
         # assertions on inputs:
@@ -79,6 +83,12 @@ def __init__(
                 nn.BatchNorm2d(out_planes),
             )
 
+        self.se = (
+            SqueezeAndExcitationLayer(out_planes, reduction_ratio=se_reduction_ratio)
+            if use_se
+            else None
+        )
+
     def forward(self, x):
 
         # if required, perform downsampling along shortcut connection:
@@ -92,6 +102,10 @@ def forward(self, x):
 
         if self.final_bn_relu:
             out = self.bn(out)
+
+        if self.se is not None:
+            out = self.se(out)
+
         # add residual connection, perform rely + batchnorm, and return result:
         out += residual
         if self.final_bn_relu:
@@ -101,7 +115,7 @@ def forward(self, x):
 
 class BasicLayer(GenericLayer):
     """
-        ResNeXt bottleneck layer with `in_planes` input planes and `out_planes`
+        ResNeXt layer with `in_planes` input planes and `out_planes`
         output planes.
     """
 
@@ -113,6 +127,8 @@ def __init__(
         mid_planes_and_cardinality=None,
         reduction=4,
         final_bn_relu=True,
+        use_se=False,
+        se_reduction_ratio=16,
     ):
 
         # assertions on inputs:
@@ -128,13 +144,15 @@ def __init__(
         )
 
         # call constructor of generic layer:
-        super(BasicLayer, self).__init__(
+        super().__init__(
             convolutional_block,
             in_planes,
             out_planes,
             stride=stride,
             reduction=reduction,
             final_bn_relu=final_bn_relu,
+            use_se=use_se,
+            se_reduction_ratio=se_reduction_ratio,
         )
 
 
@@ -152,6 +170,8 @@ def __init__(
         mid_planes_and_cardinality=None,
         reduction=4,
         final_bn_relu=True,
+        use_se=False,
+        se_reduction_ratio=16,
     ):
 
         # assertions on inputs:
@@ -185,6 +205,8 @@ def __init__(
             stride=stride,
             reduction=reduction,
             final_bn_relu=final_bn_relu,
+            use_se=use_se,
+            se_reduction_ratio=se_reduction_ratio,
         )
 
 
@@ -236,14 +258,20 @@ def __init__(
         basic_layer: bool = False,
         final_bn_relu: bool = True,
         bn_weight_decay: Optional[bool] = False,
+        use_se: bool = False,
+        se_reduction_ratio: int = 16,
     ):
         """
             Implementation of `ResNeXt <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
-            Set ``small_input`` to `True` for 32x32 sized image inputs.
-
-            Set ``final_bn_relu`` to `False` to exclude the final batchnorm and
-            ReLU layers. These settings are useful when training Siamese networks.
+            Args:
+                small_input: set to `True` for 32x32 sized image inputs.
+                final_bn_relu: set to `False` to exclude the final batchnorm and
+                    ReLU layers. These settings are useful when training Siamese
+                    networks.
+                use_se: Enable squeeze and excitation
+                se_reduction_ratio: The reduction ratio to apply in the excitation
+                    stage. Only used if `use_se` is `True`.
         """
         super().__init__()
 
@@ -263,6 +291,7 @@ def __init__(
             and is_pos_int(base_width_and_cardinality[0])
             and is_pos_int(base_width_and_cardinality[1])
         )
+        assert isinstance(use_se, bool), "use_se has to be a boolean"
 
         # Chooses whether to apply weight decay to batch norm
         # parameters. This improves results in some situations,
@@ -295,6 +324,8 @@ def __init__(
                 mid_planes_and_cardinality=mid_planes_and_cardinality,
                 reduction=reduction,
                 final_bn_relu=final_bn_relu or (idx != (len(out_planes) - 1)),
+                use_se=use_se,
+                se_reduction_ratio=se_reduction_ratio,
             )
             blocks.append(nn.Sequential(*new_block))
         self.blocks = nn.Sequential(*blocks)
@@ -337,6 +368,8 @@ def _make_resolution_block(
         mid_planes_and_cardinality=None,
         reduction=4,
         final_bn_relu=True,
+        use_se=False,
+        se_reduction_ratio=16,
     ):
 
         # add the desired number of residual blocks:
@@ -352,6 +385,8 @@ def _make_resolution_block(
                         mid_planes_and_cardinality=mid_planes_and_cardinality,
                         reduction=reduction,
                         final_bn_relu=final_bn_relu or (idx != (num_blocks - 1)),
+                        use_se=use_se,
+                        se_reduction_ratio=se_reduction_ratio,
                     ),
                 )
             )
@@ -379,6 +414,8 @@ def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
             "final_bn_relu": config.get("final_bn_relu", True),
             "zero_init_bn_residuals": config.get("zero_init_bn_residuals", False),
             "bn_weight_decay": config.get("bn_weight_decay", False),
+            "use_se": config.get("use_se", False),
+            "se_reduction_ratio": config.get("se_reduction_ratio", 16),
         }
         return cls(**config)
 
@@ -421,65 +458,68 @@ def model_depth(self):
         return sum(self.num_blocks)
 
 
+class _ResNeXt(ResNeXt):
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
+        config = copy.deepcopy(config)
+        config.pop("name")
+        return cls(**config)
+
+
 @register_model("resnet18")
-class ResNet18(ResNeXt):
-    def __init__(self):
+class ResNet18(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
-            num_blocks=[2, 2, 2, 2], basic_layer=True, zero_init_bn_residuals=True
+            num_blocks=[2, 2, 2, 2],
+            basic_layer=True,
+            zero_init_bn_residuals=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 @register_model("resnet34")
 class ResNet34(ResNeXt):
-    def __init__(self):
+    def __init__(self, **kwargs):
         super().__init__(
-            num_blocks=[3, 4, 6, 3], basic_layer=True, zero_init_bn_residuals=True
+            num_blocks=[3, 4, 6, 3],
+            basic_layer=True,
+            zero_init_bn_residuals=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 @register_model("resnet50")
-class ResNet50(ResNeXt):
-    def __init__(self):
+class ResNet50(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
-            num_blocks=[3, 4, 6, 3], basic_layer=False, zero_init_bn_residuals=True
+            num_blocks=[3, 4, 6, 3],
+            basic_layer=False,
+            zero_init_bn_residuals=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 @register_model("resnet101")
-class ResNet101(ResNeXt):
-    def __init__(self):
+class ResNet101(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
-            num_blocks=[3, 4, 23, 3], basic_layer=False, zero_init_bn_residuals=True
+            num_blocks=[3, 4, 23, 3],
+            basic_layer=False,
+            zero_init_bn_residuals=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 @register_model("resnet152")
-class ResNet152(ResNeXt):
-    def __init__(self):
+class ResNet152(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
-            num_blocks=[3, 8, 36, 3], basic_layer=False, zero_init_bn_residuals=True
+            num_blocks=[3, 8, 36, 3],
+            basic_layer=False,
+            zero_init_bn_residuals=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 # Note, the ResNeXt models all have weight decay enabled for the batch
 # norm parameters. We have found empirically that this gives better
@@ -488,48 +528,39 @@ def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
 # training on other datasets, we have observed losses in accuracy (for
 # example, the dataset used in https://arxiv.org/abs/1805.00932).
 @register_model("resnext50_32x4d")
-class ResNeXt50(ResNeXt):
-    def __init__(self):
+class ResNeXt50(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
             num_blocks=[3, 4, 6, 3],
             basic_layer=False,
             zero_init_bn_residuals=True,
             base_width_and_cardinality=(4, 32),
             bn_weight_decay=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 @register_model("resnext101_32x4d")
-class ResNeXt101(ResNeXt):
-    def __init__(self):
+class ResNeXt101(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
             num_blocks=[3, 4, 23, 3],
             basic_layer=False,
             zero_init_bn_residuals=True,
             base_width_and_cardinality=(4, 32),
             bn_weight_decay=True,
+            **kwargs,
         )
 
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()
-
 
 @register_model("resnext152_32x4d")
-class ResNeXt152(ResNeXt):
-    def __init__(self):
+class ResNeXt152(_ResNeXt):
+    def __init__(self, **kwargs):
         super().__init__(
             num_blocks=[3, 8, 36, 3],
             basic_layer=False,
             zero_init_bn_residuals=True,
             base_width_and_cardinality=(4, 32),
             bn_weight_decay=True,
+            **kwargs,
         )
-
-    @classmethod
-    def from_config(cls, config: Dict[str, Any]) -> "ResNeXt":
-        return cls()

classy_vision/models/squeeze_and_excitation_layer.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+import torch.nn as nn
+
+
+class SqueezeAndExcitationLayer(nn.Module):
+    """Squeeze and excitation layer, as per https://arxiv.org/pdf/1709.01507.pdf"""
+
+    def __init__(self, in_planes, reduction_ratio=16):
+        super().__init__()
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        reduced_planes = in_planes // reduction_ratio
+        self.excitation = nn.Sequential(
+            nn.Conv2d(in_planes, reduced_planes, kernel_size=1, stride=1, bias=True),
+            nn.ReLU(),
+            nn.Conv2d(reduced_planes, in_planes, kernel_size=1, stride=1, bias=True),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x):
+        x_squeezed = self.avgpool(x)
+        x_excited = self.excitation(x_squeezed)
+        x_scaled = x * x_excited
+        return x_scaled