Load model with mismatched sizes

[qubvel]

Allows to load pretrained model with different number of channels

For example:

import segmentation_models_pytorch as smp

model = smp.from_pretrained("smp-hub/segformer-b2-1024x1024-city-160k", classes=5, strict=False)

Reported in:

• KeyError when using Segformer architecture in segmentation_models_pytorch #1019

[codecov]

Codecov Report

Attention: Patch coverage is 70.37037% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
segmentation_models_pytorch/base/model.py	70.37%	8 Missing ⚠️

Files with missing lines	Coverage Δ
segmentation_models_pytorch/base/model.py	83.56% <70.37%> (+5.22%)	⬆️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull Request Overview

This PR adds support for loading pretrained models with mismatched channel sizes by filtering out incompatible weights during state dict loading. Key changes include:

• Adding a test in tests/test_base.py to verify mismatched keys are handled correctly.
• Updating load_state_dict in segmentation_models_pytorch/base/model.py to filter mismatched weights and issue a warning.
• Adjusting the inference notebook to use up-to-date installation instructions.

Reviewed Changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 1 comment.

File	Description
tests/test_base.py	Adds a test to ensure models load with mismatched keys while only specific layers are adjusted.
segmentation_models_pytorch/base/model.py	Modifies load_state_dict to filter out mismatched weights and warn the user.
examples/segformer_inference_pretrained.ipynb	Updates installation instructions to ensure the latest library versions are used.

Files not reviewed (1)

• docs/save_load.rst: Language not supported

Comments suppressed due to low confidence (2)

segmentation_models_pytorch/base/model.py:138

• [nitpick] The warning message uses all caps and strong language, which may not be clear; consider rephrasing it to a more neutral and descriptive message (e.g., 'Mismatched key shapes detected; please retrain the model to update these layers:').

text = f"\n\n !!!!!! Mismatched keys !!!!!!\n\nYou should TRAIN the model to use it:\n{str_keys}\n"

segmentation_models_pytorch/base/model.py:139

• [nitpick] Using 'stacklevel=-1' may not provide the expected context for the warning; consider using a more conventional stacklevel (such as 2) to point to the relevant caller.

warnings.warn(text, stacklevel=-1)