Huge fix for torch scripting (except Unet++ and UperNet)

Author: qubvel

segmentation_models_pytorch/base/hub_mixin.py

@@ -1,3 +1,4 @@
+import torch
 import json
 from pathlib import Path
 from typing import Optional, Union
@@ -114,6 +115,7 @@ def save_pretrained(
         return result
 
     @property
+    @torch.jit.unused
     def config(self) -> dict:
         return self._hub_mixin_config
 

segmentation_models_pytorch/base/model.py

@@ -11,8 +11,7 @@
 class SegmentationModel(torch.nn.Module, SMPHubMixin):
     """Base class for all segmentation models."""
 
-    # if model supports shape not divisible by 2 ^ n
-    # set to False
+    # if model supports shape not divisible by 2 ^ n set to False
     requires_divisible_input_shape = True
 
     # Fix type-hint for models, to avoid HubMixin signature
@@ -30,6 +29,9 @@ def check_input_shape(self, x):
         """Check if the input shape is divisible by the output stride.
         If not, raise a RuntimeError.
         """
+        if not self.requires_divisible_input_shape:
+            return
+
         h, w = x.shape[-2:]
         output_stride = self.encoder.output_stride
         if h % output_stride != 0 or w % output_stride != 0:
@@ -51,15 +53,13 @@ def check_input_shape(self, x):
     def forward(self, x):
         """Sequentially pass `x` trough model`s encoder, decoder and heads"""
 
-        if (
-            not torch.jit.is_tracing()
-            and not is_torch_compiling()
-            and self.requires_divisible_input_shape
+        if not (
+            torch.jit.is_scripting() or torch.jit.is_tracing() or is_torch_compiling()
         ):
             self.check_input_shape(x)
 
         features = self.encoder(x)
-        decoder_output = self.decoder(*features)
+        decoder_output = self.decoder(features)
 
         masks = self.segmentation_head(decoder_output)
 

segmentation_models_pytorch/decoders/deeplabv3/decoder.py

@@ -31,7 +31,7 @@
 """
 
 from collections.abc import Iterable, Sequence
-from typing import Literal
+from typing import Literal, List
 
 import torch
 from torch import nn
@@ -49,21 +49,42 @@ def __init__(
         aspp_separable: bool,
         aspp_dropout: float,
     ):
-        super().__init__(
-            ASPP(
-                in_channels,
-                out_channels,
-                atrous_rates,
-                separable=aspp_separable,
-                dropout=aspp_dropout,
-            ),
-            nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False),
-            nn.BatchNorm2d(out_channels),
-            nn.ReLU(),
+        super().__init__()
+        self.aspp = ASPP(
+            in_channels,
+            out_channels,
+            atrous_rates,
+            separable=aspp_separable,
+            dropout=aspp_dropout,
         )
-
-    def forward(self, *features):
-        return super().forward(features[-1])
+        self.conv = nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False)
+        self.bn = nn.BatchNorm2d(out_channels)
+        self.relu = nn.ReLU()
+
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:
+        x = features[-1]
+        x = self.aspp(x)
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.relu(x)
+        return x
+
+    def load_state_dict(self, state_dict, *args, **kwargs):
+        # For backward compatibility, previously this module was Sequential
+        # and was not scriptable.
+        keys = list(state_dict.keys())
+        for key in keys:
+            new_key = key
+            if key.startswith("0."):
+                new_key = "aspp." + key[2:]
+            elif key.startswith("1."):
+                new_key = "conv." + key[2:]
+            elif key.startswith("2."):
+                new_key = "bn." + key[2:]
+            elif key.startswith("3."):
+                new_key = "relu." + key[2:]
+            state_dict[new_key] = state_dict.pop(key)
+        super().load_state_dict(state_dict, *args, **kwargs)
 
 
 class DeepLabV3PlusDecoder(nn.Module):
@@ -124,7 +145,7 @@ def __init__(
             nn.ReLU(),
         )
 
-    def forward(self, *features):
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:
         aspp_features = self.aspp(features[-1])
         aspp_features = self.up(aspp_features)
         high_res_features = self.block1(features[2])
@@ -174,7 +195,7 @@ def __init__(self, in_channels: int, out_channels: int):
             nn.ReLU(),
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         size = x.shape[-2:]
         for mod in self:
             x = mod(x)
@@ -216,7 +237,7 @@ def __init__(
             nn.Dropout(dropout),
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         res = []
         for conv in self.convs:
             res.append(conv(x))

segmentation_models_pytorch/decoders/fpn/decoder.py

@@ -2,9 +2,11 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
+from typing import List, Literal
+
 
 class Conv3x3GNReLU(nn.Module):
-    def __init__(self, in_channels, out_channels, upsample=False):
+    def __init__(self, in_channels: int, out_channels: int, upsample: bool = False):
         super().__init__()
         self.upsample = upsample
         self.block = nn.Sequential(
@@ -15,27 +17,27 @@ def __init__(self, in_channels, out_channels, upsample=False):
             nn.ReLU(inplace=True),
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.block(x)
         if self.upsample:
-            x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
+            x = F.interpolate(x, scale_factor=2.0, mode="bilinear", align_corners=True)
         return x
 
 
 class FPNBlock(nn.Module):
-    def __init__(self, pyramid_channels, skip_channels):
+    def __init__(self, pyramid_channels: int, skip_channels: int):
         super().__init__()
         self.skip_conv = nn.Conv2d(skip_channels, pyramid_channels, kernel_size=1)
 
-    def forward(self, x, skip=None):
-        x = F.interpolate(x, scale_factor=2, mode="nearest")
+    def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
         skip = self.skip_conv(skip)
         x = x + skip
         return x
 
 
 class SegmentationBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, n_upsamples=0):
+    def __init__(self, in_channels: int, out_channels: int, n_upsamples: int = 0):
         super().__init__()
 
         blocks = [Conv3x3GNReLU(in_channels, out_channels, upsample=bool(n_upsamples))]
@@ -51,36 +53,37 @@ def forward(self, x):
 
 
 class MergeBlock(nn.Module):
-    def __init__(self, policy):
+    def __init__(self, policy: Literal["add", "cat"]):
         super().__init__()
         if policy not in ["add", "cat"]:
             raise ValueError(
                 "`merge_policy` must be one of: ['add', 'cat'], got {}".format(policy)
             )
         self.policy = policy
 
-    def forward(self, x):
+    def forward(self, x: List[torch.Tensor]) -> torch.Tensor:
         if self.policy == "add":
-            return sum(x)
+            output = torch.stack(x).sum(dim=0)
         elif self.policy == "cat":
-            return torch.cat(x, dim=1)
+            output = torch.cat(x, dim=1)
         else:
             raise ValueError(
                 "`merge_policy` must be one of: ['add', 'cat'], got {}".format(
                     self.policy
                 )
             )
+        return output
 
 
 class FPNDecoder(nn.Module):
     def __init__(
         self,
-        encoder_channels,
-        encoder_depth=5,
-        pyramid_channels=256,
-        segmentation_channels=128,
-        dropout=0.2,
-        merge_policy="add",
+        encoder_channels: List[int],
+        encoder_depth: int = 5,
+        pyramid_channels: int = 256,
+        segmentation_channels: int = 128,
+        dropout: float = 0.2,
+        merge_policy: Literal["add", "cat"] = "add",
     ):
         super().__init__()
 
@@ -116,17 +119,20 @@ def __init__(
         self.merge = MergeBlock(merge_policy)
         self.dropout = nn.Dropout2d(p=dropout, inplace=True)
 
-    def forward(self, *features):
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:
         c2, c3, c4, c5 = features[-4:]
 
         p5 = self.p5(c5)
         p4 = self.p4(p5, c4)
         p3 = self.p3(p4, c3)
         p2 = self.p2(p3, c2)
 
-        feature_pyramid = [
-            seg_block(p) for seg_block, p in zip(self.seg_blocks, [p5, p4, p3, p2])
-        ]
+        s5 = self.seg_blocks[0](p5)
+        s4 = self.seg_blocks[1](p4)
+        s3 = self.seg_blocks[2](p3)
+        s2 = self.seg_blocks[3](p2)
+
+        feature_pyramid = [s5, s4, s3, s2]
         x = self.merge(feature_pyramid)
         x = self.dropout(x)
 

segmentation_models_pytorch/decoders/linknet/decoder.py

@@ -1,10 +1,12 @@
+import torch
 import torch.nn as nn
 
+from typing import List, Optional
 from segmentation_models_pytorch.base import modules
 
 
 class TransposeX2(nn.Sequential):
-    def __init__(self, in_channels, out_channels, use_batchnorm=True):
+    def __init__(self, in_channels: int, out_channels: int, use_batchnorm: bool = True):
         super().__init__()
         layers = [
             nn.ConvTranspose2d(
@@ -20,7 +22,7 @@ def __init__(self, in_channels, out_channels, use_batchnorm=True):
 
 
 class DecoderBlock(nn.Module):
-    def __init__(self, in_channels, out_channels, use_batchnorm=True):
+    def __init__(self, in_channels: int, out_channels: int, use_batchnorm: bool = True):
         super().__init__()
 
         self.block = nn.Sequential(
@@ -41,7 +43,9 @@ def __init__(self, in_channels, out_channels, use_batchnorm=True):
             ),
         )
 
-    def forward(self, x, skip=None):
+    def forward(
+        self, x: torch.Tensor, skip: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
         x = self.block(x)
         if skip is not None:
             x = x + skip
@@ -50,7 +54,11 @@ def forward(self, x, skip=None):
 
 class LinknetDecoder(nn.Module):
     def __init__(
-        self, encoder_channels, prefinal_channels=32, n_blocks=5, use_batchnorm=True
+        self,
+        encoder_channels: List[int],
+        prefinal_channels: int = 32,
+        n_blocks: int = 5,
+        use_batchnorm: bool = True,
     ):
         super().__init__()
 
@@ -68,7 +76,7 @@ def __init__(
             ]
         )
 
-    def forward(self, *features):
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:
         features = features[1:]  # remove first skip
         features = features[::-1]  # reverse channels to start from head of encoder