Migrate tests from TorchScript to torch.export

Add new test suites for export_utils, convert_to_export, and the
export_checkpoint bundle CLI.  Migrate ~60 existing test files from
test_script (torch.jit.script) to test_export (torch.export.export),
removing TorchScript-specific skip decorators where no longer needed.

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

Author: aymuos15

tests/apps/detection/networks/test_retinanet.py

@@ -20,7 +20,7 @@
 from monai.networks import eval_mode
 from monai.networks.nets import resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200
 from monai.utils import ensure_tuple, optional_import
-from tests.test_utils import dict_product, skip_if_quick, test_onnx_save, test_script_save
+from tests.test_utils import dict_product, skip_if_quick, test_export_save, test_onnx_save
 
 _, has_torchvision = optional_import("torchvision")
 _, has_onnxruntime = optional_import("onnxruntime")
@@ -92,7 +92,7 @@
 MODEL_LIST = [resnet10, resnet18, resnet34, resnet50, resnet101, resnet152, resnet200]
 
 TEST_CASES = [[params["model"], *params["case"]] for params in dict_product(model=MODEL_LIST, case=CASE_LIST)]
-TEST_CASES_TS = [[params["model"], *params["case"]] for params in dict_product(model=MODEL_LIST, case=[TEST_CASE_1])]
+TEST_CASES_EXPORT = [[params["model"], *params["case"]] for params in dict_product(model=MODEL_LIST, case=[TEST_CASE_1])]
 
 
 @unittest.skipUnless(has_torchvision, "Requires torchvision")
@@ -136,18 +136,18 @@ def test_retina_shape(self, model, input_param, input_shape):
             self.assertEqual(tuple(cc.shape for cc in result[net.cls_key]), expected_cls_shape)
             self.assertEqual(tuple(cc.shape for cc in result[net.box_reg_key]), expected_box_shape)
 
-    @parameterized.expand(TEST_CASES_TS)
-    def test_script(self, model, input_param, input_shape):
+    @parameterized.expand(TEST_CASES_EXPORT)
+    def test_export(self, model, input_param, input_shape):
         try:
-            idx = int(self.id().split("test_script_")[-1])
+            idx = int(self.id().split("test_export_")[-1])
         except BaseException:
             idx = 0
         idx %= 3
-        # test whether support torchscript
+        # test whether support torch.export
         data = torch.randn(input_shape)
         backbone = model(**input_param)
         if idx == 0:
-            test_script_save(backbone, data)
+            test_export_save(backbone, data)
             return
         feature_extractor = resnet_fpn_feature_extractor(
             backbone=backbone,
@@ -157,7 +157,7 @@ def test_script(self, model, input_param, input_shape):
             returned_layers=[1, 2],
         )
         if idx == 1:
-            test_script_save(feature_extractor, data)
+            test_export_save(feature_extractor, data)
             return
         net = RetinaNet(
             spatial_dims=input_param["spatial_dims"],
@@ -167,17 +167,17 @@ def test_script(self, model, input_param, input_shape):
             size_divisible=32,
         )
         if idx == 2:
-            test_script_save(net, data)
+            test_export_save(net, data)
 
-    @parameterized.expand(TEST_CASES_TS)
+    @parameterized.expand(TEST_CASES_EXPORT)
     @unittest.skipUnless(has_onnxruntime, "onnxruntime not installed")
     def test_onnx(self, model, input_param, input_shape):
         try:
             idx = int(self.id().split("test_onnx_")[-1])
         except BaseException:
             idx = 0
         idx %= 3
-        # test whether support torchscript
+        # test whether support torch.export
         data = torch.randn(input_shape)
         backbone = model(**input_param)
         if idx == 0:

tests/apps/detection/networks/test_retinanet_detector.py

@@ -21,7 +21,7 @@
 from monai.apps.detection.utils.anchor_utils import AnchorGeneratorWithAnchorShape
 from monai.networks import eval_mode, train_mode
 from monai.utils import optional_import
-from tests.test_utils import skip_if_quick, test_script_save
+from tests.test_utils import skip_if_quick, test_export_save
 
 _, has_torchvision = optional_import("torchvision")
 
@@ -89,7 +89,7 @@
 TEST_CASES = []
 TEST_CASES = [TEST_CASE_1, TEST_CASE_2, TEST_CASE_2_A]
 
-TEST_CASES_TS = [TEST_CASE_1]
+TEST_CASES_EXPORT = [TEST_CASE_1]
 
 
 class NaiveNetwork(torch.nn.Module):
@@ -183,9 +183,9 @@ def test_naive_retina_detector_shape(self, input_param, input_shape):
             targets = [one_target] * len(input_data)
             result = detector.forward(input_data, targets)
 
-    @parameterized.expand(TEST_CASES_TS)
-    def test_script(self, input_param, input_shape):
-        # test whether support torchscript
+    @parameterized.expand(TEST_CASES_EXPORT)
+    def test_export(self, input_param, input_shape):
+        # test whether support torch.export
         returned_layers = [1]
         anchor_generator = AnchorGeneratorWithAnchorShape(
             feature_map_scales=(1, 2), base_anchor_shapes=((8,) * input_param["spatial_dims"],)
@@ -195,7 +195,7 @@ def test_script(self, input_param, input_shape):
         )
         with eval_mode(detector):
             input_data = torch.randn(input_shape)
-            test_script_save(detector.network, input_data)
+            test_export_save(detector.network, input_data)
 
 
 if __name__ == "__main__":

tests/apps/detection/utils/test_anchor_box.py

@@ -18,7 +18,7 @@
 
 from monai.apps.detection.utils.anchor_utils import AnchorGenerator, AnchorGeneratorWithAnchorShape
 from monai.utils import optional_import
-from tests.test_utils import assert_allclose, test_script_save
+from tests.test_utils import assert_allclose, test_export_save
 
 _, has_torchvision = optional_import("torchvision")
 
@@ -67,20 +67,20 @@ def test_anchor_2d(self, input_param, image_shape, feature_maps_shapes):
             assert_allclose(a, a_f, type_test=True, device_test=False, atol=0.1)
 
     @parameterized.expand(TEST_CASES_2D)
-    def test_script_2d(self, input_param, image_shape, feature_maps_shapes):
-        # test whether support torchscript
+    def test_export_2d(self, input_param, image_shape, feature_maps_shapes):
+        # test whether support torch.export
         anchor = AnchorGenerator(**input_param, indexing="xy")
         images = torch.rand(image_shape)
         feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes)
-        test_script_save(anchor, images, feature_maps)
+        test_export_save(anchor, images, feature_maps)
 
     @parameterized.expand(TEST_CASES_SHAPE_3D)
-    def test_script_3d(self, input_param, image_shape, feature_maps_shapes):
-        # test whether support torchscript
+    def test_export_3d(self, input_param, image_shape, feature_maps_shapes):
+        # test whether support torch.export
         anchor = AnchorGeneratorWithAnchorShape(**input_param, indexing="ij")
         images = torch.rand(image_shape)
         feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes)
-        test_script_save(anchor, images, feature_maps)
+        test_export_save(anchor, images, feature_maps)
 
 
 if __name__ == "__main__":

tests/bundle/test_bundle_export_checkpoint.py

@@ -0,0 +1,110 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+import unittest
+from pathlib import Path
+
+from parameterized import parameterized
+
+from monai.bundle import ConfigParser
+from monai.data import load_exported_program
+from monai.networks import save_state
+from tests.test_utils import command_line_tests, skip_if_windows
+
+TESTS_PATH = Path(__file__).parents[1]
+
+# key_in_ckpt
+TEST_CASE_1 = [""]
+TEST_CASE_2 = ["model"]
+
+
+@skip_if_windows
+class TestExportCheckpoint(unittest.TestCase):
+
+    def setUp(self):
+        self._orig_cuda_env = os.environ.get("CUDA_VISIBLE_DEVICES")
+
+    def tearDown(self):
+        if self._orig_cuda_env is not None:
+            os.environ["CUDA_VISIBLE_DEVICES"] = self._orig_cuda_env
+        else:
+            os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_export(self, key_in_ckpt):
+        meta_file = os.path.join(TESTS_PATH, "testing_data", "metadata.json")
+        config_file = os.path.join(TESTS_PATH, "testing_data", "inference.json")
+        with tempfile.TemporaryDirectory() as tempdir:
+            def_args = {"meta_file": "will be replaced by `meta_file` arg"}
+            def_args_file = os.path.join(tempdir, "def_args.yaml")
+
+            ckpt_file = os.path.join(tempdir, "model.pt")
+            pt2_file = os.path.join(tempdir, "model.pt2")
+
+            parser = ConfigParser()
+            parser.export_config_file(config=def_args, filepath=def_args_file)
+            parser.read_config(config_file)
+            net = parser.get_parsed_content("network_def")
+            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]",
+            ]
+            command_line_tests(cmd)
+            self.assertTrue(os.path.exists(pt2_file))
+
+            _, _metadata, extra_files = load_exported_program(
+                pt2_file, more_extra_files=["inference.json", "def_args.json"]
+            )
+            self.assertIn("meta_file", json.loads(extra_files["def_args.json"]))
+            self.assertIn("network_def", json.loads(extra_files["inference.json"]))
+
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
+    def test_default_value(self, key_in_ckpt):
+        config_file = os.path.join(TESTS_PATH, "testing_data", "inference.json")
+        with tempfile.TemporaryDirectory() as tempdir:
+            def_args = {"meta_file": "will be replaced by `meta_file` arg"}
+            def_args_file = os.path.join(tempdir, "def_args.yaml")
+            ckpt_file = os.path.join(tempdir, "models", "model.pt")
+            pt2_file = os.path.join(tempdir, "models", "model.pt2")
+
+            parser = ConfigParser()
+            parser.export_config_file(config=def_args, filepath=def_args_file)
+            parser.read_config(config_file)
+            net = parser.get_parsed_content("network_def")
+            save_state(src=net if key_in_ckpt == "" else {key_in_ckpt: net}, path=ckpt_file)
+
+            # 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]",
+            ]
+            command_line_tests(cmd)
+            self.assertTrue(os.path.exists(pt2_file))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/data/meta_tensor/test_meta_tensor.py

@@ -224,21 +224,22 @@ def test_get_set_meta_fns(self):
         self.assertTrue(get_track_meta())
 
     @parameterized.expand(TEST_DEVICES)
-    def test_torchscript(self, device):
+    def test_export(self, device):
         shape = (1, 3, 10, 8)
         im, _ = self.get_im(shape, device=device)
         conv = torch.nn.Conv2d(im.shape[1], 5, 3)
         conv.to(device)
         im_conv = conv(im)
-        traced_fn = torch.jit.trace(conv, im.as_tensor())
+        exported = torch.export.export(conv, args=(im.as_tensor(),))
         # save it, load it, use it
         with tempfile.TemporaryDirectory() as tmp_dir:
-            fname = os.path.join(tmp_dir, "im.pt")
-            torch.jit.save(traced_fn, f=fname)
-            traced_fn = torch.jit.load(fname)
-            out = traced_fn(im)
+            fname = os.path.join(tmp_dir, "im.pt2")
+            torch.export.save(exported, fname)
+            loaded = torch.export.load(fname)
+            out = loaded.module()(im.as_tensor())
             self.assertIsInstance(out, torch.Tensor)
-        self.check(out, im_conv, ids=False)
+        # exported module returns plain Tensor, compare values only
+        assert_allclose(out, im_conv)
 
     def test_pickling(self):
         m, _ = self.get_im()

tests/data/test_export_utils.py

@@ -0,0 +1,104 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import tempfile
+import unittest
+
+import torch
+
+from monai.config import get_config_values
+from monai.data import load_exported_program, save_exported_program
+from monai.utils import ExportMetadataKeys
+
+
+class TestModule(torch.nn.Module):
+    __test__ = False
+
+    def forward(self, x):
+        return x + 10
+
+
+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),))
+        with tempfile.TemporaryDirectory() as tempdir:
+            save_exported_program(ep, f"{tempdir}/test")
+            self.assertTrue(os.path.isfile(f"{tempdir}/test.pt2"))
+
+    def test_save_exported_program_ext(self):
+        """Save an exported program to a file with custom extension."""
+        ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),))
+        with tempfile.TemporaryDirectory() as tempdir:
+            save_exported_program(ep, f"{tempdir}/test.zip")
+            self.assertTrue(os.path.isfile(f"{tempdir}/test.zip"))
+
+    def test_save_with_metadata(self):
+        """Save an exported program with metadata to a file."""
+        ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),))
+        test_metadata = {"foo": [1, 2], "bar": "string"}
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            save_exported_program(ep, f"{tempdir}/test", meta_values=test_metadata)
+            self.assertTrue(os.path.isfile(f"{tempdir}/test.pt2"))
+
+    def test_load_exported_program(self):
+        """Save then load an exported program with no extra metadata."""
+        ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),))
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            save_exported_program(ep, f"{tempdir}/test")
+            loaded_ep, meta, extra_files = load_exported_program(f"{tempdir}/test.pt2")
+
+        del meta[ExportMetadataKeys.TIMESTAMP.value]
+        self.assertEqual(meta, get_config_values())
+        self.assertEqual(extra_files, {})
+
+        # Verify the loaded program produces the same output
+        result = loaded_ep.module()(torch.tensor(5.0))
+        self.assertEqual(result.item(), 15.0)
+
+    def test_load_with_metadata(self):
+        """Save then load an exported program with metadata."""
+        ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),))
+        test_metadata = {"foo": [1, 2], "bar": "string"}
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            save_exported_program(ep, f"{tempdir}/test", meta_values=test_metadata)
+            _, meta, extra_files = load_exported_program(f"{tempdir}/test.pt2")
+
+        del meta[ExportMetadataKeys.TIMESTAMP.value]
+
+        test_compare = get_config_values()
+        test_compare.update(test_metadata)
+        self.assertEqual(meta, test_compare)
+        self.assertEqual(extra_files, {})
+
+    def test_save_load_more_extra_files(self):
+        """Save then load extra file data from an exported program."""
+        ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),))
+        test_metadata = {"foo": [1, 2], "bar": "string"}
+        more_extra_files = {"test.txt": "This is test data"}
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            save_exported_program(ep, f"{tempdir}/test", meta_values=test_metadata, more_extra_files=more_extra_files)
+            self.assertTrue(os.path.isfile(f"{tempdir}/test.pt2"))
+
+            _, _, loaded_extra_files = load_exported_program(f"{tempdir}/test.pt2", more_extra_files=("test.txt",))
+            self.assertEqual(more_extra_files["test.txt"], loaded_extra_files["test.txt"])
+
+
+if __name__ == "__main__":
+    unittest.main()