Add ability to use Activated BatchNorm in decoder. (#81)

Author: SiarheiFedartsou

segmentation_models_pytorch/common/blocks.py

@@ -8,14 +8,30 @@ def __init__(self, in_channels, out_channels, kernel_size, padding=0,
         super().__init__()
 
         layers = [
-            nn.Conv2d(in_channels, out_channels, kernel_size,
-                              stride=stride, padding=padding, bias=not (use_batchnorm)),
-            nn.ReLU(inplace=True),
+            nn.Conv2d(
+                in_channels, 
+                out_channels, 
+                kernel_size,
+                stride=stride, 
+                padding=padding, 
+                bias=not (use_batchnorm)
+            )
         ]
 
-        if use_batchnorm:
-            layers.insert(1, nn.BatchNorm2d(out_channels, **batchnorm_params))
-
+        if use_batchnorm == 'inplace':
+            try:
+                from inplace_abn import InPlaceABN
+            except ImportError:
+                raise RuntimeError("In order to use `use_batchnorm='inplace'` inplace_abn package must be installed. To install see: https://github.com/mapillary/inplace_abn")
+            
+            layers.append(InPlaceABN(out_channels, activation='leaky_relu', activation_param=0.0, **batchnorm_params))
+        elif use_batchnorm:
+            layers.append(nn.BatchNorm2d(out_channels, **batchnorm_params))
+            layers.append(nn.ReLU(inplace=True))  
+        else:
+            layers.append(nn.ReLU(inplace=True))
+                
+            
         self.block = nn.Sequential(*layers)
 
     def forward(self, x):

segmentation_models_pytorch/linknet/model.py

@@ -14,11 +14,11 @@ class Linknet(EncoderDecoder):
             extractor to build segmentation model.
         encoder_weights: one of ``None`` (random initialization), ``imagenet`` (pre-training on ImageNet).
         decoder_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
-            is used.
+            is used. If 'inplace' InplaceABN will be used, allows to decrease memory consumption.
+            One of [True, False, 'inplace']
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
         activation: activation function used in ``.predict(x)`` method for inference.
             One of [``sigmoid``, ``softmax``, callable, None]
-
     Returns:
         ``torch.nn.Module``: **Linknet**
 

segmentation_models_pytorch/pspnet/model.py

@@ -14,13 +14,13 @@ class PSPNet(EncoderDecoder):
             to construct PSP module on it.
         psp_out_channels: number of filters in PSP block.
         psp_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
-                is used.
+            is used. If 'inplace' InplaceABN will be used, allows to decrease memory consumption.
+            One of [True, False, 'inplace']
         psp_aux_output: if ``True`` add auxiliary classification output for encoder training
         psp_dropout: spatial dropout rate between 0 and 1.
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
         activation: activation function used in ``.predict(x)`` method for inference.
             One of [``sigmoid``, ``softmax``, callable, None]
-
     Returns:
         ``torch.nn.Module``: **PSPNet**
 

segmentation_models_pytorch/unet/model.py

@@ -12,14 +12,14 @@ class Unet(EncoderDecoder):
         encoder_weights: one of ``None`` (random initialization), ``imagenet`` (pre-training on ImageNet).
         decoder_channels: list of numbers of ``Conv2D`` layer filters in decoder blocks
         decoder_use_batchnorm: if ``True``, ``BatchNormalisation`` layer between ``Conv2D`` and ``Activation`` layers
-            is used.
+            is used. If 'inplace' InplaceABN will be used, allows to decrease memory consumption.
+            One of [True, False, 'inplace']
         classes: a number of classes for output (output shape - ``(batch, classes, h, w)``).
         activation: activation function used in ``.predict(x)`` method for inference.
             One of [``sigmoid``, ``softmax``, callable, None]
         center: if ``True`` add ``Conv2dReLU`` block on encoder head (useful for VGG models)
         attention_type: attention module used in decoder of the model
             One of [``None``, ``scse``]
-
     Returns:
         ``torch.nn.Module``: **Unet**
 

tests/test_models.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 import random
+import importlib
 
 # mock detection module 
 sys.modules['torchvision._C'] = mock.Mock()
@@ -35,18 +36,17 @@ def _select_names(names, k=2):
         return names
 
 
-def _test_forward_backward(model_fn, encoder_name):
-
-    model = model_fn(encoder_name, encoder_weights=None)
+def _test_forward_backward(model_fn, encoder_name, **model_params):
+    model = model_fn(encoder_name, encoder_weights=None, **model_params)
 
     x = torch.ones((1, 3, 64, 64))
     y = model.forward(x)
     l = y.mean()
     l.backward()
 
 
-def _test_pretrained_model(model_fn, encoder_name, encoder_weights):
-    model = model_fn(encoder_name, encoder_weights=encoder_weights)
+def _test_pretrained_model(model_fn, encoder_name, encoder_weights, **model_params):
+    model = model_fn(encoder_name, encoder_weights=encoder_weights, **model_params)
 
     x = torch.ones((1, 3, 64, 64))
     y = model.predict(x)
@@ -82,5 +82,11 @@ def test_pspnet(encoder_name):
     _test_pretrained_model(smp.PSPNet, encoder_name, get_pretrained_weights_name(encoder_name))
 
 
+@pytest.mark.skipif(importlib.util.find_spec('inplace_abn') is None, reason='')
+def test_inplace_abn():
+    _test_forward_backward(smp.Unet, 'resnet18', decoder_use_batchnorm='inplace')
+    _test_pretrained_model(smp.Unet, 'resnet18', get_pretrained_weights_name('resnet18'), decoder_use_batchnorm='inplace')
+
+
 if __name__ == '__main__':
     pytest.main([__file__])