test(models): update the test suite for the models module

Author: okunator

cellseg_models_pytorch/models/tests/test_models.py

@@ -4,44 +4,55 @@
 from cellseg_models_pytorch.models import MultiTaskUnet, get_model
 
 
+@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
 @pytest.mark.parametrize("style_channels", [None, 32])
 @pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_cppnet_fwdbwd(model_type, style_channels, add_stem_skip):
+def test_cppnet_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
     n_rays = 3
-    x = torch.rand([1, 3, 32, 32])
+    x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="cppnet",
         type=model_type,
+        enc_name=enc_name,
         n_rays=n_rays,
         ntypes=3,
         ntissues=3,
         style_channels=style_channels,
         add_stem_skip=add_stem_skip,
+        enc_pretrain=False,
     )
 
     y = model(x)
     y["stardist_refined"].mean().backward()
 
     assert y["type"].shape == x.shape
-    assert y["stardist_refined"].shape == torch.Size([1, n_rays, 32, 32])
+    assert y["stardist_refined"].shape == torch.Size([1, n_rays, 64, 64])
 
     if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 32, 32])
+        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
 
 
+@pytest.mark.parametrize(
+    "enc_name",
+    [
+        "samvit_base_patch16",
+        "samvit_base_patch16_224",
+        "samvit_huge_patch16",
+        "samvit_large_patch16",
+    ],
+)
 @pytest.mark.parametrize("model_type", ["base", "plus", "small_plus", "small"])
 @pytest.mark.parametrize("style_channels", [None, 32])
-@pytest.mark.parametrize("enc_name", ["sam_vit_b", "sam_vit_h", "sam_vit_l"])
-def test_cellvit_fwdbwd(model_type, style_channels, enc_name):
+def test_cellvit_fwdbwd(enc_name, model_type, style_channels):
     x = torch.rand([1, 3, 32, 32])
     model = get_model(
         name="cellvit",
         type=model_type,
+        enc_name=enc_name,
         ntypes=3,
         ntissues=3,
         style_channels=style_channels,
-        enc_name=enc_name,
         enc_pretrain=False,
     )
     model.freeze_encoder()
@@ -55,18 +66,21 @@ def test_cellvit_fwdbwd(model_type, style_channels, enc_name):
         assert y["sem"].shape == torch.Size([1, 3, 32, 32])
 
 
+@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus", "small_plus", "small"])
 @pytest.mark.parametrize("style_channels", [None, 32])
 @pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_hovernet_fwdbwd(model_type, style_channels, add_stem_skip):
+def test_hovernet_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="hovernet",
         type=model_type,
+        enc_name=enc_name,
         ntypes=3,
         ntissues=3,
         style_channels=style_channels,
         add_stem_skip=add_stem_skip,
+        enc_pretrain=False,
     )
 
     y = model(x)
@@ -78,42 +92,48 @@ def test_hovernet_fwdbwd(model_type, style_channels, add_stem_skip):
         assert y["sem"].shape == torch.Size([1, 3, 64, 64])
 
 
+@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
 @pytest.mark.parametrize("style_channels", [None, 32])
 @pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_stardist_fwdbwd(model_type, style_channels, add_stem_skip):
+def test_stardist_fwdbwd(enc_name, model_type, style_channels, add_stem_skip):
     n_rays = 3
-    x = torch.rand([1, 3, 32, 32])
+    x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="stardist",
         type=model_type,
         n_rays=n_rays,
+        enc_name=enc_name,
         ntypes=3,
         ntissues=3,
         style_channels=style_channels,
         add_stem_skip=add_stem_skip,
+        enc_pretrain=False,
     )
 
     y = model(x)
     y["stardist"].mean().backward()
 
     assert y["type"].shape == x.shape
-    assert y["stardist"].shape == torch.Size([1, n_rays, 32, 32])
+    assert y["stardist"].shape == torch.Size([1, n_rays, 64, 64])
 
     if "sem" in y.keys():
-        assert y["sem"].shape == torch.Size([1, 3, 32, 32])
+        assert y["sem"].shape == torch.Size([1, 3, 64, 64])
 
 
+@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
 @pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_cellpose_fwdbwd(model_type, add_stem_skip):
+def test_cellpose_fwdbwd(enc_name, model_type, add_stem_skip):
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="cellpose",
         type=model_type,
+        enc_name=enc_name,
         ntypes=3,
         ntissues=3,
         add_stem_skip=add_stem_skip,
+        enc_pretrain=False,
     )
 
     y = model(x)
@@ -126,16 +146,19 @@ def test_cellpose_fwdbwd(model_type, add_stem_skip):
         assert y["sem"].shape == torch.Size([1, 3, 64, 64])
 
 
+@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("model_type", ["base", "plus"])
 @pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_cellpose_fwdbwd(model_type, add_stem_skip):
+def test_cellpose_fwdbwd(enc_name, model_type, add_stem_skip):
     x = torch.rand([1, 3, 64, 64])
     model = get_model(
         name="omnipose",
         type=model_type,
+        enc_name=enc_name,
         ntypes=3,
         ntissues=3,
         add_stem_skip=add_stem_skip,
+        enc_pretrain=False,
     )
 
     y = model(x)
@@ -148,8 +171,9 @@ def test_cellpose_fwdbwd(model_type, add_stem_skip):
         assert y["sem"].shape == torch.Size([1, 3, 64, 64])
 
 
+@pytest.mark.parametrize("enc_name", ["resnet18", "samvit_base_patch16"])
 @pytest.mark.parametrize("add_stem_skip", [False, True])
-def test_multitaskunet_fwdbwd(add_stem_skip):
+def test_multitaskunet_fwdbwd(enc_name, add_stem_skip):
     x = torch.rand([1, 3, 64, 64])
     m = MultiTaskUnet(
         decoders=("sem",),
@@ -160,6 +184,8 @@ def test_multitaskunet_fwdbwd(add_stem_skip):
         long_skips={"sem": "unet"},
         dec_params={"sem": None},
         add_stem_skip=add_stem_skip,
+        enc_name=enc_name,
+        enc_pretrain=False,
     )
     y = m(x)
     y["sem"].mean().backward()