Add test

Author: qubvel

tests/test_base.py

@@ -0,0 +1,36 @@
+import torch
+import tempfile
+import segmentation_models_pytorch as smp
+
+import pytest
+
+
+def test_from_pretrained_with_mismatched_keys():
+    orginal_model = smp.Unet(classes=1)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        orginal_model.save_pretrained(temp_dir)
+
+        # we should catch warning here and check if there specific keys there
+        with pytest.warns(UserWarning):
+            restored_model = smp.from_pretrained(temp_dir, classes=2, strict=False)
+
+    assert restored_model.segmentation_head[0].out_channels == 2
+
+    # verify all the weight are the same expect mismatched ones
+    original_state_dict = orginal_model.state_dict()
+    restored_state_dict = restored_model.state_dict()
+
+    expected_mismatched_keys = [
+        "segmentation_head.0.weight",
+        "segmentation_head.0.bias",
+    ]
+    mismatched_keys = []
+    for key in original_state_dict:
+        if key not in expected_mismatched_keys:
+            assert torch.allclose(original_state_dict[key], restored_state_dict[key])
+        else:
+            mismatched_keys.append(key)
+
+    assert len(mismatched_keys) == 2
+    assert sorted(mismatched_keys) == sorted(expected_mismatched_keys)