Fix torch compile, script, export

.github/workflows/tests.yml

@@ -51,7 +51,7 @@ jobs:
       run: uv pip list
 
     - name: Test with PyTest
-      run: uv run pytest -v -rsx -n 2 --cov=segmentation_models_pytorch --cov-report=xml --cov-config=pyproject.toml -k "not logits_match"
+      run: uv run pytest -v -rsx -n 2 --cov=segmentation_models_pytorch --cov-report=xml --cov-config=pyproject.toml --non-marked-only
 
     - name: Upload coverage reports to Codecov
       uses: codecov/codecov-action@v5
@@ -73,7 +73,52 @@ jobs:
     - name: Show installed packages
       run: uv pip list
     - name: Test with PyTest
-      run: RUN_SLOW=1 uv run pytest -v -rsx -n 2 -k "logits_match"
+      run: RUN_SLOW=1 uv run pytest -v -rsx -n 2 -m "logits_match"
+
+  test_torch_compile:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: astral-sh/setup-uv@v5
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: uv pip install -r requirements/required.txt -r requirements/test.txt
+    - name: Show installed packages
+      run: uv pip list
+    - name: Test with PyTest
+      run: uv run pytest -v -rsx -n 2 -m "compile"
+
+  test_torch_export:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: astral-sh/setup-uv@v5
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: uv pip install -r requirements/required.txt -r requirements/test.txt
+    - name: Show installed packages
+      run: uv pip list
+    - name: Test with PyTest
+      run: uv run pytest -v -rsx -n 2 -m "torch_export"
+
+  test_torch_script:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: astral-sh/setup-uv@v5
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: uv pip install -r requirements/required.txt -r requirements/test.txt
+    - name: Show installed packages
+      run: uv pip list
+    - name: Test with PyTest
+      run: uv run pytest -v -rsx -n 2 -m "torch_script"
 
   minimum:
     runs-on: ubuntu-latest
@@ -88,4 +133,4 @@ jobs:
     - name: Show installed packages
       run: uv pip list
     - name: Test with pytest
-      run: uv run pytest -v -rsx -n 2 -k "not logits_match"
+      run: uv run pytest -v -rsx -n 2 --non-marked-only

Makefile

@@ -7,7 +7,7 @@ install_dev: .venv
 	.venv/bin/pip install -e ".[test]"
 
 test: .venv
-	.venv/bin/pytest -v -rsx -n 2 tests/ -k "not logits_match"
+	.venv/bin/pytest -v -rsx -n 2 tests/ --non-marked-only
 
 test_all: .venv
 	RUN_SLOW=1 .venv/bin/pytest -v -rsx -n 2 tests/

pyproject.toml

@@ -39,11 +39,13 @@ docs = [
     'sphinx-book-theme',
 ]
 test = [
+    'gitpython',
     'packaging',
     'pytest',
     'pytest-cov',
     'pytest-xdist',
     'ruff>=0.9',
+    'setuptools',
 ]
 
 [project.urls]
@@ -61,18 +63,10 @@ include = ['segmentation_models_pytorch*']
 
 [tool.pytest.ini_options]
 markers = [
-    "deeplabv3",
-    "deeplabv3plus",
-    "fpn",
-    "linknet",
-    "manet",
-    "pan",
-    "psp",
-    "segformer",
-    "unet",
-    "unetplusplus",
-    "upernet",
     "logits_match",
+    "compile",
+    "torch_export",
+    "torch_script",
 ]
 
 [tool.coverage.run]

requirements/test.txt

@@ -1,5 +1,7 @@
+gitpython==3.1.44
 packaging==24.2
 pytest==8.3.4
 pytest-xdist==3.6.1
 pytest-cov==6.0.0
 ruff==0.9.1
+setuptools==75.8.0

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

@@ -3,15 +3,15 @@
 
 from . import initialization as init
 from .hub_mixin import SMPHubMixin
+from .utils import is_torch_compiling
 
 T = TypeVar("T", bound="SegmentationModel")
 
 
 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
@@ -29,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:
@@ -50,11 +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 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/base/utils.py

@@ -0,0 +1,14 @@
+import torch
+
+
+@torch.jit.unused
+def is_torch_compiling():
+    try:
+        return torch.compiler.is_compiling()
+    except Exception:
+        try:
+            import torch._dynamo as dynamo  # noqa: F401
+
+            return dynamo.is_compiling()
+        except Exception:
+            return False

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 @@
         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 @@
             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 @@
             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 @@
             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))