From 7741301f77fb68c9d89306e362b6b518e8eee0fe Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 9 Jul 2026 13:22:40 +0100 Subject: [PATCH] perf(test): cache SwinUNETR SSL weights instead of re-downloading Signed-off-by: Soumya Snigdha Kundu --- tests/networks/nets/test_swin_unetr.py | 29 +++++++++++++------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/networks/nets/test_swin_unetr.py b/tests/networks/nets/test_swin_unetr.py index ba94aab4f90..534b582c204 100644 --- a/tests/networks/nets/test_swin_unetr.py +++ b/tests/networks/nets/test_swin_unetr.py @@ -12,8 +12,8 @@ from __future__ import annotations import os -import tempfile import unittest +from pathlib import Path from unittest import skipUnless import torch @@ -35,6 +35,8 @@ einops, has_einops = optional_import("einops") +TESTS_PATH = Path(__file__).parents[2] + test_merging_mode = ["mergingv2", "merging", PatchMerging, PatchMergingV2] checkpoint_vals = [True, False] @@ -111,19 +113,18 @@ def test_patch_merging(self): @skip_if_no_cuda def test_filter_swinunetr(self, input_param, key, value): with skip_if_downloading_fails(): - with tempfile.TemporaryDirectory() as tempdir: - file_name = "ssl_pretrained_weights.pth" - data_spec = testing_data_config("models", f"{file_name.split('.', 1)[0]}") - weight_path = os.path.join(tempdir, file_name) - download_url( - data_spec["url"], weight_path, hash_val=data_spec["hash_val"], hash_type=data_spec["hash_type"] - ) - - ssl_weight = torch.load(weight_path, weights_only=True)["model"] - net = SwinUNETR(**input_param) - dst_dict, loaded, not_loaded = copy_model_state(net, ssl_weight, filter_func=filter_swinunetr) - assert_allclose(dst_dict[key][:8], value, atol=1e-4, rtol=1e-4, type_test=False) - self.assertTrue(len(loaded) == 157 and len(not_loaded) == 2) + file_name = "ssl_pretrained_weights.pth" + data_spec = testing_data_config("models", f"{file_name.split('.', 1)[0]}") + weight_path = os.path.join(TESTS_PATH, "testing_data", f"temp_{file_name}") + download_url( + data_spec["url"], weight_path, hash_val=data_spec["hash_val"], hash_type=data_spec["hash_type"] + ) + + ssl_weight = torch.load(weight_path, weights_only=True)["model"] + net = SwinUNETR(**input_param) + dst_dict, loaded, not_loaded = copy_model_state(net, ssl_weight, filter_func=filter_swinunetr) + assert_allclose(dst_dict[key][:8], value, atol=1e-4, rtol=1e-4, type_test=False) + self.assertTrue(len(loaded) == 157 and len(not_loaded) == 2) if __name__ == "__main__":