Add optional final upsampling for FPN (#72)

Author: qubvel

segmentation_models_pytorch/fpn/decoder.py

@@ -65,11 +65,12 @@ def __init__(
             encoder_channels,
             pyramid_channels=256,
             segmentation_channels=128,
+            final_upsampling=4,
             final_channels=1,
             dropout=0.2,
     ):
         super().__init__()
-
+        self.final_upsampling = final_upsampling
         self.conv1 = nn.Conv2d(encoder_channels[0], pyramid_channels, kernel_size=(1, 1))
 
         self.p4 = FPNBlock(pyramid_channels, encoder_channels[1])
@@ -104,5 +105,6 @@ def forward(self, x):
         x = self.dropout(x)
         x = self.final_conv(x)
 
-        x = F.interpolate(x, scale_factor=4, mode='bilinear', align_corners=True)
+        if self.final_upsampling is not None and self.final_upsampling > 1:
+            x = F.interpolate(x, scale_factor=self.final_upsampling, mode='bilinear', align_corners=True)
         return x

segmentation_models_pytorch/fpn/model.py

@@ -15,6 +15,8 @@ class FPN(EncoderDecoder):
         dropout: spatial dropout rate in range (0, 1).
         activation: activation function used in ``.predict(x)`` method for inference.
             One of [``sigmoid``, ``softmax``, callable, None]
+        final_upsampling: optional, final upsampling factor
+            (default is 4 to preserve input -> output spatial shape identity)
 
     Returns:
         ``torch.nn.Module``: **FPN**
@@ -33,6 +35,7 @@ def __init__(
             classes=1,
             dropout=0.2,
             activation='sigmoid',
+            final_upsampling=4,
     ):
         encoder = get_encoder(
             encoder_name,
@@ -45,6 +48,7 @@ def __init__(
             segmentation_channels=decoder_segmentation_channels,
             final_channels=classes,
             dropout=dropout,
+            final_upsampling=final_upsampling,
         )
 
         super().__init__(encoder, decoder, activation)