Add ability to use concatenation and addition in FPN (#82)

Author: SiarheiFedartsou

segmentation_models_pytorch/fpn/decoder.py

@@ -68,8 +68,14 @@ def __init__(
             final_upsampling=4,
             final_channels=1,
             dropout=0.2,
+            merge_policy='add'
     ):
         super().__init__()
+        
+        if merge_policy not in ['add', 'cat']:
+            raise ValueError("`merge_policy` must be one of: ['add', 'cat'], got {}".format(merge_policy))
+        self.merge_policy = merge_policy
+
         self.final_upsampling = final_upsampling
         self.conv1 = nn.Conv2d(encoder_channels[0], pyramid_channels, kernel_size=(1, 1))
 
@@ -83,6 +89,10 @@ def __init__(
         self.s2 = SegmentationBlock(pyramid_channels, segmentation_channels, n_upsamples=0)
 
         self.dropout = nn.Dropout2d(p=dropout, inplace=True)
+
+        if self.merge_policy == 'cat':
+            segmentation_channels *= 4
+        
         self.final_conv = nn.Conv2d(segmentation_channels, final_channels, kernel_size=1, padding=0)
 
         self.initialize()
@@ -100,7 +110,10 @@ def forward(self, x):
         s3 = self.s3(p3)
         s2 = self.s2(p2)
 
-        x = s5 + s4 + s3 + s2
+        if self.merge_policy == 'add':
+            x = s5 + s4 + s3 + s2
+        elif self.merge_policy == 'cat':
+            x = torch.cat([s5, s4, s3, s2], dim=1)
 
         x = self.dropout(x)
         x = self.final_conv(x)

segmentation_models_pytorch/fpn/model.py

@@ -17,7 +17,8 @@ class FPN(EncoderDecoder):
             One of [``sigmoid``, ``softmax``, callable, None]
         final_upsampling: optional, final upsampling factor
             (default is 4 to preserve input -> output spatial shape identity)
-
+        decoder_merge_policy: determines how to merge outputs inside FPN.
+            One of [``add``, ``cat``]
     Returns:
         ``torch.nn.Module``: **FPN**
 
@@ -36,6 +37,7 @@ def __init__(
             dropout=0.2,
             activation='sigmoid',
             final_upsampling=4,
+            decoder_merge_policy='add'
     ):
         encoder = get_encoder(
             encoder_name,
@@ -49,6 +51,7 @@ def __init__(
             final_channels=classes,
             dropout=dropout,
             final_upsampling=final_upsampling,
+            merge_policy=decoder_merge_policy
         )
 
         super().__init__(encoder, decoder, activation)

tests/test_models.py

@@ -65,6 +65,10 @@ def test_fpn(encoder_name):
     _test_forward_backward(smp.FPN, encoder_name)
     _test_pretrained_model(smp.FPN, encoder_name, get_pretrained_weights_name(encoder_name))
 
+    from functools import partial
+    _test_forward_backward(partial(smp.FPN, decoder_merge_policy='cat'), encoder_name)
+    _test_pretrained_model(partial(smp.FPN, decoder_merge_policy='cat'), encoder_name, get_pretrained_weights_name(encoder_name))
+
 
 @pytest.mark.parametrize('encoder_name', _select_names(ENCODERS, k=1))
 def test_linknet(encoder_name):