CI and linting fix

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/data/__init__.py

@@ -47,6 +47,7 @@
     load_decathlon_datalist,
     load_decathlon_properties,
 )
+from .export_utils import load_exported_program, save_exported_program
 from .folder_layout import FolderLayout, FolderLayoutBase
 from .grid_dataset import GridPatchDataset, PatchDataset, PatchIter, PatchIterd
 from .image_dataset import ImageDataset
@@ -77,7 +78,6 @@
 from .test_time_augmentation import TestTimeAugmentation
 from .thread_buffer import ThreadBuffer, ThreadDataLoader
 from .torchscript_utils import load_net_with_metadata, save_net_with_metadata
-from .export_utils import load_exported_program, save_exported_program
 from .utils import (
     affine_to_spacing,
     compute_importance_map,

monai/networks/nets/dints.py

@@ -376,8 +376,7 @@ def __init__(
         # Pre-compute node activation flags as Python booleans for torch.export compatibility.
         # NOTE: node_a must not be mutated after construction.
         self._node_flags: list[list[bool]] = [
-            [bool(self.node_a[b, d]) for d in range(self.num_depths)]
-            for b in range(self.node_a.shape[0])
+            [bool(self.node_a[b, d]) for d in range(self.num_depths)] for b in range(self.node_a.shape[0])
         ]
 
         # define stem operations for every block
@@ -675,8 +674,7 @@ def __init__(
         # NOTE: arch_code_a must not be mutated after construction; this class
         # is only used at inference/re-training time, not during architecture search.
         self._active_flags: list[list[bool]] = [
-            [bool(self.arch_code_a[b, r]) for r in range(self.arch_code_a.shape[1])]
-            for b in range(self.num_blocks)
+            [bool(self.arch_code_a[b, r]) for r in range(self.arch_code_a.shape[1])] for b in range(self.num_blocks)
         ]
 
     def forward(self, x: list[torch.Tensor]) -> list[torch.Tensor]:

tests/bundle/test_bundle_export_checkpoint.py

@@ -33,7 +33,6 @@
 
 @skip_if_windows
 class TestExportCheckpoint(unittest.TestCase):
-
     def setUp(self):
         self._orig_cuda_env = os.environ.get("CUDA_VISIBLE_DEVICES")
 
@@ -61,14 +60,26 @@ def test_export(self, key_in_ckpt):
             save_state(src=net if key_in_ckpt == "" else {key_in_ckpt: net}, path=ckpt_file)
 
             cmd = [
-                "coverage", "run", "-m", "monai.bundle", "export_checkpoint",
-                "network_def", "--filepath", pt2_file,
-                "--meta_file", meta_file,
-                "--config_file", f"['{config_file}','{def_args_file}']",
-                "--ckpt_file", ckpt_file,
-                "--key_in_ckpt", key_in_ckpt,
-                "--args_file", def_args_file,
-                "--input_shape", "[1, 1, 96, 96, 96]",
+                "coverage",
+                "run",
+                "-m",
+                "monai.bundle",
+                "export_checkpoint",
+                "network_def",
+                "--filepath",
+                pt2_file,
+                "--meta_file",
+                meta_file,
+                "--config_file",
+                f"['{config_file}','{def_args_file}']",
+                "--ckpt_file",
+                ckpt_file,
+                "--key_in_ckpt",
+                key_in_ckpt,
+                "--args_file",
+                def_args_file,
+                "--input_shape",
+                "[1, 1, 96, 96, 96]",
             ]
             command_line_tests(cmd)
             self.assertTrue(os.path.exists(pt2_file))
@@ -96,11 +107,19 @@ def test_default_value(self, key_in_ckpt):
 
             # check with default value
             cmd = [
-                "coverage", "run", "-m", "monai.bundle", "export_checkpoint",
-                "--key_in_ckpt", key_in_ckpt,
-                "--config_file", config_file,
-                "--bundle_root", tempdir,
-                "--input_shape", "[1, 1, 96, 96, 96]",
+                "coverage",
+                "run",
+                "-m",
+                "monai.bundle",
+                "export_checkpoint",
+                "--key_in_ckpt",
+                key_in_ckpt,
+                "--config_file",
+                config_file,
+                "--bundle_root",
+                tempdir,
+                "--input_shape",
+                "[1, 1, 96, 96, 96]",
             ]
             command_line_tests(cmd)
             self.assertTrue(os.path.exists(pt2_file))

tests/data/test_export_utils.py

@@ -30,7 +30,6 @@ def forward(self, x):
 
 
 class TestExportUtils(unittest.TestCase):
-
     def test_save_exported_program(self):
         """Save an exported program without metadata to a file."""
         ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),))

tests/losses/test_generalized_wasserstein_dice_loss.py

@@ -20,7 +20,7 @@
 import torch.optim as optim
 
 from monai.losses import GeneralizedWassersteinDiceLoss
-from tests.test_utils import test_export_save
+from tests.test_utils import test_script_save
 
 
 class TestGeneralizedWassersteinDiceLoss(unittest.TestCase):
@@ -303,7 +303,7 @@ def test_batch_size_different_samples(self):
                 float(loss_batch[1]), loss_b, places=5, msg=f"Batch loss[1] != loss_b for weighting_mode={w_mode}"
             )
 
-    def test_export(self):
+    def test_script(self):
         target = torch.tensor([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]])
 
         # add another dimension corresponding to the batch (batch size = 1 here)
@@ -312,7 +312,7 @@ def test_export(self):
 
         loss = GeneralizedWassersteinDiceLoss(dist_matrix=np.array([[0.0, 1.0], [1.0, 0.0]]), weighting_mode="default")
 
-        test_export_save(loss, pred_very_good, target)
+        test_script_save(loss, pred_very_good, target)
 
 
 if __name__ == "__main__":

tests/losses/test_masked_loss.py

@@ -19,7 +19,7 @@
 from monai.losses.dice import DiceFocalLoss, DiceLoss
 from monai.losses.spatial_mask import MaskedLoss
 from monai.utils import set_determinism
-from tests.test_utils import test_export_save
+from tests.test_utils import test_script_save
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -73,14 +73,14 @@ def test_ill_opts(self):
             masked = MaskedLoss(loss=dice_loss)
             masked(input=torch.zeros((3, 3, 2, 2)), target=torch.zeros((3, 2, 2, 2)), mask=torch.zeros((3, 3, 2, 2)))
 
-    def test_export(self):
+    def test_script(self):
         input_param, expected_val = TEST_CASES[0]
         size = [3, 3, 5, 5]
         label = torch.randint(low=0, high=2, size=size)
         label = torch.argmax(label, dim=1, keepdim=True)
         pred = torch.randn(size)
         loss = MaskedLoss(**input_param)
-        test_export_save(loss, pred, label)
+        test_script_save(loss, pred, label)
 
 
 if __name__ == "__main__":

tests/networks/blocks/test_fpn_block.py

@@ -21,7 +21,7 @@
 from monai.networks.blocks.feature_pyramid_network import FeaturePyramidNetwork
 from monai.networks.nets.resnet import resnet50
 from monai.utils import optional_import
-from tests.test_utils import test_export_save
+from tests.test_utils import test_script_save
 
 _, has_torchvision = optional_import("torchvision")
 
@@ -55,13 +55,13 @@ def test_fpn_block(self, input_param, input_shape, expected_shape):
         self.assertEqual(result["feat1"].shape, expected_shape[1])
 
     @parameterized.expand(TEST_CASES)
-    def test_export(self, input_param, input_shape, expected_shape):
-        # test whether support torch.export
+    def test_script(self, input_param, input_shape, expected_shape):
+        # test whether support torchscript
         net = FeaturePyramidNetwork(**input_param)
         data = OrderedDict()
         data["feat0"] = torch.rand(input_shape[0])
         data["feat1"] = torch.rand(input_shape[1])
-        test_export_save(net, data)
+        test_script_save(net, data)
 
 
 @unittest.skipUnless(has_torchvision, "Requires torchvision")
@@ -75,11 +75,11 @@ def test_fpn(self, input_param, input_shape, expected_shape):
         self.assertEqual(result["pool"].shape, expected_shape[1])
 
     @parameterized.expand(TEST_CASES2)
-    def test_export(self, input_param, input_shape, expected_shape):
-        # test whether support torch.export
+    def test_script(self, input_param, input_shape, expected_shape):
+        # test whether support torchscript
         net = _resnet_fpn_extractor(backbone=resnet50(), spatial_dims=input_param["spatial_dims"], returned_layers=[1])
         data = torch.rand(input_shape)
-        test_export_save(net, data)
+        test_script_save(net, data)
 
 
 if __name__ == "__main__":

tests/networks/nets/regunet/test_localnet.py

@@ -18,7 +18,7 @@
 
 from monai.networks import eval_mode
 from monai.networks.nets.regunet import LocalNet
-from tests.test_utils import test_export_save
+from tests.test_utils import test_script_save
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -75,11 +75,11 @@ def test_extract_levels(self, input_param, input_shape, expected_shape):
         self.assertEqual(len(net.decode_deconvs), len(input_param["extract_levels"]) - 1)
         self.assertEqual(len(net.decode_convs), len(input_param["extract_levels"]) - 1)
 
-    def test_export(self):
+    def test_script(self):
         input_param, input_shape, _ = TEST_CASE_LOCALNET_2D[0]
         net = LocalNet(**input_param)
         test_data = torch.randn(input_shape)
-        test_export_save(net, test_data)
+        test_script_save(net, test_data)
 
 
 if __name__ == "__main__":

tests/networks/nets/regunet/test_regunet.py

@@ -18,7 +18,7 @@
 
 from monai.networks import eval_mode
 from monai.networks.nets.regunet import RegUNet
-from tests.test_utils import test_export_save
+from tests.test_utils import test_script_save
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -77,11 +77,11 @@ def test_ill_shape(self):
             net = RegUNet(**input_param).to(device)
             net.forward(torch.randn(input_shape).to(device))
 
-    def test_export(self):
+    def test_script(self):
         input_param, input_shape, _ = TEST_CASE_REGUNET_2D[0]
         net = RegUNet(**input_param)
         test_data = torch.randn(input_shape)
-        test_export_save(net, test_data)
+        test_script_save(net, test_data)
 
 
 if __name__ == "__main__":

tests/networks/nets/test_ahnet.py

@@ -19,7 +19,7 @@
 from monai.networks import eval_mode
 from monai.networks.blocks import FCN, MCFCN
 from monai.networks.nets import AHNet
-from tests.test_utils import skip_if_quick, test_pretrained_networks, test_export_save
+from tests.test_utils import skip_if_quick, test_export_save, test_pretrained_networks
 
 device = "cuda" if torch.cuda.is_available() else "cpu"