Adding EfficientNetsB0-B7 support (#1938)

* adding init efficientnet support

Signed-off-by: masadcv <[email protected]>

* fixing flake8 and further refactoring

Signed-off-by: masadcv <[email protected]>

* adding unittests for efficiennet

Signed-off-by: masadcv <[email protected]>

* making unittests backwards compatible python<3.8

Signed-off-by: masadcv <[email protected]>

* fixed kitty unittests file path

Signed-off-by: masadcv <[email protected]>

* adding docstrings and minor refactoring

Signed-off-by: masadcv <[email protected]>

* fix flake8-py3 failing test

Signed-off-by: masadcv <[email protected]>

* generalize drop_connect for n-dim, fix/add unittests, remove assert

Signed-off-by: masadcv <[email protected]>

* fix failing unittest, CC0-license image for test

Signed-off-by: masadcv <[email protected]>

* refactoring code for review

Signed-off-by: masadcv <[email protected]>

* WIP fix mypy type hint errors

Signed-off-by: masadcv <[email protected]>

* fix cuda test error

Signed-off-by: masadcv <[email protected]>

* WIP fix test errors

Signed-off-by: masadcv <[email protected]>

* adding non-default shape tests

Signed-off-by: masadcv <[email protected]>

* remove 3d case from non-default shape test

Signed-off-by: masadcv <[email protected]>

* refactoring and updating docs

Signed-off-by: masadcv <[email protected]>

Co-authored-by: Yiheng Wang <[email protected]>
Co-authored-by: Wenqi Li <[email protected]>

Author: 3 people

docs/source/networks.rst

@@ -35,6 +35,11 @@ Blocks
 .. autoclass:: Swish
     :members:
 
+`MemoryEfficientSwish`
+~~~~~~~~~~~~~~~~~~~~~~
+.. autoclass:: MemoryEfficientSwish
+    :members:
+
 `Mish`
 ~~~~~~
 .. autoclass:: Mish
@@ -292,6 +297,11 @@ Nets
 .. autoclass:: DenseNet
   :members:
 
+`EfficientNet`
+~~~~~~~~~~~~~~
+.. autoclass:: EfficientNet
+  :members:
+
 `SegResNet`
 ~~~~~~~~~~~
 .. autoclass:: SegResNet

monai/networks/blocks/__init__.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from .acti_norm import ADN
-from .activation import Mish, Swish
+from .activation import MemoryEfficientSwish, Mish, Swish
 from .aspp import SimpleASPP
 from .convolutions import Convolution, ResidualUnit
 from .crf import CRF

monai/networks/blocks/activation.py

@@ -17,7 +17,7 @@ class Swish(nn.Module):
     r"""Applies the element-wise function:
 
     .. math::
-        \text{Swish}(x) = x * \text{Sigmoid}(\alpha * x) for constant value alpha.
+        \text{Swish}(x) = x * \text{Sigmoid}(\alpha * x) ~~~~\text{for constant value}~ \alpha.
 
     Citation: Searching for Activation Functions, Ramachandran et al., 2017, https://arxiv.org/abs/1710.05941.
 
@@ -43,6 +43,57 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         return input * torch.sigmoid(self.alpha * input)
 
 
+class SwishImplementation(torch.autograd.Function):
+    r"""Memory efficient implementation for training
+    Follows recommendation from:
+    https://github.com/lukemelas/EfficientNet-PyTorch/issues/18#issuecomment-511677853
+
+    Results in ~ 30% memory saving during training as compared to Swish()
+    """
+
+    @staticmethod
+    def forward(ctx, input):
+        result = input * torch.sigmoid(input)
+        ctx.save_for_backward(input)
+        return result
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input = ctx.saved_tensors[0]
+        sigmoid_input = torch.sigmoid(input)
+        return grad_output * (sigmoid_input * (1 + input * (1 - sigmoid_input)))
+
+
+class MemoryEfficientSwish(nn.Module):
+    r"""Applies the element-wise function:
+
+    .. math::
+        \text{Swish}(x) = x * \text{Sigmoid}(\alpha * x) ~~~~\text{for constant value}~ \alpha=1.
+
+    Memory efficient implementation for training following recommendation from:
+    https://github.com/lukemelas/EfficientNet-PyTorch/issues/18#issuecomment-511677853
+
+    Results in ~ 30% memory saving during training as compared to Swish()
+
+    Citation: Searching for Activation Functions, Ramachandran et al., 2017, https://arxiv.org/abs/1710.05941.
+
+    Shape:
+        - Input: :math:`(N, *)` where `*` means, any number of additional
+          dimensions
+        - Output: :math:`(N, *)`, same shape as the input
+
+
+    Examples::
+
+        >>> m = Act['memswish']()
+        >>> input = torch.randn(2)
+        >>> output = m(input)
+    """
+
+    def forward(self, input: torch.Tensor):
+        return SwishImplementation.apply(input)
+
+
 class Mish(nn.Module):
     r"""Applies the element-wise function:
 

monai/networks/layers/factories.py

@@ -256,6 +256,13 @@ def swish_factory():
     return Swish
 
 
+@Act.factory_function("memswish")
+def memswish_factory():
+    from monai.networks.blocks.activation import MemoryEfficientSwish
+
+    return MemoryEfficientSwish
+
+
 @Act.factory_function("mish")
 def mish_factory():
     from monai.networks.blocks.activation import Mish

monai/networks/nets/__init__.py

@@ -15,6 +15,7 @@
 from .classifier import Classifier, Critic, Discriminator
 from .densenet import DenseNet, DenseNet121, DenseNet169, DenseNet201, DenseNet264
 from .dynunet import DynUNet, DynUnet, Dynunet
+from .efficientnet import EfficientNet, EfficientNetBN, drop_connect, get_efficientnet_image_size
 from .fullyconnectednet import FullyConnectedNet, VarFullyConnectedNet
 from .generator import Generator
 from .highresnet import HighResBlock, HighResNet