-
-
Notifications
You must be signed in to change notification settings - Fork 1.7k
Load model with mismatched sizes #1107
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
Codecov ReportAttention: Patch coverage is
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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)
Allows to load pretrained model with different number of channels
For example:
Reported in: