Fix torch compile, script, export #1031

qubvel · 2025-01-13T18:21:04Z

Huge PR, but many things are dependent, so include everything here:

Deprecate timm- encoders (map weights to tu- except for EfficientNet and SKNet).
Fix torch.compile for encoders and add tests (only EfficientNet is currently skipped, but it can be fixed once we copy-paste the code).
Fix and add tests for torch.export.export
Fix and add tests for torch.jit.script

Fixes:

codecov · 2025-01-13T19:27:41Z

Codecov Report

Attention: Patch coverage is 93.98496% with 40 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...entation_models_pytorch/encoders/timm_universal.py	40.00%	12 Missing ⚠️
segmentation_models_pytorch/base/model.py	62.50%	9 Missing ⚠️
segmentation_models_pytorch/base/utils.py	45.45%	6 Missing ⚠️
segmentation_models_pytorch/encoders/_base.py	78.94%	4 Missing ⚠️
...ntation_models_pytorch/decoders/deeplabv3/model.py	75.00%	3 Missing ⚠️
...ntation_models_pytorch/encoders/mix_transformer.py	96.92%	2 Missing ⚠️
...egmentation_models_pytorch/decoders/fpn/decoder.py	94.73%	1 Missing ⚠️
...egmentation_models_pytorch/decoders/pan/decoder.py	94.73%	1 Missing ⚠️
...ation_models_pytorch/decoders/segformer/decoder.py	83.33%	1 Missing ⚠️
...ation_models_pytorch/encoders/timm_efficientnet.py	96.87%	1 Missing ⚠️

Files with missing lines	Coverage Δ
segmentation_models_pytorch/base/hub_mixin.py	`98.33% <100.00%> (+0.05%)`	⬆️
...ation_models_pytorch/decoders/deeplabv3/decoder.py	`98.68% <100.00%> (+0.17%)`	⬆️
...ntation_models_pytorch/decoders/linknet/decoder.py	`100.00% <100.00%> (ø)`
...mentation_models_pytorch/decoders/manet/decoder.py	`97.75% <100.00%> (ø)`
...entation_models_pytorch/decoders/pspnet/decoder.py	`100.00% <100.00%> (ø)`
...gmentation_models_pytorch/decoders/unet/decoder.py	`91.37% <100.00%> (ø)`
...on_models_pytorch/decoders/unetplusplus/decoder.py	`92.85% <100.00%> (+0.10%)`	⬆️
...tion_models_pytorch/decoders/unetplusplus/model.py	`95.00% <100.00%> (+0.26%)`	⬆️
...ntation_models_pytorch/decoders/upernet/decoder.py	`98.00% <100.00%> (ø)`
...mentation_models_pytorch/decoders/upernet/model.py	`100.00% <100.00%> (ø)`
... and 23 more

... and 1 file with indirect coverage changes

brianhou0208 · 2025-01-14T14:07:46Z

segmentation_models_pytorch/encoders/dpn.py

-    def get_stages(self):
-        return [
-            nn.Identity(),
-            nn.Sequential(
-                self.features[0].conv, self.features[0].bn, self.features[0].act
-            ),
-            nn.Sequential(
-                self.features[0].pool, self.features[1 : self._stage_idxs[0]]
-            ),
-            self.features[self._stage_idxs[0] : self._stage_idxs[1]],
-            self.features[self._stage_idxs[1] : self._stage_idxs[2]],
-            self.features[self._stage_idxs[2] : self._stage_idxs[3]],
-        ]
-
-    def forward(self, x):
-        stages = self.get_stages()
-
-        features = []
-        for i in range(self._depth + 1):
-            x = stages[i](x)
-            if isinstance(x, (list, tuple)):
-                features.append(F.relu(torch.cat(x, dim=1), inplace=True))
-            else:
-                features.append(x)
+    def get_stages(self) -> Dict[int, Sequence[torch.nn.Module]]:
+        return {
+            16: [self.features[self._stage_idxs[1] : self._stage_idxs[2]]],
+            32: [self.features[self._stage_idxs[2] : self._stage_idxs[3]]],
+        }
+
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        features = [x]
+
+        if self._depth >= 1:
+            x = self.features[0].conv(x)
+            x = self.features[0].bn(x)
+            x = self.features[0].act(x)
+            features.append(x)
+
+        if self._depth >= 2:
+            x = self.features[0].pool(x)
+            x = self.features[1 : self._stage_idxs[0]](x)
+            skip = F.relu(torch.cat(x, dim=1), inplace=True)
+            features.append(skip)
+
+        if self._depth >= 3:
+            x = self.features[self._stage_idxs[0] : self._stage_idxs[1]](x)
+            skip = F.relu(torch.cat(x, dim=1), inplace=True)
+            features.append(skip)
+
+        if self._depth >= 4:
+            x = self.features[self._stage_idxs[1] : self._stage_idxs[2]](x)
+            skip = F.relu(torch.cat(x, dim=1), inplace=True)
+            features.append(skip)
+
+        if self._depth >= 5:
+            x = self.features[self._stage_idxs[2] : self._stage_idxs[3]](x)
+            features.append(x)


This PR refactors the self.get_stages method, which previously used a for loop to return multi-scale features. Now, the self.forward method uses if-else statements to handle different feature scales.
What prompted this change? Is it intended to better support more complex models?

Hey @brianhou0208! This is prompted by compatibility with Torch script/export. I'm not sure if it's easier to read, but it is still not that complicated and is very explicit.

qubvel · 2025-01-15T11:24:30Z

@adamjstewart would you like to have a look? 😄 otherwise it's ready to be merged.

The next one will move all encoders to the hf-hub for faster loading and download stats. Actually I moved them already and just need to update URLs and the way of loading

qubvel · 2025-01-15T11:32:36Z

btw, additionally tested all models and encoders against the main branch with the following script to ensure output match and weight can be loaded with no issues

import os
import torch
import segmentation_models_pytorch as smp

from tqdm import tqdm

TMP_FOLDER = "tmp-model"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def run_on_main():
    
    encoder_names = sorted(smp.encoders.get_encoder_names())

    for encoder_name in tqdm(encoder_names):

        model = smp.Unet(encoder_name, decoder_channels=[4, 4, 4, 4, 4], encoder_weights=None)
        model = model.eval().to(DEVICE)
        sample = torch.randn(2, 3, 256, 256).to(DEVICE)

        with torch.no_grad():
            output = model(sample)

        os.makedirs(os.path.join(TMP_FOLDER, encoder_name), exist_ok=True)
        torch.save(sample, os.path.join(TMP_FOLDER, encoder_name, "input.pth"))
        torch.save(output, os.path.join(TMP_FOLDER, encoder_name, "output.pth"))
        torch.save(model.state_dict(), os.path.join(TMP_FOLDER, encoder_name, "state_dict.pth"))


def run_on_branch():
    
    encoder_names = os.listdir(TMP_FOLDER)

    for encoder_name in tqdm(encoder_names):

        sample = torch.load(os.path.join(TMP_FOLDER, encoder_name, "input.pth"), weights_only=True)
        expected_output = torch.load(os.path.join(TMP_FOLDER, encoder_name, "output.pth"), weights_only=True)
        state_dict = torch.load(os.path.join(TMP_FOLDER, encoder_name, "state_dict.pth"), weights_only=True)
        
        model = smp.Unet(encoder_name, decoder_channels=[4, 4, 4, 4, 4], encoder_weights=None).eval().to(DEVICE)
        try:
            model.load_state_dict(state_dict)
        except Exception as e:
            print(f"Error loading state dict for {encoder_name}: {e}")
            raise e

        with torch.no_grad():
            output = model(sample)

        if not torch.allclose(output, expected_output):
            diff = torch.abs(output - expected_output).max().item()
            print(f"Encoder {encoder_name} has different output with max diff {diff:.6f}")

if __name__ == "__main__":

    import git
    repo = git.Repo(".")

    if repo.active_branch.name == "main":
        print("\n--- Running on main branch ---\n")
        run_on_main()
    else:
        print(f"\n--- Running on {repo.active_branch.name} branch ---\n ")
        run_on_branch()

adamjstewart

This PR is a bit too big for me to properly review, but I added a few comments on things that caught my eye. Thanks for adding more type hints!

adamjstewart · 2025-01-15T11:37:18Z

segmentation_models_pytorch/decoders/fpn/decoder.py

        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:


Technically List is a bit too strict, it could be any collections.abc.Sequence. This includes things like tuples. Likely true for a lot of other places in the code base as well.

adamjstewart · 2025-01-15T11:40:14Z

segmentation_models_pytorch/decoders/pan/decoder.py

@@ -220,7 +247,7 @@ def __init__(
                upscale_mode=upscale_mode,
            )

-    def forward(self, *features):
+    def forward(self, features: List[torch.Tensor]) -> torch.Tensor:


This is a pretty big change. You went from model(x1, x2, x3) to model([x1, x2, x3]). Is this intentional, or was this an accidental type hint change?

This was an intentional change for torchscript compatibility. Yeah, while the model interface does not change, decoder interface has been changed, so might break smth for those who use building blocks

I'm fine with that if it's only the building blocks that changed and not the outward-facing encoders/decoders.

I hope it will be fine as well. While Decoders themselves are not private, they are not advertised as a public API. The main use case is the model API, which is fully backward compatible.

adamjstewart · 2025-01-15T11:40:41Z

segmentation_models_pytorch/decoders/pspnet/decoder.py

+    def __init__(
+        self,
+        in_channels: int,
+        sizes: Tuple[int, ...] = (1, 2, 3, 6),


Would any collections.abc.Sequence be valid here?

segmentation_models_pytorch/decoders/pspnet/decoder.py

segmentation_models_pytorch/encoders/__init__.py

adamjstewart · 2025-01-15T11:45:31Z

segmentation_models_pytorch/encoders/densenet.py

+            x = self.features.transition3.pool(x)
+            x = self.features.denseblock4(x)
+            x = self.features.norm5(x)
+            features.append(x)


I honestly prefer the for-loop here, but I'm guessing that makes it not possible to compile. What happens if depth > 5? Is there an assert statement to prevent that invalid input?

Add validation in d121fec

Yes, more like torhcscript limitations, I left loop in a few places, but it's not possible to loop over indexes, only over layers + break is not supported in torchscript

brianhou0208 · 2025-01-15T12:09:36Z

I also agree with what @adamjstewart said.

Maybe it can be divided into multiple PR for easier review, such as type hints check or delete timm- encoders.

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

qubvel · 2025-01-15T13:00:10Z

@brianhou0208, thanks for the review! typehint corrections were needed for torchscript as well.. so yeah.. not happy PR become so huge either

qubvel · 2025-01-15T13:01:06Z

believe in previously improved tests and my additional testing against main branch 🤞

qubvel · 2025-01-15T13:02:43Z

Agree with @adamjstewart, typhints are not ideal, and can be improved further, but out of scope of this PR, so gonna merge them as is

qubvel added 25 commits January 12, 2025 02:59

Move tests

643c0b6

Add compile test for encoders (to be optimized)

7a937ab

densnet

9a7c768

dpn

34b8533

efficientnet

a3618fa

inceptionresnetv2

e3f6c70

inceptionv4

20b28be

mix-transformer

d996165

mobilenet

9e38154

mobileone

c6e5d53

resnet

5a76722

senet

36d056b

vgg

aefcfd4

xception

e9628bf

Deprecate timm- encoders, remap to tu- most of them

70262e5

Add tiny encoders and compile mark

0b0b1c4

Add conftest

4c11682

Fix features

70168b4

Merge branch 'main' into torch-compile-export

8aed7ef

Add triggering compile tests on diff

50c40d1

Remove marks

0764d5e

Add test_compile stage to CI

7cab4be

Update requirements

2622e0e

Update makefile

e12ee8d

Update get_stages

da0cd19

qubvel added 4 commits January 13, 2025 20:43

Fix weight loading for deprecate encoders

7752969

Fix weight loading for mobilenetv3

409b820

Format

ae3cb8a

Add compile test for models

ff278c9

brianhou0208 reviewed Jan 14, 2025

View reviewed changes

qubvel added 16 commits January 14, 2025 14:22

Make from_pretrained strict by default

31bee79

Fix DeepLabV3 BC

556b3aa

Fix scripting for encoders

f70d861

Refactor test do not skip

ead24b4

Fix encoders (mobilenet, inceptionv4)

d44509a

Update encoders table

b2c13f1

Fix export test

73809e3

Fix docs

bc1319e

Update warning

d25dd47

Move pretrained settings

4f3b37e

Add BC for timm- encoders

06199b0

Fixing table

51e0a67

Update compile test

524bcae

Change compile backend to eager

a2b97d8

Update docs

17a4b70

Fixup

20564f2

adamjstewart mentioned this pull request Jan 15, 2025

Remove efficientnet-pytorch dependency #1018

Closed

adamjstewart reviewed Jan 15, 2025

View reviewed changes

Fix batchnorm typo

5bbb1db

qubvel and others added 3 commits January 15, 2025 12:33

Add depth validation

d121fec

Update segmentation_models_pytorch/encoders/__init__.py

7bb9d37

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

Style

da24de9

qubvel merged commit 456871a into main Jan 15, 2025
17 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix torch compile, script, export #1031

Fix torch compile, script, export #1031

qubvel commented Jan 13, 2025 •

edited

Loading

codecov bot commented Jan 13, 2025 •

edited

Loading

brianhou0208 Jan 14, 2025

qubvel Jan 14, 2025 •

edited

Loading

qubvel commented Jan 15, 2025 •

edited

Loading

qubvel commented Jan 15, 2025

adamjstewart left a comment

adamjstewart Jan 15, 2025

adamjstewart Jan 15, 2025

qubvel Jan 15, 2025

adamjstewart Jan 15, 2025

qubvel Jan 15, 2025

adamjstewart Jan 15, 2025

adamjstewart Jan 15, 2025

qubvel Jan 15, 2025

qubvel Jan 15, 2025 •

edited

Loading

brianhou0208 commented Jan 15, 2025 •

edited

Loading

qubvel commented Jan 15, 2025

qubvel commented Jan 15, 2025

qubvel commented Jan 15, 2025

Fix torch compile, script, export #1031

Fix torch compile, script, export #1031

Conversation

qubvel commented Jan 13, 2025 • edited Loading

codecov bot commented Jan 13, 2025 • edited Loading

Codecov Report

Choose a reason for hiding this comment

qubvel Jan 14, 2025 • edited Loading

Choose a reason for hiding this comment

qubvel commented Jan 15, 2025 • edited Loading

qubvel commented Jan 15, 2025

adamjstewart left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

qubvel Jan 15, 2025 • edited Loading

Choose a reason for hiding this comment

brianhou0208 commented Jan 15, 2025 • edited Loading

qubvel commented Jan 15, 2025

qubvel commented Jan 15, 2025

qubvel commented Jan 15, 2025

qubvel commented Jan 13, 2025 •

edited

Loading

codecov bot commented Jan 13, 2025 •

edited

Loading

qubvel Jan 14, 2025 •

edited

Loading

qubvel commented Jan 15, 2025 •

edited

Loading

qubvel Jan 15, 2025 •

edited

Loading

brianhou0208 commented Jan 15, 2025 •

edited

Loading