Repository: Project-MONAI/MONAI
File: tests/test_swin_unetr.py (Python)

Code before:
from monai.networks import eval_mode
from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, filter_swinunetr
from monai.networks.utils import copy_model_state
from monai.utils import optional_import
from tests.utils import (
    SkipIfBeforePyTorchVersion,
    assert_allclose,
    skip_if_downloading_fails,
    skip_if_no_cuda,
    skip_if_quick,
    testing_data_config,
)

einops, has_einops = optional_import("einops")

TEST_CASE_SWIN_UNETR = []
case_idx = 0
test_merging_mode = ["mergingv2", "merging", PatchMerging, PatchMergingV2]
for attn_drop_rate in [0.4]:
    for in_channels in [1]:
        for depth in [[2, 1, 1, 1], [1, 2, 1, 1]]:
            for out_channels in [2]:
                for img_size in ((64, 32, 192), (96, 32)):
                    for feature_size in [12]:
                        for norm_name in ["instance"]:
                            for use_checkpoint in [True, False]:
                                test_case = [
                                    {
                                        "spatial_dims": len(img_size),
                                        "in_channels": in_channels,
                                        "out_channels": out_channels,
                                        "img_size": img_size,
                                        "feature_size": feature_size,
                                        "depths": depth,
                                        "norm_name": norm_name,
                                        "attn_drop_rate": attn_drop_rate,
                                        "downsample": test_merging_mode[case_idx % 4],
                                        "use_checkpoint": use_checkpoint,
                                    },
                                    (2, in_channels, *img_size),
                                    (2, out_channels, *img_size),
                                ]
                                case_idx += 1
                                TEST_CASE_SWIN_UNETR.append(test_case)

TEST_CASE_FILTER = [
    [
        {"img_size": (96, 96, 96), "in_channels": 1, "out_channels": 14, "feature_size": 48, "use_checkpoint": True},
        "swinViT.layers1.0.blocks.0.norm1.weight",
        torch.tensor([0.9473, 0.9343, 0.8566, 0.8487, 0.8065, 0.7779, 0.6333, 0.5555]),
    ]

Review comment (suggestion):
Instead of skipping all of the tests for older versions of Pytorch, we can just make the value to loop over here just `[False]` if the Pytorch version is <=1.10...

Code after:
from monai.networks import eval_mode
from monai.networks.nets.swin_unetr import PatchMerging, PatchMergingV2, SwinUNETR, filter_swinunetr
from monai.networks.utils import copy_model_state
from monai.utils import optional_import
from tests.utils import (
    assert_allclose,
    pytorch_after,
    skip_if_downloading_fails,
    skip_if_no_cuda,
    skip_if_quick,
    testing_data_config,
)

einops, has_einops = optional_import("einops")

TEST_CASE_SWIN_UNETR = []
case_idx = 0
test_merging_mode = ["mergingv2", "merging", PatchMerging, PatchMergingV2]
checkpoint_vals = [True, False] if pytorch_after(1, 11) else [False]
for attn_drop_rate in [0.4]:
    for in_channels in [1]:
        for depth in [[2, 1, 1, 1], [1, 2, 1, 1]]:
            for out_channels in [2]:
                for img_size in ((64, 32, 192), (96, 32)):
                    for feature_size in [12]:
                        for norm_name in ["instance"]:
                            for use_checkpoint in checkpoint_vals:
                                test_case = [
                                    {
                                        "spatial_dims": len(img_size),
                                        "in_channels": in_channels,
                                        "out_channels": out_channels,
                                        "img_size": img_size,
                                        "feature_size": feature_size,
                                        "depths": depth,
                                        "norm_name": norm_name,
                                        "attn_drop_rate": attn_drop_rate,
                                        "downsample": test_merging_mode[case_idx % 4],
                                        "use_checkpoint": use_checkpoint,
                                    },
                                    (2, in_channels, *img_size),
                                    (2, out_channels, *img_size),
                                ]
                                case_idx += 1
                                TEST_CASE_SWIN_UNETR.append(test_case)

TEST_CASE_FILTER = [
    [
        {"img_size": (96, 96, 96), "in_channels": 1, "out_channels": 14, "feature_size": 48, "use_checkpoint": True},
        "swinViT.layers1.0.blocks.0.norm1.weight",
        torch.tensor([0.9473, 0.9343, 0.8566, 0.8487, 0.8065, 0.7779, 0.6333, 0.5555]),