From 6a06c1747ebadbd0d64e9d62098fba3180f05085 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:11:15 +0000 Subject: [PATCH 01/18] Add save_exported_program and load_exported_program Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/data/__init__.py | 1 + monai/data/export_utils.py | 125 ++++++++++++++++++++++++++++++ tests/data/test_export_utils.py | 133 ++++++++++++++++++++++++++++++++ 3 files changed, 259 insertions(+) create mode 100644 monai/data/export_utils.py create mode 100644 tests/data/test_export_utils.py diff --git a/monai/data/__init__.py b/monai/data/__init__.py index 562289edcb..6e38285520 100644 --- a/monai/data/__init__.py +++ b/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 diff --git a/monai/data/export_utils.py b/monai/data/export_utils.py new file mode 100644 index 0000000000..8e4162f6b5 --- /dev/null +++ b/monai/data/export_utils.py @@ -0,0 +1,125 @@ +# 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 datetime +import json +import os +from collections.abc import Mapping, Sequence +from typing import IO, Any + +import torch + +from monai.config import get_config_values +from monai.data.torchscript_utils import METADATA_FILENAME +from monai.utils import JITMetadataKeys + +__all__ = ["load_exported_program", "save_exported_program"] + + +def save_exported_program( + exported_program: torch.export.ExportedProgram, + filename_prefix_or_stream: str | os.PathLike | IO[bytes], + include_config_vals: bool = True, + append_timestamp: bool = False, + meta_values: Mapping[str, Any] | None = None, + more_extra_files: Mapping[str, Any] | None = None, +) -> None: + """ + Save an ``ExportedProgram`` produced by :func:`torch.export.export` with metadata included + as a JSON file inside the ``.pt2`` archive. + + Examples:: + + import torch + from monai.networks.nets import UNet + + net = UNet(spatial_dims=2, in_channels=1, out_channels=1, channels=[8, 16], strides=[2]) + ep = torch.export.export(net, args=(torch.rand(1, 1, 32, 32),)) + + meta = {"name": "Test UNet", "input_dims": 2} + save_exported_program(ep, "test", meta_values=meta) + + loaded_ep, loaded_meta, _ = load_exported_program("test.pt2") + + Args: + exported_program: an ``ExportedProgram`` returned by :func:`torch.export.export`. + filename_prefix_or_stream: filename or file-like stream object. + If a string filename has no extension it becomes ``.pt2``. + include_config_vals: if True, MONAI, PyTorch, and NumPy versions are included in metadata. + append_timestamp: if True, a timestamp is appended to the filename before the extension. + meta_values: metadata values to store, compatible with JSON serialization. + more_extra_files: additional data items to include in the archive, as strings or UTF-8 encoded bytes. + + Raises: + ValueError: When ``more_extra_files`` uses the reserved metadata filename. + """ + now = datetime.datetime.now() + metadict: dict[str, Any] = {} + + if include_config_vals: + metadict.update(get_config_values()) + metadict[JITMetadataKeys.TIMESTAMP.value] = now.astimezone().isoformat() + + if meta_values is not None: + metadict.update(meta_values) + + json_data = json.dumps(metadict) + + extra_files: dict[str, Any] = {METADATA_FILENAME: json_data} + + if more_extra_files is not None: + if METADATA_FILENAME in more_extra_files: + raise ValueError(f"'{METADATA_FILENAME}' is reserved and cannot be used in more_extra_files.") + extra_files.update(more_extra_files) + + # torch.export.save only accepts str values + extra_files = {k: v.decode() if isinstance(v, bytes) else v for k, v in extra_files.items()} + + if isinstance(filename_prefix_or_stream, (str, os.PathLike)): + filename_prefix_or_stream = str(filename_prefix_or_stream) + filename_no_ext, ext = os.path.splitext(filename_prefix_or_stream) + if ext == "": + ext = ".pt2" + + if append_timestamp: + filename_prefix_or_stream = now.strftime(f"{filename_no_ext}_%Y%m%d%H%M%S{ext}") + else: + filename_prefix_or_stream = filename_no_ext + ext + + torch.export.save(exported_program, filename_prefix_or_stream, extra_files=extra_files) + + +def load_exported_program( + filename_prefix_or_stream: str | os.PathLike | IO[bytes], more_extra_files: Sequence[str] = () +) -> tuple[torch.export.ExportedProgram, dict, dict]: + """ + Load an ``ExportedProgram`` from a ``.pt2`` file and convert the stored JSON metadata back to a dict object. + This will produce an empty dict if the metadata file is not present. + + :func:`torch.export.load` uses pickle, only load files from a source you trust. + + Args: + filename_prefix_or_stream: filename or file-like stream object. + more_extra_files: additional extra file names to load from the archive. + + Returns: + Triple of (ExportedProgram, metadata dict, extra files dict). :func:`torch.export.load` fills the extra files + dict with every other file in the archive, and a requested name that is not present has an empty string. + """ + extra_files: dict[str, Any] = dict.fromkeys(more_extra_files, "") + extra_files[METADATA_FILENAME] = "" + + exported_program = torch.export.load(filename_prefix_or_stream, extra_files=extra_files) + json_data_dict = json.loads(extra_files.pop(METADATA_FILENAME) or "{}") + + return exported_program, json_data_dict, extra_files diff --git a/tests/data/test_export_utils.py b/tests/data/test_export_utils.py new file mode 100644 index 0000000000..1edb1c1722 --- /dev/null +++ b/tests/data/test_export_utils.py @@ -0,0 +1,133 @@ +# 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.data.torchscript_utils import METADATA_FILENAME +from monai.utils import JITMetadataKeys + + +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[JITMetadataKeys.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[JITMetadataKeys.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"]) + + def test_load_without_metadata(self): + """Load a file saved directly with torch.export.save, which has no metadata.""" + ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),)) + + with tempfile.TemporaryDirectory() as tempdir: + torch.export.save(ep, f"{tempdir}/test.pt2") + _, meta, extra_files = load_exported_program(f"{tempdir}/test.pt2") + + self.assertEqual(meta, {}) + self.assertEqual(extra_files, {}) + + def test_reserved_extra_file_name(self): + """The metadata filename cannot be used for extra files.""" + ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),)) + + with tempfile.TemporaryDirectory() as tempdir: + with self.assertRaisesRegex(ValueError, METADATA_FILENAME): + save_exported_program(ep, f"{tempdir}/test", more_extra_files={METADATA_FILENAME: "{}"}) + + def test_bytes_extra_file(self): + """Extra files given as UTF-8 bytes load back as strings.""" + ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),)) + + with tempfile.TemporaryDirectory() as tempdir: + save_exported_program(ep, f"{tempdir}/test", more_extra_files={"test.json": b"{}"}) + _, _, extra_files = load_exported_program(f"{tempdir}/test.pt2", more_extra_files=("test.json",)) + + self.assertEqual(extra_files["test.json"], "{}") + + +if __name__ == "__main__": + unittest.main() From 8e6022f1d7640f659750b966598ca11ec27bae12 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:11:23 +0000 Subject: [PATCH 02/18] Add convert_to_export, verifying the saved file Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/networks/__init__.py | 1 + monai/networks/utils.py | 81 ++++++++++++ tests/networks/test_convert_to_export.py | 157 +++++++++++++++++++++++ 3 files changed, 239 insertions(+) create mode 100644 tests/networks/test_convert_to_export.py diff --git a/monai/networks/__init__.py b/monai/networks/__init__.py index 5a240021d6..4271bbd821 100644 --- a/monai/networks/__init__.py +++ b/monai/networks/__init__.py @@ -14,6 +14,7 @@ from .trt_compiler import trt_compile from .utils import ( add_casts_around_norms, + convert_to_export, convert_to_onnx, convert_to_torchscript, convert_to_trt, diff --git a/monai/networks/utils.py b/monai/networks/utils.py index 83815fc543..ed2189d87e 100644 --- a/monai/networks/utils.py +++ b/monai/networks/utils.py @@ -15,6 +15,7 @@ from __future__ import annotations import io +import itertools import re import tempfile import warnings @@ -54,6 +55,7 @@ "save_state", "convert_to_onnx", "convert_to_torchscript", + "convert_to_export", "convert_to_trt", "meshgrid_ij", "meshgrid_xy", @@ -865,6 +867,85 @@ def convert_to_torchscript( return script_module +def convert_to_export( + model: nn.Module, + filename_or_obj: Any | None = None, + extra_files: dict | None = None, + verify: bool = False, + inputs: Sequence[Any] | None = None, + dynamic_shapes: dict | tuple | None = None, + device: str | torch.device | None = None, + rtol: float = 1e-4, + atol: float = 0.0, + **kwargs, +) -> torch.export.ExportedProgram: + """ + Utility to export a model using :func:`torch.export.export` and optionally save to a ``.pt2`` file, + with optional input/output data verification. + + Without ``dynamic_shapes`` the exported program only accepts inputs with the shapes of ``inputs``. Modules from + :meth:`torch.export.ExportedProgram.module` reject ``.eval()`` and ``.train()``, the model is exported in eval mode. + + Args: + model: source PyTorch model to export. + filename_or_obj: if not None, a file path string to save the exported program. + extra_files: map from filename to string contents to store in the saved archive. + verify: whether to verify that the exported program, reloaded from ``filename_or_obj`` if given, + produces the same outputs as ``model``. + inputs: input test data for export and verification. Should be a sequence of + tensors that map to positional arguments of ``model()``. + dynamic_shapes: dynamic shape specifications passed to :func:`torch.export.export`. + device: target device to verify the model. If None, uses CUDA if available. + rtol: the relative tolerance when comparing outputs. + atol: the absolute tolerance when comparing outputs. + kwargs: additional keyword arguments for :func:`torch.export.export`. + + Returns: + A :class:`torch.export.ExportedProgram` representing the exported model. + + Raises: + ValueError: When ``inputs`` is None. + AssertionError: When ``verify`` is True and the outputs differ in structure or value. + """ + if inputs is None: + raise ValueError("Input data is required for torch.export.export.") + + model.eval() + with torch.no_grad(): + exported = torch.export.export(model, args=tuple(inputs), dynamic_shapes=dynamic_shapes, **kwargs) + + if filename_or_obj is not None: + torch.export.save(exported, filename_or_obj, extra_files=extra_files) + + if verify: + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + verify_inputs = [i.to(device) if isinstance(i, torch.Tensor) else i for i in inputs] + if filename_or_obj is not None: + loaded_module = torch.export.load(filename_or_obj).module() + else: + loaded_module = deepcopy(exported).module() + loaded_module.to(device) + + # the model and the returned program share weights, so the model is moved back after the check + model_device = next(itertools.chain(model.parameters(), model.buffers()), torch.empty(0)).device + model.to(device) + try: + with torch.no_grad(): + set_determinism(seed=0) + torch_out = model(*verify_inputs) + set_determinism(seed=0) + export_out = loaded_module(*verify_inputs) + set_determinism(seed=None) + finally: + model.to(model_device) + + torch.testing.assert_close(export_out, torch_out, rtol=rtol, atol=atol) # type: ignore + + return exported + + def _onnx_trt_compile( onnx_model, min_shape: Sequence[int], diff --git a/tests/networks/test_convert_to_export.py b/tests/networks/test_convert_to_export.py new file mode 100644 index 0000000000..cb908993a6 --- /dev/null +++ b/tests/networks/test_convert_to_export.py @@ -0,0 +1,157 @@ +# 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 +from unittest.mock import patch + +import torch +from parameterized import parameterized + +from monai.networks import convert_to_export +from monai.networks.nets import UNet +from monai.utils.module import pytorch_after +from tests.test_utils import skip_if_no_cuda + + +class _Wrapped(torch.nn.Module): + def __init__(self, wrap): + super().__init__() + self.conv = torch.nn.Conv2d(1, 1, 1) + self.wrap = wrap + + def forward(self, x): + return self.wrap(self.conv(x)) + + +class _KeepsLastOutput(torch.nn.Module): + """Keeps its last training output, a tensor that cannot be deep-copied, as some networks do.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(1, 1, 1) + self.last = None + + def forward(self, x): + out = self.conv(x) + if self.training: + self.last = out + return out + + +def _build_unet(): + return UNet(spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2), num_res_units=0) + + +class TestConvertToExport(unittest.TestCase): + def test_basic_export(self): + """Export a UNet and verify output matches.""" + model = _build_unet() + with tempfile.TemporaryDirectory() as tempdir: + exported = convert_to_export( + model=model, + filename_or_obj=os.path.join(tempdir, "model.pt2"), + verify=True, + inputs=[torch.randn((16, 1, 32, 32), requires_grad=False)], + device="cpu", + rtol=1e-3, + atol=1e-4, + ) + self.assertIsInstance(exported, torch.export.ExportedProgram) + + def test_verify_checks_saved_file(self): + """Verification must use the saved file, so a saved model with other weights fails.""" + model = _build_unet() + inputs = [torch.randn((2, 1, 32, 32))] + other = torch.export.export(_build_unet().eval(), args=tuple(inputs)) + with tempfile.TemporaryDirectory() as tempdir: + with patch("torch.export.load", return_value=other) as load: + with self.assertRaises(AssertionError): + convert_to_export( + model=model, + filename_or_obj=os.path.join(tempdir, "model.pt2"), + verify=True, + inputs=inputs, + device="cpu", + ) + load.assert_called_once() + + @parameterized.expand([[lambda x: {"out": x}], [lambda x: ([x],)]]) + def test_verify_compares_nested_outputs(self, wrap): + """Verification must reach tensors inside dicts and nested sequences.""" + inputs = [torch.randn((2, 1, 8, 8))] + other = torch.export.export(_Wrapped(wrap).eval(), args=tuple(inputs)) + with tempfile.TemporaryDirectory() as tempdir: + with patch("torch.export.load", return_value=other): + with self.assertRaises(AssertionError): + convert_to_export( + _Wrapped(wrap), os.path.join(tempdir, "model.pt2"), verify=True, inputs=inputs, device="cpu" + ) + + def test_verify_model_after_training_step(self): + """Verification must not need to copy the model, which can hold tensors from earlier training steps.""" + net = _KeepsLastOutput() + data = torch.randn(2, 1, 8, 8) + net(data).sum().backward() + convert_to_export(net, verify=True, inputs=[data], device="cpu") + + @skip_if_no_cuda + def test_verify_leaves_devices_alone(self): + """Verifying on another device moves neither the model nor the returned program.""" + model = _build_unet() + exported = convert_to_export(model, verify=True, inputs=[torch.randn((2, 1, 32, 32))], device="cuda:0") + self.assertEqual(next(model.parameters()).device.type, "cpu") + self.assertEqual(next(iter(exported.state_dict.values())).device.type, "cpu") + + def test_extra_files(self): + """Extra files are stored in the saved archive.""" + with tempfile.TemporaryDirectory() as tempdir: + filename = os.path.join(tempdir, "model.pt2") + convert_to_export( + _build_unet(), filename, extra_files={"a.txt": "data"}, inputs=[torch.randn((2, 1, 32, 32))] + ) + loaded = {"a.txt": ""} + torch.export.load(filename, extra_files=loaded) + self.assertEqual(loaded["a.txt"], "data") + + def test_export_without_save(self): + """Export a model without saving to disk.""" + model = _build_unet() + exported = convert_to_export(model=model, inputs=[torch.randn((2, 1, 32, 32))]) + self.assertIsInstance(exported, torch.export.ExportedProgram) + out = exported.module()(torch.randn(2, 1, 32, 32)) + self.assertEqual(out.shape, torch.Size([2, 3, 32, 32])) + + def test_missing_inputs_raises(self): + """Verify that missing inputs raise ValueError.""" + model = _build_unet() + with self.assertRaises(ValueError): + convert_to_export(model=model) + + @unittest.skipUnless(pytorch_after(2, 9), "torch.export.Dim.DYNAMIC requires PyTorch >= 2.9") + def test_export_with_dynamic_shapes(self): + """Export with dynamic batch dimension.""" + model = _build_unet() + dynamic = torch.export.Dim.DYNAMIC + static = torch.export.Dim.STATIC + exported = convert_to_export( + model=model, inputs=[torch.randn((2, 1, 32, 32))], dynamic_shapes=((dynamic, static, dynamic, dynamic),) + ) + # Verify works with different batch size and spatial dims + out = exported.module()(torch.randn(4, 1, 64, 64)) + self.assertEqual(out.shape, torch.Size([4, 3, 64, 64])) + + +if __name__ == "__main__": + unittest.main() From 7eca946ba5a7e66b6014a89c812a3eac3756cc0d Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:11:30 +0000 Subject: [PATCH 03/18] Add test_export_save test helper Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- tests/test_utils.py | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 9d245f9778..c700675011 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -49,7 +49,7 @@ from monai.config.type_definitions import NdarrayOrTensor from monai.data import create_test_image_2d, create_test_image_3d from monai.data.meta_tensor import MetaTensor, get_track_meta -from monai.networks import convert_to_onnx, convert_to_torchscript +from monai.networks import convert_to_export, convert_to_onnx, convert_to_torchscript from monai.utils import optional_import from monai.utils.misc import MONAIEnvVars from monai.utils.module import compute_capabilities_after, pytorch_after @@ -781,6 +781,25 @@ def test_script_save(net, *inputs, device=None, rtol=1e-4, atol=0.0): ) +def test_export_save(net, *inputs, dynamic_shapes=None, rtol=1e-4, atol=0.0): + """ + Test the ability to save `net` with `torch.export`, reload it, and apply inference. The value `inputs` is + forward-passed through the original and loaded copy of the network and their results compared. + The forward pass for both is done without gradient accumulation, on CPU as with `test_script_save`. + """ + with tempfile.TemporaryDirectory() as tempdir: + convert_to_export( + model=net, + filename_or_obj=os.path.join(tempdir, "model.pt2"), + verify=True, + inputs=[i.to("cpu") if isinstance(i, torch.Tensor) else i for i in inputs], + dynamic_shapes=dynamic_shapes, + device="cpu", + rtol=rtol, + atol=atol, + ) + + def test_onnx_save(net, *inputs, device=None, rtol=1e-4, atol=0.0): """ Test the ability to save `net` in ONNX format, reload it and validate with runtime. From b6ef51aa4d5a7e78b86c2ddf727ecf6eb6937aef Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:13:15 +0000 Subject: [PATCH 04/18] Build reduce_axis with list(range()) in losses Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/losses/dice.py | 4 ++-- monai/losses/mcc_loss.py | 2 +- monai/losses/tversky.py | 2 +- tests/losses/test_calibration_loss.py | 8 ++++++++ tests/losses/test_dice_ce_loss.py | 6 ++++++ tests/losses/test_dice_focal_loss.py | 7 ++++++- tests/losses/test_dice_loss.py | 7 ++++++- tests/losses/test_ds_loss.py | 7 ++++++- tests/losses/test_generalized_dice_focal_loss.py | 7 ++++++- tests/losses/test_generalized_dice_loss.py | 7 ++++++- .../test_generalized_wasserstein_dice_loss.py | 13 ++++++++++++- tests/losses/test_mcc_loss.py | 7 ++++++- tests/losses/test_multi_scale.py | 7 ++++++- tests/losses/test_spectral_loss.py | 7 ++++++- tests/losses/test_tversky_loss.py | 7 ++++++- 15 files changed, 84 insertions(+), 14 deletions(-) diff --git a/monai/losses/dice.py b/monai/losses/dice.py index 5561012f0e..a1082fca29 100644 --- a/monai/losses/dice.py +++ b/monai/losses/dice.py @@ -197,7 +197,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: input, target = mask_loss_inputs(input, target, self.ignore_index, mask=mask) # reducing only spatial dimensions (not batch nor channels) - reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist() + reduce_axis: list[int] = list(range(2, len(input.shape))) if self.batch: # reducing spatial dimensions and batch reduce_axis = [0] + reduce_axis @@ -446,7 +446,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})") # reducing only spatial dimensions (not batch nor channels) - reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist() + reduce_axis: list[int] = list(range(2, len(input.shape))) if self.batch: reduce_axis = [0] + reduce_axis diff --git a/monai/losses/mcc_loss.py b/monai/losses/mcc_loss.py index 17323f8941..0b57096fea 100644 --- a/monai/losses/mcc_loss.py +++ b/monai/losses/mcc_loss.py @@ -157,7 +157,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})") # reducing only spatial dimensions (not batch nor channels) - reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist() + reduce_axis: list[int] = list(range(2, len(input.shape))) if self.batch: reduce_axis = [0] + reduce_axis diff --git a/monai/losses/tversky.py b/monai/losses/tversky.py index a76c54aa35..f54c86149d 100644 --- a/monai/losses/tversky.py +++ b/monai/losses/tversky.py @@ -159,7 +159,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})") # reducing only spatial dimensions (not batch nor channels) - reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist() + reduce_axis: list[int] = list(range(2, len(input.shape))) if self.batch: # reducing spatial dimensions and batch reduce_axis = [0] + reduce_axis diff --git a/tests/losses/test_calibration_loss.py b/tests/losses/test_calibration_loss.py index 1504a0029e..528b469a06 100644 --- a/tests/losses/test_calibration_loss.py +++ b/tests/losses/test_calibration_loss.py @@ -20,6 +20,7 @@ from monai.losses.calibration import _hard_binned_calibration, _soft_binned_calibration from monai.metrics import calibration_binning from tests.test_utils import skip_if_no_cuda +from tests.test_utils import test_export_save as check_export_save from tests.test_utils import test_script_save as check_script_save LOSS_TYPES = (HardL1ACELoss, SoftL1ACELoss) @@ -236,6 +237,13 @@ def test_script_save(self, loss_type): target = target.movedim(-1, 1).float() check_script_save(loss_type(num_bins=5), prediction, target) + @parameterized.expand([HardL1ACELoss, SoftL1ACELoss]) + def test_export_save(self, loss_type): + prediction = torch.rand(2, 3, 4, 5, 6) + target = torch.nn.functional.one_hot(torch.randint(0, 3, (2, 4, 5, 6)), num_classes=3) + target = target.movedim(-1, 1).float() + check_export_save(loss_type(num_bins=5), prediction, target) + @parameterized.expand([(loss_type, dtype) for loss_type in LOSS_TYPES for dtype in (torch.float32, torch.float64)]) @skip_if_no_cuda def test_cuda_float32_and_float64(self, loss_type, dtype): diff --git a/tests/losses/test_dice_ce_loss.py b/tests/losses/test_dice_ce_loss.py index 97c7ae5050..761ab0db78 100644 --- a/tests/losses/test_dice_ce_loss.py +++ b/tests/losses/test_dice_ce_loss.py @@ -18,6 +18,7 @@ from parameterized import parameterized from monai.losses import DiceCELoss +from tests.test_utils import test_export_save TEST_CASES = [ [ # shape: (2, 2, 3), (2, 1, 3) @@ -118,6 +119,11 @@ def test_ill_shape3(self): # test_input = torch.ones(2, 2, 8, 8) # test_script_save(loss, test_input, test_input) + def test_export(self): + loss = DiceCELoss() + test_input = torch.rand(2, 2, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_dice_focal_loss.py b/tests/losses/test_dice_focal_loss.py index 98ea475ded..d78ef9ecae 100644 --- a/tests/losses/test_dice_focal_loss.py +++ b/tests/losses/test_dice_focal_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import DiceFocalLoss, DiceLoss, FocalLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save class TestDiceFocalLoss(unittest.TestCase): @@ -90,6 +90,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = DiceFocalLoss() + test_input = torch.rand(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + @parameterized.expand( [ ("sum_None_0.5_0.25", "sum", None, 0.5, 0.25), diff --git a/tests/losses/test_dice_loss.py b/tests/losses/test_dice_loss.py index d8fb5e1195..8c634d44c7 100644 --- a/tests/losses/test_dice_loss.py +++ b/tests/losses/test_dice_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import DiceLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) @@ -217,6 +217,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = DiceLoss() + test_input = torch.ones(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_ds_loss.py b/tests/losses/test_ds_loss.py index 586826aafe..376eeb1669 100644 --- a/tests/losses/test_ds_loss.py +++ b/tests/losses/test_ds_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import DeepSupervisionLoss, DiceCELoss, DiceFocalLoss, DiceLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES_DICECE = [ [ @@ -156,6 +156,11 @@ def test_script(self): test_input = torch.ones(2, 2, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = DeepSupervisionLoss(DiceCELoss()) + test_input = torch.ones(2, 2, 8, 8) + test_export_save(loss, test_input, test_input) + class TestDSLossDiceCE2(unittest.TestCase): @parameterized.expand(TEST_CASES_DICECE2) diff --git a/tests/losses/test_generalized_dice_focal_loss.py b/tests/losses/test_generalized_dice_focal_loss.py index 2af4aa68db..df0cb13ad5 100644 --- a/tests/losses/test_generalized_dice_focal_loss.py +++ b/tests/losses/test_generalized_dice_focal_loss.py @@ -17,7 +17,7 @@ import torch from monai.losses import FocalLoss, GeneralizedDiceFocalLoss, GeneralizedDiceLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save class TestGeneralizedDiceFocalLoss(unittest.TestCase): @@ -80,6 +80,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = GeneralizedDiceFocalLoss() + test_input = torch.rand(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_generalized_dice_loss.py b/tests/losses/test_generalized_dice_loss.py index 7d60b99932..223a4209fe 100644 --- a/tests/losses/test_generalized_dice_loss.py +++ b/tests/losses/test_generalized_dice_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import GeneralizedDiceLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) @@ -211,6 +211,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = GeneralizedDiceLoss() + test_input = torch.rand(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_generalized_wasserstein_dice_loss.py b/tests/losses/test_generalized_wasserstein_dice_loss.py index 1ed05c137d..73bc2ffd5c 100644 --- a/tests/losses/test_generalized_wasserstein_dice_loss.py +++ b/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_script_save +from tests.test_utils import test_export_save, test_script_save class TestGeneralizedWassersteinDiceLoss(unittest.TestCase): @@ -314,6 +314,17 @@ def test_script(self): test_script_save(loss, pred_very_good, target) + def test_export(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) + target = target.unsqueeze(0) + pred_very_good = 1000 * F.one_hot(target, num_classes=2).permute(0, 3, 1, 2).float() + + 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) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_mcc_loss.py b/tests/losses/test_mcc_loss.py index 2451cf76d4..38a5ae3ddc 100644 --- a/tests/losses/test_mcc_loss.py +++ b/tests/losses/test_mcc_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import MCCLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2), sigmoid @@ -149,6 +149,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = MCCLoss() + test_input = torch.rand(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_multi_scale.py b/tests/losses/test_multi_scale.py index 87ccca7676..0d0260e69c 100644 --- a/tests/losses/test_multi_scale.py +++ b/tests/losses/test_multi_scale.py @@ -18,7 +18,7 @@ from monai.losses import DiceLoss from monai.losses.multi_scale import MultiScaleLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save dice_loss = DiceLoss(include_background=True, sigmoid=True, smooth_nr=1e-5, smooth_dr=1e-5) device = "cuda" if torch.cuda.is_available() else "cpu" @@ -81,6 +81,11 @@ def test_script(self): loss = MultiScaleLoss(**input_param) test_script_save(loss, input_data["y_pred"], input_data["y_true"]) + def test_export(self): + input_param, input_data, expected_val = TEST_CASES[0] + loss = MultiScaleLoss(**input_param) + test_export_save(loss, input_data["y_pred"], input_data["y_true"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_spectral_loss.py b/tests/losses/test_spectral_loss.py index 8a4988a30d..035ac99625 100644 --- a/tests/losses/test_spectral_loss.py +++ b/tests/losses/test_spectral_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import JukeboxLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES = [ [ @@ -81,6 +81,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = JukeboxLoss(spatial_dims=2) + test_input = torch.ones(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() diff --git a/tests/losses/test_tversky_loss.py b/tests/losses/test_tversky_loss.py index 32303434ca..a97c000b65 100644 --- a/tests/losses/test_tversky_loss.py +++ b/tests/losses/test_tversky_loss.py @@ -18,7 +18,7 @@ from parameterized import parameterized from monai.losses import TverskyLoss -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES = [ [ # shape: (1, 1, 2, 2), (1, 1, 2, 2) @@ -193,6 +193,11 @@ def test_script(self): test_input = torch.ones(2, 1, 8, 8) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = TverskyLoss() + test_input = torch.rand(2, 1, 8, 8) + test_export_save(loss, test_input, test_input) + if __name__ == "__main__": unittest.main() From a9db1d8c4d52caf11168f78c97599c2747b809aa Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:15:06 +0000 Subject: [PATCH 05/18] Move Dice and Focal weight checks to __init__ Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/losses/dice.py | 5 +++-- monai/losses/focal_loss.py | 10 +++++----- tests/losses/test_dice_loss.py | 17 ++++++++++++++++- tests/losses/test_focal_loss.py | 17 ++++++++++++++--- 4 files changed, 38 insertions(+), 11 deletions(-) diff --git a/monai/losses/dice.py b/monai/losses/dice.py index a1082fca29..858d706cae 100644 --- a/monai/losses/dice.py +++ b/monai/losses/dice.py @@ -111,6 +111,7 @@ def __init__( TypeError: When ``other_act`` is not an ``Optional[Callable]``. ValueError: When more than 1 of [``sigmoid=True``, ``softmax=True``, ``other_act is not None``]. Incompatible values. + ValueError: When ``weight`` has negative values. """ super().__init__(reduction=LossReduction(reduction).value) @@ -130,6 +131,8 @@ def __init__( self.batch = batch self.ignore_index = ignore_index weight = torch.as_tensor(weight) if weight is not None else None + if weight is not None and (weight < 0).any(): + raise ValueError(f"the value/values of the `weight` should be no less than 0, got {weight}.") self.register_buffer("class_weight", weight) self.class_weight: None | torch.Tensor self.soft_label = soft_label @@ -223,8 +226,6 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: "The length of the `weight` sequence should be the same as the number of classes. " "If `include_background=False`, the weight should not include the background category class 0." ) - if self.class_weight.min() < 0: - raise ValueError("the value/values of the `weight` should be no less than 0.") # apply class_weight to loss f = f * self.class_weight.to(f) diff --git a/monai/losses/focal_loss.py b/monai/losses/focal_loss.py index 1e7063e06f..05fa15dd08 100644 --- a/monai/losses/focal_loss.py +++ b/monai/losses/focal_loss.py @@ -114,6 +114,9 @@ def __init__( >>> grnd = torch.tensor([[0], [1], [0]], dtype=torch.int64) >>> fl = FocalLoss(to_onehot_y=True) >>> fl(pred, grnd) + + Raises: + ValueError: When ``weight`` has negative values. """ super().__init__(reduction=LossReduction(reduction).value) self.include_background = include_background @@ -129,6 +132,8 @@ def __init__( else: self.alpha = torch.as_tensor(alpha) weight = torch.as_tensor(weight) if weight is not None else None + if weight is not None and (weight < 0).any(): + raise ValueError(f"class_weight values must be non-negative, got {weight}.") self.register_buffer("class_weight", weight) self.class_weight: None | torch.Tensor self.ignore_index = ignore_index @@ -210,11 +215,6 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor: f"The number of class_weight ({cw.numel()}) must match the number of " f"output channels ({num_classes})." ) - if (cw < 0).any(): - raise ValueError("class_weight values must be non-negative.") - else: - if cw < 0: - raise ValueError("class_weight values must be non-negative.") if cw.ndim == 0: loss = loss * cw diff --git a/tests/losses/test_dice_loss.py b/tests/losses/test_dice_loss.py index 8c634d44c7..28222d99ea 100644 --- a/tests/losses/test_dice_loss.py +++ b/tests/losses/test_dice_loss.py @@ -212,6 +212,21 @@ def test_input_warnings(self): loss = DiceLoss(to_onehot_y=True) loss.forward(chn_input, chn_target) + @parameterized.expand([[(1.0, -1.0)], [-1.0]]) + def test_negative_weight(self, weight): + with self.assertRaisesRegex(ValueError, "no less than 0"): + DiceLoss(weight=weight) + + def test_empty_weight_keeps_length_error(self): + loss = DiceLoss(weight=[]) + with self.assertRaisesRegex(ValueError, "same as the number of classes"): + loss(torch.ones(1, 2, 3, 3), torch.ones(1, 2, 3, 3)) + + def test_export_with_weight(self): + loss = DiceLoss(weight=(1.0, 2.0)) + test_input = torch.rand(2, 2, 8, 8) + test_export_save(loss, test_input, test_input) + def test_script(self): loss = DiceLoss() test_input = torch.ones(2, 1, 8, 8) @@ -219,7 +234,7 @@ def test_script(self): def test_export(self): loss = DiceLoss() - test_input = torch.ones(2, 1, 8, 8) + test_input = torch.rand(2, 1, 8, 8) test_export_save(loss, test_input, test_input) diff --git a/tests/losses/test_focal_loss.py b/tests/losses/test_focal_loss.py index 35017ec898..21580bbb07 100644 --- a/tests/losses/test_focal_loss.py +++ b/tests/losses/test_focal_loss.py @@ -21,7 +21,7 @@ from monai.losses import FocalLoss from monai.networks import one_hot -from tests.test_utils import TEST_DEVICES, test_script_save +from tests.test_utils import TEST_DEVICES, test_export_save, test_script_save TEST_CASES = [] for case in TEST_DEVICES: @@ -356,8 +356,11 @@ def test_ill_class_weight(self): FocalLoss(include_background=True, weight=(1.0, 1.0, 2.0))(chn_input, chn_target) with self.assertRaisesRegex(ValueError, ""): FocalLoss(include_background=False, weight=(1.0, 1.0, 1.0, 1.0))(chn_input, chn_target) - with self.assertRaisesRegex(ValueError, ""): - FocalLoss(include_background=False, weight=(1.0, 1.0, -1.0))(chn_input, chn_target) + + @parameterized.expand([[(1.0, 1.0, -1.0)], [-1.0]]) + def test_negative_weight(self, weight): + with self.assertRaisesRegex(ValueError, "non-negative"): + FocalLoss(weight=weight) def test_warnings(self): with self.assertWarns(Warning): @@ -376,6 +379,14 @@ def test_warnings(self): loss = FocalLoss(include_background=False, use_softmax=True, alpha=0.5) loss(chn_input, chn_target) + @parameterized.expand( + [[True, None], [False, None], [True, (1.0, 2.0)], [False, (1.0, 2.0)], [True, 0.5], [False, 0.5]] + ) + def test_export(self, use_softmax, weight): + loss = FocalLoss(use_softmax=use_softmax, weight=weight) + test_input = torch.ones(2, 2, 8, 8) + test_export_save(loss, test_input, test_input) + def test_script(self): for use_softmax in [True, False]: loss = FocalLoss(use_softmax=use_softmax) From 0e3135fbc15a56aa5a1f341e2ed0ae54cc70994f Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:15:19 +0000 Subject: [PATCH 06/18] Make CoilSensitivityModel export and script Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- .../networks/nets/coil_sensitivity_model.py | 40 +++++-- .../nets/test_coil_sensitivity_model.py | 100 ++++++++++++++++++ tests/networks/test_varnet.py | 15 ++- 3 files changed, 148 insertions(+), 7 deletions(-) create mode 100644 tests/apps/reconstruction/nets/test_coil_sensitivity_model.py diff --git a/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py b/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py index 91a9f3d8d3..d2b48765a2 100644 --- a/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py +++ b/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py @@ -24,6 +24,38 @@ reshape_channel_to_batch_dim, ) from monai.networks.blocks.fft_utils_t import ifftn_centered_t +from monai.utils import deprecated + + +def _fully_sampled_mask(mask: Tensor) -> Tensor: + """ + Compute a boolean mask of the Auto Calibration Region (ACR), the fully-sampled center of the k-space. + This uses tensor operations rather than loops over mask values, which ``torch.export`` cannot trace, and creates + no tensors with an explicit device, which an exported program would keep and so could not be moved. + + Args: + mask: the under-sampling mask of shape (..., S, 1) where S denotes the sampling dimension. + + Returns: + A boolean tensor with the same shape as ``mask``. + + Raises: + ValueError: When ``mask`` holds more than one sampling vector. + """ + num_samples = mask.shape[-2] + if mask.numel() != num_samples: + raise ValueError(f"mask must hold a single sampling vector, got shape {mask.shape}.") + + center = num_samples // 2 + sampled = mask.reshape(-1) != 0 + right = torch.cumprod(sampled[center:], dim=0).sum() + left = torch.cumprod(sampled[: center + 1].flip([0]), dim=0).sum() + num_low_freqs = left + right - 1 # size of the fully-sampled center, which is counted twice + + start = (num_samples - num_low_freqs + 1) // 2 # this marks the start of center extraction + positions = torch.cumsum(torch.ones_like(sampled, dtype=torch.int64), dim=0) - 1 + region = (positions >= start) & (positions < start + num_low_freqs) + return region.view(mask.shape) class CoilSensitivityModel(nn.Module): @@ -85,6 +117,7 @@ def __init__( self.spatial_dims = spatial_dims self.coil_dim = coil_dim + @deprecated(since="1.7.0", removed="1.9.0", msg_suffix="`forward` no longer uses it.") def get_fully_sampled_region(self, mask: Tensor) -> tuple[int, int]: """ Extracts the size of the fully-sampled part of the kspace. Note that when a kspace @@ -122,13 +155,8 @@ def forward(self, masked_kspace: Tensor, mask: Tensor) -> Tensor: Returns: predicted coil sensitivity maps with shape (B,C,H,W,2) for 2D data or (B,C,H,W,D,2) for 3D data. """ - left, right = self.get_fully_sampled_region(mask) - num_low_freqs = right - left # size of the fully-sampled center - # take out the fully-sampled region and set the rest of the data to zero - x = torch.zeros_like(masked_kspace) - start = (mask.shape[-2] - num_low_freqs + 1) // 2 # this marks the start of center extraction - x[..., start : start + num_low_freqs, :] = masked_kspace[..., start : start + num_low_freqs, :] + x = masked_kspace.masked_fill(~_fully_sampled_mask(mask), 0) # apply inverse fourier to the extracted fully-sampled data x = ifftn_centered_t(x, spatial_dims=self.spatial_dims, is_complex=True) diff --git a/tests/apps/reconstruction/nets/test_coil_sensitivity_model.py b/tests/apps/reconstruction/nets/test_coil_sensitivity_model.py new file mode 100644 index 0000000000..8e792c445d --- /dev/null +++ b/tests/apps/reconstruction/nets/test_coil_sensitivity_model.py @@ -0,0 +1,100 @@ +# 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 unittest + +import torch +from parameterized import parameterized + +from monai.apps.reconstruction.networks.nets.coil_sensitivity_model import CoilSensitivityModel, _fully_sampled_mask +from tests.test_utils import skip_if_no_cuda, test_export_save, test_script_save + +FEATURES = (2, 2, 2, 2, 2, 2) + +# number of samples, then the sampled positions as (start, stop), then the expected region as (start, stop) +TEST_REGIONS = [ + [20, [(8, 12)], (8, 12)], + [21, [(8, 13)], (8, 13)], + [32, [(12, 20), (2, 4), (28, 30)], (12, 20)], + [32, [(10, 20)], (11, 21)], # off-center regions are recentered, as they have always been + [32, [(0, 32)], (0, 32)], + [32, [(0, 20)], (6, 26)], + [32, [(2, 4)], (0, 0)], + [32, [(16, 32)], (8, 24)], # used to raise IndexError as the region reaches the last index + [32, [(0, 20), (30, 32)], (6, 26)], # used to wrap past index 0 into the samples at the end +] + + +def _mask(num_samples, sampled): + mask = torch.zeros(1, 1, 1, num_samples, 1, dtype=torch.bool) + for start, stop in sampled: + mask[..., start:stop, :] = True + return mask + + +class TestCoilSensitivityModel(unittest.TestCase): + @parameterized.expand(TEST_REGIONS) + def test_fully_sampled_mask(self, num_samples, sampled, expected): + expected_mask = _mask(num_samples, [expected]) + result = _fully_sampled_mask(_mask(num_samples, sampled)) + self.assertTrue(torch.equal(result, expected_mask)) + + @parameterized.expand( + [[20, [(8, 12)]], [21, [(8, 13)]], [32, [(12, 20), (2, 4), (28, 30)]], [32, [(10, 20)]], [32, [(2, 4)]]] + ) + def test_matches_get_fully_sampled_region(self, num_samples, sampled): + model = CoilSensitivityModel(spatial_dims=2) + mask = _mask(num_samples, sampled) + left, right = model.get_fully_sampled_region(mask) + start = (num_samples - (right - left) + 1) // 2 + expected = torch.zeros_like(mask) + expected[..., start : start + right - left, :] = True + self.assertTrue(torch.equal(_fully_sampled_mask(mask), expected)) + + def test_multiple_sampling_vectors(self): + mask = torch.ones(2, 1, 1, 32, 1, dtype=torch.bool) + with self.assertRaisesRegex(ValueError, "single sampling vector"): + _fully_sampled_mask(mask) + + def test_values_outside_region_are_dropped(self): + net = CoilSensitivityModel(spatial_dims=2, features=FEATURES).eval() + data = torch.randn(1, 3, 32, 32, 2) + data[..., 2, :] = float("inf") + with torch.no_grad(): + self.assertTrue(torch.isfinite(net(data, _mask(32, [(12, 20)]))).all()) + + @skip_if_no_cuda + def test_exported_program_moves_to_gpu(self): + net = CoilSensitivityModel(spatial_dims=2, features=FEATURES).eval() + data, mask = torch.randn(1, 3, 32, 32, 2), _mask(32, [(12, 20)]) + exported = torch.export.export(net, args=(data, mask)).module().to("cuda:0") + self.assertEqual(exported(data.to("cuda:0"), mask.to("cuda:0")).device.type, "cuda") + + def test_script(self): + net = CoilSensitivityModel(spatial_dims=2, features=FEATURES) + test_script_save(net, torch.randn(1, 3, 32, 32, 2), _mask(32, [(12, 20)])) + + def test_export(self): + net = CoilSensitivityModel(spatial_dims=2, features=FEATURES) + test_export_save(net, torch.randn(1, 3, 32, 32, 2), _mask(32, [(12, 20)])) + + def test_export_follows_mask(self): + net = CoilSensitivityModel(spatial_dims=2, features=FEATURES).eval() + data = torch.randn(1, 3, 32, 32, 2) + exported = torch.export.export(net, args=(data, _mask(32, [(12, 20)]))).module() + other_mask = _mask(32, [(10, 22)]) + torch.testing.assert_close(exported(data, other_mask), net(data, other_mask)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/networks/test_varnet.py b/tests/networks/test_varnet.py index 9cadf15ce5..4d1f56d054 100644 --- a/tests/networks/test_varnet.py +++ b/tests/networks/test_varnet.py @@ -20,7 +20,7 @@ from monai.apps.reconstruction.networks.nets.complex_unet import ComplexUnet from monai.apps.reconstruction.networks.nets.varnet import VariationalNetworkModel from monai.networks import eval_mode -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = torch.device("cuda" if torch.cuda.is_available() else "cpu") coil_sens_model = CoilSensitivityModel(spatial_dims=2, features=[8, 16, 32, 64, 128, 8]) @@ -57,6 +57,19 @@ def test_script(self, coil_sens_model, refinement_model, num_cascades, input_sha test_script_save(net, test_data, mask.bool()) + @parameterized.expand(TESTS) + def test_export(self, coil_sens_model, refinement_model, num_cascades, input_shape, expected_shape): + net = VariationalNetworkModel(coil_sens_model, refinement_model, num_cascades) + + mask_shape = [1 for _ in input_shape] + mask_shape[-2] = input_shape[-2] + mask = torch.zeros(mask_shape) + mask[..., mask_shape[-2] // 2 - 5 : mask_shape[-2] // 2 + 5, :] = 1 + + test_data = torch.randn(input_shape) + + test_export_save(net, test_data, mask.bool()) + if __name__ == "__main__": unittest.main() From e84081b82f2d4ca880ccc2e459f7798368fb45be Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:17:09 +0000 Subject: [PATCH 07/18] Extract _prepare_ckpt_export from ckpt_export Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/bundle/scripts.py | 70 +++++++++++++++++++++++++++-------------- 1 file changed, 46 insertions(+), 24 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index b280919128..ea86697658 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1453,6 +1453,49 @@ def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str, **kwargs: Any) -> N ) +def _prepare_ckpt_export( + _args: dict, + config_file: str | Sequence[str], + meta_file: str | Sequence[str] | None, + filepath: PathLike | None, + ckpt_file: str | None, + net_id: str | None, + default_filename: str, +) -> tuple[ConfigParser, PathLike, str, str]: + """ + Shared setup for checkpoint export: parse the config and metadata, apply the remaining ``_args`` as config + overrides, resolve the default output/checkpoint paths and network id, and validate that the network + definition can be found. Returns ``(parser, filepath, ckpt_file, net_id)``. + """ + bundle_root = _args.get("bundle_root", os.getcwd()) + + parser = ConfigParser() + parser.read_config(f=config_file) + meta_file = os.path.join(bundle_root, "configs", "metadata.json") if meta_file is None else meta_file + # a single missing file is skipped since bundles may have no metadata, a missing file in a list is an error + if not isinstance(meta_file, (str, os.PathLike)) or os.path.exists(meta_file): + parser.read_meta(f=meta_file) + + # the rest key-values in the _args are to override config content + for k, v in _args.items(): + parser[k] = v + + filepath = os.path.join(bundle_root, "models", default_filename) if filepath is None else filepath + ckpt_file = os.path.join(bundle_root, "models", "model.pt") if ckpt_file is None else ckpt_file + if not os.path.exists(ckpt_file): + raise FileNotFoundError(f'Checkpoint file "{ckpt_file}" not found, please specify it in argument "ckpt_file".') + + net_id = "network_def" if net_id is None else net_id + try: + parser.get_parsed_content(net_id) + except ValueError as e: + raise ValueError( + f'Network definition "{net_id}" cannot be found in "{config_file}", specify name with argument "net_id".' + ) from e + + return parser, filepath, ckpt_file, net_id + + def ckpt_export( net_id: str | None = None, filepath: PathLike | None = None, @@ -1538,30 +1581,9 @@ def ckpt_export( input_shape=None, converter_kwargs={}, ) - bundle_root = _args.get("bundle_root", os.getcwd()) - - parser = ConfigParser() - parser.read_config(f=config_file_) - meta_file_ = os.path.join(bundle_root, "configs", "metadata.json") if meta_file_ is None else meta_file_ - if os.path.exists(meta_file_): - parser.read_meta(f=meta_file_) - - # the rest key-values in the _args are to override config content - for k, v in _args.items(): - parser[k] = v - - filepath_ = os.path.join(bundle_root, "models", "model.ts") if filepath_ is None else filepath_ - ckpt_file_ = os.path.join(bundle_root, "models", "model.pt") if ckpt_file_ is None else ckpt_file_ - if not os.path.exists(ckpt_file_): - raise FileNotFoundError(f'Checkpoint file "{ckpt_file_}" not found, please specify it in argument "ckpt_file".') - - net_id_ = "network_def" if net_id_ is None else net_id_ - try: - parser.get_parsed_content(net_id_) - except ValueError as e: - raise ValueError( - f'Network definition "{net_id_}" cannot be found in "{config_file_}", specify name with argument "net_id".' - ) from e + parser, filepath_, ckpt_file_, net_id_ = _prepare_ckpt_export( + _args, config_file_, meta_file_, filepath_, ckpt_file_, net_id_, default_filename="model.ts" + ) # When export through torch.jit.trace without providing input_shape, will try to parse one from the parser. if (not input_shape_) and use_trace: From 8a5fb7547306bc371851ab538883073834605069 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:20:24 +0000 Subject: [PATCH 08/18] Add export_checkpoint bundle command Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/bundle/__init__.py | 1 + monai/bundle/__main__.py | 1 + monai/bundle/scripts.py | 119 ++++++++++++++- tests/bundle/test_bundle_export_checkpoint.py | 143 ++++++++++++++++++ tests/min_tests.py | 1 + 5 files changed, 263 insertions(+), 2 deletions(-) create mode 100644 tests/bundle/test_bundle_export_checkpoint.py diff --git a/monai/bundle/__init__.py b/monai/bundle/__init__.py index 3f3c8d545e..194d3deacd 100644 --- a/monai/bundle/__init__.py +++ b/monai/bundle/__init__.py @@ -20,6 +20,7 @@ create_workflow, download, download_large_files, + export_checkpoint, get_all_bundles_list, get_bundle_info, get_bundle_versions, diff --git a/monai/bundle/__main__.py b/monai/bundle/__main__.py index 778c9ef2f0..9593296978 100644 --- a/monai/bundle/__main__.py +++ b/monai/bundle/__main__.py @@ -15,6 +15,7 @@ ckpt_export, download, download_large_files, + export_checkpoint, init_bundle, onnx_export, run, diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index ea86697658..8ca657ba12 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -35,8 +35,9 @@ from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA, merge_kv from monai.bundle.workflows import BundleWorkflow, ConfigWorkflow from monai.config import PathLike -from monai.data import load_net_with_metadata, save_net_with_metadata +from monai.data import load_net_with_metadata, save_exported_program, save_net_with_metadata from monai.networks import ( + convert_to_export, convert_to_onnx, convert_to_torchscript, convert_to_trt, @@ -1295,7 +1296,7 @@ def _export( (extracted from the parser), and a dictionary of extra JSON files (name -> contents) as input. parser: a ConfigParser of the bundle to be converted. net_id: ID name of the network component in the parser, it must be `torch.nn.Module`. - filepath: filepath to export, if filename has no extension, it becomes `.ts`. + filepath: filepath to export, the `saver` decides the extension of a filename without one. ckpt_file: filepath of the model checkpoint to load. config_file: filepath of the config file to save in the converted model,the saved key in the converted model is the config filename without extension, and the saved config value is always serialized in @@ -1609,6 +1610,120 @@ def ckpt_export( ) +def export_checkpoint( + net_id: str | None = None, + filepath: PathLike | None = None, + ckpt_file: str | None = None, + meta_file: str | Sequence[str] | None = None, + config_file: str | Sequence[str] | None = None, + key_in_ckpt: str | None = None, + input_shape: Sequence[int] | None = None, + dynamic_shapes: dict | tuple | None = None, + args_file: str | None = None, + converter_kwargs: Mapping | None = None, + **override: Any, +) -> None: + """ + Export the model checkpoint to a ``.pt2`` file using :func:`torch.export.export`, with metadata and + config included. Without ``dynamic_shapes`` the saved program only accepts inputs of exactly ``input_shape``, + including the batch size. The network is exported on CPU. + + Typical usage examples: + + .. code-block:: bash + + python -m monai.bundle export_checkpoint network --filepath --ckpt_file ... + + Args: + net_id: ID name of the network component in the config, it must be ``torch.nn.Module``. + Default to ``"network_def"``. + filepath: filepath to export. If filename has no extension it becomes ``.pt2``. + Default to ``"models/model.pt2"`` under ``"os.getcwd()"`` if ``bundle_root`` is not specified. + ckpt_file: filepath of the model checkpoint to load. + Default to ``"models/model.pt"`` under ``"os.getcwd()"`` if ``bundle_root`` is not specified. + meta_file: filepath of the metadata file. If it is a list of file paths, contents will be merged. + Default to ``"configs/metadata.json"`` under ``"os.getcwd()"`` if ``bundle_root`` is not specified. + config_file: filepath of the config file to save in the exported model. The saved key is the + config filename without extension; the value is always serialized in JSON format. + It can be a single file or a list of files. If ``None``, must be provided in ``args_file``. + key_in_ckpt: for nested checkpoints like ``{"model": XXX, "optimizer": XXX, ...}``, specify the + key of model weights. If not nested, no need to set. + input_shape: a shape used to generate random input for the network, e.g. ``[N, C, H, W]`` or + ``[N, C, H, W, D]``. If not given, will try to parse from ``metadata``. + dynamic_shapes: dynamic shape specifications passed to :func:`torch.export.export`. These hold + ``torch.export.Dim`` objects, so they can only be given when calling this function from Python. + args_file: a JSON or YAML file to provide default values for all the parameters. + converter_kwargs: extra arguments for :func:`~monai.networks.utils.convert_to_export`, + except ones that already exist in the input parameters. + override: id-value pairs to override or add the corresponding config content. + + Raises: + ValueError: When a config file is named ``metadata``, which is reserved for the metadata in the archive. + """ + _args = update_kwargs( + args=args_file, + net_id=net_id, + filepath=filepath, + meta_file=meta_file, + config_file=config_file, + ckpt_file=ckpt_file, + key_in_ckpt=key_in_ckpt, + input_shape=input_shape, + dynamic_shapes=dynamic_shapes, + converter_kwargs=converter_kwargs, + **override, + ) + _log_input_summary(tag="export_checkpoint", args=_args) + ( + config_file_, + filepath_, + ckpt_file_, + net_id_, + meta_file_, + key_in_ckpt_, + input_shape_, + dynamic_shapes_, + converter_kwargs_, + ) = _pop_args( + _args, + "config_file", + filepath=None, + ckpt_file=None, + net_id=None, + meta_file=None, + key_in_ckpt="", + input_shape=None, + dynamic_shapes=None, + converter_kwargs={}, + ) + parser, filepath_, ckpt_file_, net_id_ = _prepare_ckpt_export( + _args, config_file_, meta_file_, filepath_, ckpt_file_, net_id_, default_filename="model.pt2" + ) + + if not input_shape_: + input_shape_ = _get_fake_input_shape(parser=parser) + + # export on CPU, whatever device the config uses, so that the saved file can be loaded on any machine + parser.get_parsed_content(net_id_).cpu() + inputs_: Sequence[Any] = [torch.rand(input_shape_)] + + converter_kwargs_.update({"inputs": inputs_, "dynamic_shapes": dynamic_shapes_}) + + save_ep = partial(save_exported_program, include_config_vals=False, append_timestamp=False) + + _export( + convert_to_export, + save_ep, + parser, + net_id=net_id_, + filepath=filepath_, + ckpt_file=ckpt_file_, + config_file=config_file_, + key_in_ckpt=key_in_ckpt_, + **converter_kwargs_, + ) + + def trt_export( net_id: str | None = None, filepath: PathLike | None = None, diff --git a/tests/bundle/test_bundle_export_checkpoint.py b/tests/bundle/test_bundle_export_checkpoint.py new file mode 100644 index 0000000000..3b8809b960 --- /dev/null +++ b/tests/bundle/test_bundle_export_checkpoint.py @@ -0,0 +1,143 @@ +# 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, export_checkpoint +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): + @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)) + + def test_meta_files_are_merged(self): + config_file = os.path.join(TESTS_PATH, "testing_data", "inference.json") + with tempfile.TemporaryDirectory() as tempdir: + parser = ConfigParser() + parser.read_config(config_file) + ckpt_file = os.path.join(tempdir, "model.pt") + save_state(src=parser.get_parsed_content("network_def"), path=ckpt_file) + meta_files = [os.path.join(tempdir, "first.json"), os.path.join(tempdir, "second.json")] + for meta_file, value in zip(meta_files, ({"first": 1}, {"second": 2})): + parser.export_config_file(config=value, filepath=meta_file) + pt2_file = os.path.join(tempdir, "model.pt2") + + export_checkpoint( + filepath=pt2_file, + ckpt_file=ckpt_file, + meta_file=meta_files, + config_file=config_file, + input_shape=[1, 1, 96, 96, 96], + ) + _, metadata, _ = load_exported_program(pt2_file) + + self.assertEqual(metadata, {"first": 1, "second": 2}) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/min_tests.py b/tests/min_tests.py index 25a42fe4b5..b153b6aa55 100644 --- a/tests/min_tests.py +++ b/tests/min_tests.py @@ -206,6 +206,7 @@ def run_testsuit(): "test_bundle_verify_metadata", "test_bundle_verify_net", "test_bundle_ckpt_export", + "test_bundle_export_checkpoint", "test_bundle_utils", "test_bundle_init_bundle", "test_fastmri_reader", From eaacac9ec5938f13f9a227f0ea97a1822c745faa Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:20:50 +0000 Subject: [PATCH 09/18] Add load_exported_module as load()'s last argument Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/bundle/scripts.py | 28 ++++-- tests/bundle/test_bundle_export_checkpoint.py | 86 ++++++++++++++++++- 2 files changed, 105 insertions(+), 9 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index 8ca657ba12..bd870ae2b8 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -35,7 +35,7 @@ from monai.bundle.utils import DEFAULT_INFERENCE, DEFAULT_METADATA, merge_kv from monai.bundle.workflows import BundleWorkflow, ConfigWorkflow from monai.config import PathLike -from monai.data import load_net_with_metadata, save_exported_program, save_net_with_metadata +from monai.data import load_exported_program, load_net_with_metadata, save_exported_program, save_net_with_metadata from monai.networks import ( convert_to_export, convert_to_onnx, @@ -648,9 +648,10 @@ def load( args_file: str | None = None, copy_model_args: dict | None = None, net_override: dict | None = None, + load_exported_module: bool = False, ) -> object | tuple[torch.nn.Module, dict, dict] | Any: """ - Load model weights or TorchScript module of a bundle. + Load model weights, TorchScript module or exported program of a bundle. Security note: if `model` is `None`, building `network_def` requires parsing the bundle's own "{workflow_type}.json" config, which can define `"_target_"` components resolved to any importable @@ -696,12 +697,16 @@ def load( device: target device of returned weights or module, if `None`, prefer to "cuda" if existing. key_in_ckpt: for nested checkpoint like `{"model": XXX, "optimizer": XXX, ...}`, specify the key of model weights. if not nested checkpoint, no need to set. - config_files: extra filenames would be loaded. The argument only works when loading a TorchScript module, - see `_extra_files` in `torch.jit.load` for more details. + config_files: extra filenames would be loaded. The argument only works when loading a TorchScript module + or exported program, see `_extra_files` in `torch.jit.load` and `extra_files` in `torch.export.load`. workflow_name: specified bundle workflow name, should be a string or class, default to "ConfigWorkflow". args_file: a JSON or YAML file to provide default values for all the args in "download" function. copy_model_args: other arguments for the `monai.networks.copy_model_state` function. net_override: id-value pairs to override the parameters in the network of the bundle, default to `None`. + load_exported_module: a flag to specify if loading a `torch.export` `.pt2` file, "models/model.pt2" if + `model_file` is `None`. This uses pickle, so only load files from a source you trust. `device` is not + applied since `torch.export.load` has no `map_location`, use `exported_program.module().to(device)`. + Modules from `exported_program.module()` reject `.eval()` and `.train()`. Returns: 1. If `load_ts_module` is `False` and `model` is `None`, @@ -712,8 +717,17 @@ def load( 3. If `load_ts_module` is `True`, return a triple that include a TorchScript module, the corresponding metadata dict, and extra files dict. please check `monai.data.load_net_with_metadata` for more details. + 4. If `load_exported_module` is `True`, return a triple that include an exported program, + the corresponding metadata dict, and extra files dict. + please check `monai.data.load_exported_program` for more details. + + Raises: + ValueError: When both `load_ts_module` and `load_exported_module` are `True`. """ + if load_ts_module and load_exported_module: + raise ValueError("`load_ts_module` and `load_exported_module` are mutually exclusive.") + bundle_dir_ = _process_bundle_dir(bundle_dir) net_override = {} if net_override is None else net_override copy_model_args = {} if copy_model_args is None else copy_model_args @@ -721,7 +735,8 @@ def load( if device is None: device = "cuda:0" if is_available() else "cpu" if model_file is None: - model_file = os.path.join("models", "model.ts" if load_ts_module is True else "model.pt") + default_file = "model.pt2" if load_exported_module else "model.ts" if load_ts_module is True else "model.pt" + model_file = os.path.join("models", default_file) if source == "ngc": name = _add_ngc_prefix(name) if remove_prefix: @@ -739,6 +754,9 @@ def load( args_file=args_file, ) + # loading with `torch.export.load` + if load_exported_module: + return load_exported_program(full_path, more_extra_files=config_files) # loading with `torch.jit.load` if load_ts_module is True: return load_net_with_metadata(full_path, map_location=torch.device(device), more_extra_files=config_files) diff --git a/tests/bundle/test_bundle_export_checkpoint.py b/tests/bundle/test_bundle_export_checkpoint.py index 3b8809b960..7306de8878 100644 --- a/tests/bundle/test_bundle_export_checkpoint.py +++ b/tests/bundle/test_bundle_export_checkpoint.py @@ -11,18 +11,20 @@ from __future__ import annotations +import inspect import json import os import tempfile import unittest from pathlib import Path +import torch from parameterized import parameterized -from monai.bundle import ConfigParser, export_checkpoint -from monai.data import load_exported_program +from monai.bundle import ConfigParser, export_checkpoint, load +from monai.data import load_exported_program, save_exported_program from monai.networks import save_state -from tests.test_utils import command_line_tests, skip_if_windows +from tests.test_utils import command_line_tests, skip_if_no_cuda, skip_if_windows TESTS_PATH = Path(__file__).parents[1] @@ -75,12 +77,16 @@ def test_export(self, key_in_ckpt): command_line_tests(cmd) self.assertTrue(os.path.exists(pt2_file)) - _, _metadata, extra_files = load_exported_program( + _loaded, _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"])) + exported_weights = _loaded.state_dict + for name, value in net.state_dict().items(): + torch.testing.assert_close(exported_weights[name], value) + @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") @@ -115,6 +121,46 @@ def test_default_value(self, key_in_ckpt): command_line_tests(cmd) self.assertTrue(os.path.exists(pt2_file)) + def test_missing_file_in_meta_list(self): + config_file = os.path.join(TESTS_PATH, "testing_data", "inference.json") + with tempfile.TemporaryDirectory() as tempdir: + parser = ConfigParser() + parser.read_config(config_file) + ckpt_file = os.path.join(tempdir, "model.pt") + save_state(src=parser.get_parsed_content("network_def"), path=ckpt_file) + meta_file = os.path.join(tempdir, "first.json") + parser.export_config_file(config={"first": 1}, filepath=meta_file) + with self.assertRaises(FileNotFoundError): + export_checkpoint( + filepath=os.path.join(tempdir, "model.pt2"), + ckpt_file=ckpt_file, + meta_file=[meta_file, os.path.join(tempdir, "missing.json")], + config_file=config_file, + input_shape=[1, 1, 96, 96, 96], + ) + + @skip_if_no_cuda + def test_network_on_gpu_exports_on_cpu(self): + with tempfile.TemporaryDirectory() as tempdir: + config_file = os.path.join(tempdir, "inference.json") + net_def = {"_target_": "UNet", "spatial_dims": 2, "in_channels": 1, "out_channels": 2} + net_def.update({"channels": [4, 8], "strides": [2]}) + config = {"network_def": net_def, "network": "$@network_def.to('cuda:0')"} + ConfigParser.export_config_file(config=config, filepath=config_file) + ckpt_file = os.path.join(tempdir, "model.pt") + parser = ConfigParser() + parser.read_config(config_file) + save_state(src=parser.get_parsed_content("network_def"), path=ckpt_file) + pt2_file = os.path.join(tempdir, "model.pt2") + + export_checkpoint( + "network", filepath=pt2_file, ckpt_file=ckpt_file, config_file=config_file, input_shape=[1, 1, 16, 16] + ) + module = load_exported_program(pt2_file)[0].module() + + self.assertEqual(next(module.parameters()).device.type, "cpu") + self.assertEqual(module(torch.rand(1, 1, 16, 16)).shape, torch.Size([1, 2, 16, 16])) + def test_meta_files_are_merged(self): config_file = os.path.join(TESTS_PATH, "testing_data", "inference.json") with tempfile.TemporaryDirectory() as tempdir: @@ -139,5 +185,37 @@ def test_meta_files_are_merged(self): self.assertEqual(metadata, {"first": 1, "second": 2}) +class TestLoadExportedModule(unittest.TestCase): + def setUp(self): + self.net = torch.nn.Linear(2, 2).eval() + self.data = torch.rand(1, 2) + self.exported = torch.export.export(self.net, args=(self.data,)) + + def test_load_default_file(self): + with tempfile.TemporaryDirectory() as tempdir: + os.makedirs(os.path.join(tempdir, "bundle", "models")) + save_exported_program( + self.exported, + os.path.join(tempdir, "bundle", "models", "model.pt2"), + include_config_vals=False, + meta_values={"version": "1"}, + more_extra_files={"inference.json": "{}"}, + ) + loaded, meta, extra = load( + "bundle", bundle_dir=tempdir, load_exported_module=True, config_files=("inference.json",) + ) + + torch.testing.assert_close(loaded.module()(self.data), self.net(self.data)) + self.assertEqual(meta, {"version": "1"}) + self.assertEqual(extra, {"inference.json": "{}"}) + + def test_mutually_exclusive(self): + with self.assertRaisesRegex(ValueError, "mutually exclusive"): + load("bundle", load_ts_module=True, load_exported_module=True) + + def test_new_argument_is_last(self): + self.assertEqual(list(inspect.signature(load).parameters)[-1], "load_exported_module") + + if __name__ == "__main__": unittest.main() From 6a8982af78818f9e394711bfc21e493b724662fe Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:30:11 +0000 Subject: [PATCH 10/18] Add export tests for networks Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- tests/networks/nets/dints/test_dints_mixop.py | 7 +++- tests/networks/nets/regunet/test_localnet.py | 8 ++++- tests/networks/nets/regunet/test_regunet.py | 8 ++++- tests/networks/nets/test_ahnet.py | 13 ++++++- tests/networks/nets/test_autoencoder.py | 7 +++- tests/networks/nets/test_basic_unet.py | 7 +++- .../networks/nets/test_basic_unetplusplus.py | 7 +++- tests/networks/nets/test_convnext.py | 8 ++++- tests/networks/nets/test_densenet.py | 8 ++++- tests/networks/nets/test_dints_network.py | 10 +++++- tests/networks/nets/test_discriminator.py | 7 +++- tests/networks/nets/test_dynunet.py | 15 +++++++- tests/networks/nets/test_efficientnet.py | 14 +++++++- tests/networks/nets/test_generator.py | 7 +++- tests/networks/nets/test_globalnet.py | 8 ++++- tests/networks/nets/test_highresnet.py | 9 ++++- tests/networks/nets/test_hovernet.py | 8 ++++- tests/networks/nets/test_milmodel.py | 8 ++++- tests/networks/nets/test_net_adapter.py | 12 ++++++- .../nets/test_patch_gan_dicriminator.py | 36 ++++++++++++++++++- tests/networks/nets/test_quicknat.py | 7 +++- tests/networks/nets/test_resnet.py | 7 ++++ tests/networks/nets/test_segresnet.py | 14 +++++++- tests/networks/nets/test_segresnet_ds.py | 8 ++++- tests/networks/nets/test_senet.py | 14 +++++++- tests/networks/nets/test_unet.py | 22 +++++++++++- tests/networks/nets/test_unetr.py | 8 ++++- tests/networks/nets/test_varautoencoder.py | 9 ++++- tests/networks/nets/test_vit.py | 8 ++++- tests/networks/nets/test_vnet.py | 7 +++- tests/networks/nets/test_voxelmorph.py | 13 ++++++- tests/networks/nets/test_vqvae.py | 18 ++++++++++ 32 files changed, 312 insertions(+), 30 deletions(-) diff --git a/tests/networks/nets/dints/test_dints_mixop.py b/tests/networks/nets/dints/test_dints_mixop.py index ea78514fa5..74a6687146 100644 --- a/tests/networks/nets/dints/test_dints_mixop.py +++ b/tests/networks/nets/dints/test_dints_mixop.py @@ -17,7 +17,7 @@ from parameterized import parameterized from monai.networks.nets.dints import Cell, MixedOp -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES_3D = [ [ @@ -79,6 +79,11 @@ def test_script(self, input_param, ops, weight, input_shape, expected_shape): net = MixedOp(ops=Cell.OPS3D, **input_param) test_script_save(net, torch.randn(input_shape), weight) + @parameterized.expand(TEST_CASES_3D) + def test_export(self, input_param, ops, weight, input_shape, expected_shape): + net = MixedOp(ops=Cell.OPS3D, **input_param) + test_export_save(net, torch.randn(input_shape), weight) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/regunet/test_localnet.py b/tests/networks/nets/regunet/test_localnet.py index d255934495..d5a0936e6e 100644 --- a/tests/networks/nets/regunet/test_localnet.py +++ b/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_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -81,6 +81,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(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) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/regunet/test_regunet.py b/tests/networks/nets/regunet/test_regunet.py index 152a05fc62..d0e1397f38 100644 --- a/tests/networks/nets/regunet/test_regunet.py +++ b/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_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -83,6 +83,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(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) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_ahnet.py b/tests/networks/nets/test_ahnet.py index 7facf9af24..ef28ab8e53 100644 --- a/tests/networks/nets/test_ahnet.py +++ b/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_script_save +from tests.test_utils import skip_if_quick, test_export_save, test_pretrained_networks, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -190,6 +190,17 @@ def test_script(self): test_data = torch.randn(1, 1, 32, 32, 64) test_script_save(net, test_data) + @skip_if_quick + def test_export(self): + # test 2D network + net = AHNet(spatial_dims=2, out_channels=2) + test_data = torch.randn(1, 1, 128, 64) + test_export_save(net, test_data) + # test 3D network + net = AHNet(spatial_dims=3, out_channels=2, psp_block_num=0, upsample_mode="nearest") + test_data = torch.randn(1, 1, 32, 32, 64) + test_export_save(net, test_data) + class TestAHNETWithPretrain(unittest.TestCase): @parameterized.expand( diff --git a/tests/networks/nets/test_autoencoder.py b/tests/networks/nets/test_autoencoder.py index dcf90b809a..879d1f02a8 100644 --- a/tests/networks/nets/test_autoencoder.py +++ b/tests/networks/nets/test_autoencoder.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.layers import Act from monai.networks.nets import AutoEncoder -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") @@ -86,6 +86,11 @@ def test_script(self): test_data = torch.randn(2, 1, 32, 32) test_script_save(net, test_data) + def test_export(self): + net = AutoEncoder(spatial_dims=2, in_channels=1, out_channels=1, channels=(4, 8), strides=(2, 2)) + test_data = torch.randn(2, 1, 32, 32) + test_export_save(net, test_data) + def test_channel_stride_difference(self): with self.assertRaises(ValueError): AutoEncoder(**TEST_CASE_FAIL) diff --git a/tests/networks/nets/test_basic_unet.py b/tests/networks/nets/test_basic_unet.py index edd7b599d4..eee8c23b27 100644 --- a/tests/networks/nets/test_basic_unet.py +++ b/tests/networks/nets/test_basic_unet.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import BasicUNet -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save CASES_1D = [] for mode in ["pixelshuffle", "nontrainable", "deconv", None]: @@ -97,6 +97,11 @@ def test_script(self): test_data = torch.randn(16, 1, 32, 32) test_script_save(net, test_data) + def test_export(self): + net = BasicUNet(spatial_dims=2, in_channels=1, out_channels=3) + test_data = torch.randn(16, 1, 32, 32) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_basic_unetplusplus.py b/tests/networks/nets/test_basic_unetplusplus.py index 37a2813f2e..1c5699a626 100644 --- a/tests/networks/nets/test_basic_unetplusplus.py +++ b/tests/networks/nets/test_basic_unetplusplus.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import BasicUNetPlusPlus -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save CASES_1D = [] for mode in ["pixelshuffle", "nontrainable", "deconv", None]: @@ -104,6 +104,11 @@ def test_script(self): test_data = torch.randn(16, 1, 32, 32) test_script_save(net, test_data) + def test_export(self): + net = BasicUNetPlusPlus(spatial_dims=2, deep_supervision=True, in_channels=1, out_channels=3) + test_data = torch.randn(16, 1, 32, 32) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_convnext.py b/tests/networks/nets/test_convnext.py index fe7b9bc740..977f29a4c9 100644 --- a/tests/networks/nets/test_convnext.py +++ b/tests/networks/nets/test_convnext.py @@ -28,7 +28,7 @@ convnext_large, convnext_tiny, ) -from tests.test_utils import skip_if_quick, test_script_save +from tests.test_utils import skip_if_quick, test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -100,6 +100,12 @@ def test_script(self, model, input_param, input_shape, expected_shape): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_SCRIPT_CASES) + def test_export(self, model, input_param, input_shape, expected_shape): + net = model(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + @parameterized.expand(TEST_REFERENCE_PARAMS) @skip_if_quick def test_reference_parameter_count(self, model, expected_params): diff --git a/tests/networks/nets/test_densenet.py b/tests/networks/nets/test_densenet.py index bc6d6d0094..dbc49ef2df 100644 --- a/tests/networks/nets/test_densenet.py +++ b/tests/networks/nets/test_densenet.py @@ -21,7 +21,7 @@ from monai.networks import eval_mode from monai.networks.nets import DenseNet121, Densenet169, DenseNet264, densenet201 from monai.utils import optional_import -from tests.test_utils import skip_if_downloading_fails, skip_if_quick, test_script_save +from tests.test_utils import skip_if_downloading_fails, skip_if_quick, test_export_save, test_script_save if TYPE_CHECKING: import torchvision @@ -120,6 +120,12 @@ def test_script(self, model, input_param, input_shape, expected_shape): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_SCRIPT_CASES) + def test_export(self, model, input_param, input_shape, expected_shape): + net = model(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_dints_network.py b/tests/networks/nets/test_dints_network.py index fad32c0350..cfff81807d 100644 --- a/tests/networks/nets/test_dints_network.py +++ b/tests/networks/nets/test_dints_network.py @@ -21,7 +21,7 @@ from monai.networks.nets import DiNTS, TopologyInstance, TopologySearch from monai.networks.nets.dints import Cell -from tests.test_utils import skip_if_quick, test_script_save +from tests.test_utils import skip_if_quick, test_export_save, test_script_save TEST_CASES_3D = [ [ @@ -347,6 +347,14 @@ def test_script(self, dints_grid_params, dints_params, input_shape, _): net = DiNTS(**dints_params).to(dints_grid_params["device"]) test_script_save(net, torch.randn(input_shape).to(dints_grid_params["device"])) + @parameterized.expand(TEST_CASES_3D + TEST_CASES_2D) + def test_export(self, dints_grid_params, dints_params, input_shape, _): + grid = TopologyInstance(**dints_grid_params) + dints_grid_params["device"] = "cpu" + dints_params["dints_space"] = grid + net = DiNTS(**dints_params).to(dints_grid_params["device"]) + test_export_save(net, torch.randn(input_shape).to(dints_grid_params["device"])) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_discriminator.py b/tests/networks/nets/test_discriminator.py index 8f460a2450..b5457c5a16 100644 --- a/tests/networks/nets/test_discriminator.py +++ b/tests/networks/nets/test_discriminator.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import Discriminator -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASE_0 = [ {"in_shape": (1, 64, 64), "channels": (2, 4, 8), "strides": (2, 2, 2), "num_res_units": 0}, @@ -54,6 +54,11 @@ def test_script(self): test_data = torch.rand(16, 1, 64, 64) test_script_save(net, test_data) + def test_export(self): + net = Discriminator(in_shape=(1, 64, 64), channels=(2, 4), strides=(2, 2), num_res_units=0) + test_data = torch.rand(16, 1, 64, 64) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_dynunet.py b/tests/networks/nets/test_dynunet.py index a573956296..ebaa5124d7 100644 --- a/tests/networks/nets/test_dynunet.py +++ b/tests/networks/nets/test_dynunet.py @@ -21,7 +21,14 @@ from monai.networks import eval_mode from monai.networks.nets import DynUNet from monai.utils import optional_import -from tests.test_utils import assert_allclose, dict_product, skip_if_no_cuda, skip_if_windows, test_script_save +from tests.test_utils import ( + assert_allclose, + dict_product, + skip_if_no_cuda, + skip_if_windows, + test_export_save, + test_script_save, +) InstanceNorm3dNVFuser, _ = optional_import("apex.normalization", name="InstanceNorm3dNVFuser") @@ -175,6 +182,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASE_DYNUNET_2D[0] + net = DynUNet(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + @skip_if_no_cuda @skip_if_windows diff --git a/tests/networks/nets/test_efficientnet.py b/tests/networks/nets/test_efficientnet.py index e76d5a6d5a..dd08160b4f 100644 --- a/tests/networks/nets/test_efficientnet.py +++ b/tests/networks/nets/test_efficientnet.py @@ -29,7 +29,13 @@ get_efficientnet_image_size, ) from monai.utils import optional_import -from tests.test_utils import skip_if_downloading_fails, skip_if_quick, test_pretrained_networks, test_script_save +from tests.test_utils import ( + skip_if_downloading_fails, + skip_if_quick, + test_export_save, + test_pretrained_networks, + test_script_save, +) TESTS_PATH = Path(__file__).parents[2] @@ -377,6 +383,12 @@ def test_script(self): test_data = torch.randn(1, 3, 224, 224) test_script_save(net, test_data) + def test_export(self): + with skip_if_downloading_fails(): + net = EfficientNetBN(model_name="efficientnet-b0", spatial_dims=2, in_channels=3, num_classes=1000) + test_data = torch.randn(1, 3, 224, 224) + test_export_save(net, test_data) + class TestExtractFeatures(unittest.TestCase): @parameterized.expand(CASE_EXTRACT_FEATURES) diff --git a/tests/networks/nets/test_generator.py b/tests/networks/nets/test_generator.py index 9ec08194e9..50ad0513c0 100644 --- a/tests/networks/nets/test_generator.py +++ b/tests/networks/nets/test_generator.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import Generator -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASE_0 = [ {"latent_shape": (64,), "start_shape": (8, 8, 8), "channels": (8, 4, 1), "strides": (2, 2, 2), "num_res_units": 0}, @@ -54,6 +54,11 @@ def test_script(self): test_data = torch.rand(16, 64) test_script_save(net, test_data) + def test_export(self): + net = Generator(latent_shape=(64,), start_shape=(8, 8, 8), channels=(8, 1), strides=(2, 2), num_res_units=2) + test_data = torch.rand(16, 64) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_globalnet.py b/tests/networks/nets/test_globalnet.py index 84e3d466de..5d504f5f55 100644 --- a/tests/networks/nets/test_globalnet.py +++ b/tests/networks/nets/test_globalnet.py @@ -21,7 +21,7 @@ from monai.networks.blocks import Warp from monai.networks.nets import GlobalNet from monai.networks.nets.regunet import AffineHead -from tests.test_utils import assert_allclose, test_script_save +from tests.test_utils import assert_allclose, test_export_save, test_script_save TEST_CASES_AFFINE_TRANSFORM = [ [ @@ -98,6 +98,12 @@ def test_script(self, input_param, input_shape, _): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_CASES_GLOBAL_NET) + def test_export(self, input_param, input_shape, _): + net = GlobalNet(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + @parameterized.expand(TEST_CASES_GLOBAL_NET) def test_half_precision_forward(self, input_param, input_shape, expected_shape): """Check that a GlobalNet cast to half precision can run a forward pass and returns float16. diff --git a/tests/networks/nets/test_highresnet.py b/tests/networks/nets/test_highresnet.py index 1384dfaeff..3ca769b8f6 100644 --- a/tests/networks/nets/test_highresnet.py +++ b/tests/networks/nets/test_highresnet.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import HighResNet -from tests.test_utils import DistTestCase, TimedCall, test_script_save +from tests.test_utils import DistTestCase, TimedCall, test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -62,6 +62,13 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data, rtol=1e-4, atol=1e-4) + @TimedCall(seconds=800, force_quit=True) + def test_export(self): + input_param, input_shape, expected_shape = TEST_CASE_1 + net = HighResNet(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data, rtol=1e-4, atol=1e-4) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_hovernet.py b/tests/networks/nets/test_hovernet.py index 58657e6dea..7064dc516c 100644 --- a/tests/networks/nets/test_hovernet.py +++ b/tests/networks/nets/test_hovernet.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode, train_mode from monai.networks.nets import HoVerNet from monai.networks.nets.hovernet import _DenseLayerDecoder -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -184,6 +184,12 @@ def test_script(self): test_data = torch.randn(1, 3, 256, 256) test_script_save(net, test_data) + @parameterized.expand([[True], [False]]) + def test_export(self, padding_flag): + net = HoVerNet(mode=HoVerNet.Mode.FAST, decoder_padding=padding_flag) + test_data = torch.randn(1, 3, 256, 256) + test_export_save(net, test_data) + def test_ill_input_shape(self): net = HoVerNet(mode=HoVerNet.Mode.FAST) with eval_mode(net): diff --git a/tests/networks/nets/test_milmodel.py b/tests/networks/nets/test_milmodel.py index 15fda15a11..3d7e4014a5 100644 --- a/tests/networks/nets/test_milmodel.py +++ b/tests/networks/nets/test_milmodel.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.nets import MILModel from monai.utils.module import optional_import -from tests.test_utils import skip_if_downloading_fails, test_script_save +from tests.test_utils import skip_if_downloading_fails, test_export_save, test_script_save models, _ = optional_import("torchvision.models") @@ -87,6 +87,12 @@ def test_script(self): test_data = torch.randn(input_shape, dtype=torch.float) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, expected_shape = TEST_CASE_MILMODEL[0] + net = MILModel(**input_param) + test_data = torch.randn(input_shape, dtype=torch.float) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_net_adapter.py b/tests/networks/nets/test_net_adapter.py index 08344900e4..0c1d275f1a 100644 --- a/tests/networks/nets/test_net_adapter.py +++ b/tests/networks/nets/test_net_adapter.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import NetAdapter, resnet18 -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -63,6 +63,16 @@ def test_script(self, input_param, input_shape, expected_shape): test_data = torch.randn(input_shape).to("cpu") test_script_save(net, test_data) + @parameterized.expand([TEST_CASE_0]) + def test_export(self, input_param, input_shape, expected_shape): + spatial_dims = input_param["dim"] + stride = (1, 2, 2)[:spatial_dims] + model = resnet18(spatial_dims=spatial_dims, conv1_t_stride=stride) + input_param["model"] = model + net = NetAdapter(**input_param).to("cpu") + test_data = torch.randn(input_shape).to("cpu") + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_patch_gan_dicriminator.py b/tests/networks/nets/test_patch_gan_dicriminator.py index 184f76fa9d..60985a580a 100644 --- a/tests/networks/nets/test_patch_gan_dicriminator.py +++ b/tests/networks/nets/test_patch_gan_dicriminator.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import MultiScalePatchDiscriminator, PatchDiscriminator -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_PATCHGAN = [ [ @@ -140,6 +140,22 @@ def test_script(self): i = torch.rand([1, 3, 256, 512]) test_script_save(net, i) + def test_export(self): + net = PatchDiscriminator( + num_layers_d=3, + spatial_dims=2, + channels=8, + in_channels=3, + out_channels=1, + kernel_size=3, + activation="LEAKYRELU", + norm="instance", + bias=False, + dropout=0.1, + ) + i = torch.rand([1, 3, 256, 512]) + test_export_save(net, i) + class TestMultiscalePatchGAN(unittest.TestCase): @parameterized.expand(TEST_MULTISCALE_PATCHGAN) @@ -174,6 +190,24 @@ def test_script(self): i = torch.rand([1, 3, 256, 512]) test_script_save(net, i) + def test_export(self): + net = MultiScalePatchDiscriminator( + num_d=2, + num_layers_d=3, + spatial_dims=2, + channels=8, + in_channels=3, + out_channels=1, + kernel_size=3, + activation="LEAKYRELU", + norm="instance", + bias=False, + dropout=0.1, + minimum_size_im=256, + ) + i = torch.rand([1, 3, 256, 512]) + test_export_save(net, i) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_quicknat.py b/tests/networks/nets/test_quicknat.py index 913f30b661..c195a4c846 100644 --- a/tests/networks/nets/test_quicknat.py +++ b/tests/networks/nets/test_quicknat.py @@ -20,7 +20,7 @@ from monai.networks.nets import Quicknat from monai.networks.nets.quicknat import SkipConnectionWithIdx from monai.utils import optional_import -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save _, has_se = optional_import("squeeze_and_excitation") @@ -109,6 +109,11 @@ def test_script(self): test_data = torch.randn(16, 1, 32, 32) test_script_save(net, test_data) + def test_export(self): + net = Quicknat(num_classes=1, num_channels=1) + test_data = torch.randn(16, 1, 32, 32) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_resnet.py b/tests/networks/nets/test_resnet.py index 241f57c78d..52e2d19f0b 100644 --- a/tests/networks/nets/test_resnet.py +++ b/tests/networks/nets/test_resnet.py @@ -43,6 +43,7 @@ skip_if_downloading_fails, skip_if_no_cuda, skip_if_quick, + test_export_save, test_script_save, ) @@ -316,6 +317,12 @@ def test_script(self, model, input_param, input_shape, expected_shape): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_SCRIPT_CASES) + def test_export(self, model, input_param, input_shape, expected_shape): + net = model(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + @SkipIfNoModule("hf_hub_download") class TestExtractFeatures(unittest.TestCase): diff --git a/tests/networks/nets/test_segresnet.py b/tests/networks/nets/test_segresnet.py index 1536d33853..c587f1e087 100644 --- a/tests/networks/nets/test_segresnet.py +++ b/tests/networks/nets/test_segresnet.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.nets import SegResNet, SegResNetVAE from monai.utils import UpsampleMode -from tests.test_utils import dict_product, test_script_save +from tests.test_utils import dict_product, test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -83,6 +83,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, expected_shape = TEST_CASE_SEGRESNET[0] + net = SegResNet(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + class TestResNetVAE(unittest.TestCase): @parameterized.expand(TEST_CASE_SEGRESNET_VAE) @@ -98,6 +104,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, expected_shape = TEST_CASE_SEGRESNET_VAE[0] + net = SegResNetVAE(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_segresnet_ds.py b/tests/networks/nets/test_segresnet_ds.py index 4a2cf40e6f..48d8365c9e 100644 --- a/tests/networks/nets/test_segresnet_ds.py +++ b/tests/networks/nets/test_segresnet_ds.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import SegResNetDS, SegResNetDS2 -from tests.test_utils import dict_product, test_script_save +from tests.test_utils import dict_product, test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -141,6 +141,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASE_SEGRESNET_DS[0] + net = SegResNetDS(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_senet.py b/tests/networks/nets/test_senet.py index 90d711d0d9..11c04ca411 100644 --- a/tests/networks/nets/test_senet.py +++ b/tests/networks/nets/test_senet.py @@ -24,7 +24,13 @@ from monai.networks import eval_mode from monai.networks.nets import SENet, SENet154, SEResNet50, SEResNet101, SEResNet152, SEResNext50, SEResNext101 from monai.utils import optional_import -from tests.test_utils import test_is_quick, test_pretrained_networks, test_script_save, testing_data_config +from tests.test_utils import ( + test_export_save, + test_is_quick, + test_pretrained_networks, + test_script_save, + testing_data_config, +) if TYPE_CHECKING: import pretrainedmodels @@ -74,6 +80,12 @@ def test_script(self, net, net_args): input_data = torch.randn(2, 2, 64, 64, 64) test_script_save(net, input_data) + @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6]) + def test_export(self, net, net_args): + net = net(**net_args) + input_data = torch.randn(2, 2, 64, 64, 64) + test_export_save(net, input_data) + class TestPretrainedSENET(unittest.TestCase): def setUp(self): diff --git a/tests/networks/nets/test_unet.py b/tests/networks/nets/test_unet.py index 7a6d0e98bb..6402534b39 100644 --- a/tests/networks/nets/test_unet.py +++ b/tests/networks/nets/test_unet.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.layers import Act, Norm from monai.networks.nets import UNet -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -179,6 +179,13 @@ def test_script(self): test_data = torch.randn(16, 1, 32, 32) test_script_save(net, test_data) + def test_export(self): + net = UNet( + spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2), num_res_units=0 + ) + test_data = torch.randn(16, 1, 32, 32) + test_export_save(net, test_data) + def test_script_without_running_stats(self): net = UNet( spatial_dims=2, @@ -192,6 +199,19 @@ def test_script_without_running_stats(self): test_data = torch.randn(16, 1, 16, 4) test_script_save(net, test_data) + def test_export_without_running_stats(self): + net = UNet( + spatial_dims=2, + in_channels=1, + out_channels=3, + channels=(16, 32, 64), + strides=(2, 2), + num_res_units=0, + norm=("batch", {"track_running_stats": False}), + ) + test_data = torch.randn(16, 1, 16, 4) + test_export_save(net, test_data) + def test_ill_input_shape(self): net = UNet(spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2)) with eval_mode(net): diff --git a/tests/networks/nets/test_unetr.py b/tests/networks/nets/test_unetr.py index 5d4faa3979..bf11dc3d5a 100644 --- a/tests/networks/nets/test_unetr.py +++ b/tests/networks/nets/test_unetr.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets.unetr import UNETR -from tests.test_utils import dict_product, skip_if_quick, test_script_save +from tests.test_utils import dict_product, skip_if_quick, test_export_save, test_script_save TEST_CASE_UNETR = [ [ @@ -124,6 +124,12 @@ def test_script(self, input_param, input_shape, _): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_CASE_UNETR) + def test_export(self, input_param, input_shape, _): + net = UNETR(**(input_param)) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_varautoencoder.py b/tests/networks/nets/test_varautoencoder.py index a8dc715301..377819149b 100644 --- a/tests/networks/nets/test_varautoencoder.py +++ b/tests/networks/nets/test_varautoencoder.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.layers import Act from monai.networks.nets import VarAutoEncoder -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") @@ -122,6 +122,13 @@ def test_script(self): test_data = torch.randn(2, 1, 32, 32) test_script_save(net, test_data, rtol=1e-3, atol=1e-3) + def test_export(self): + net = VarAutoEncoder( + spatial_dims=2, in_shape=(1, 32, 32), out_channels=1, latent_size=2, channels=(4, 8), strides=(2, 2) + ) + test_data = torch.randn(2, 1, 32, 32) + test_export_save(net, test_data, rtol=1e-3, atol=1e-3) + def test_reparameterize_eval_returns_mu(self): """A VarAutoEncoder latent code is deterministic at eval (equals mu) and stochastic at train. diff --git a/tests/networks/nets/test_vit.py b/tests/networks/nets/test_vit.py index 949b0f4409..430b8411eb 100644 --- a/tests/networks/nets/test_vit.py +++ b/tests/networks/nets/test_vit.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets.vit import ViT -from tests.test_utils import dict_product, skip_if_quick, test_script_save +from tests.test_utils import dict_product, skip_if_quick, test_export_save, test_script_save TEST_CASE_Vit = [ ( @@ -108,6 +108,12 @@ def test_script(self, input_param, input_shape, _): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand(TEST_CASE_Vit[:1]) + def test_export(self, input_param, input_shape, _): + net = ViT(**(input_param)) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + def test_access_attn_matrix(self): # input format in_channels = 1 diff --git a/tests/networks/nets/test_vnet.py b/tests/networks/nets/test_vnet.py index 6c93893480..a6b3fb087d 100644 --- a/tests/networks/nets/test_vnet.py +++ b/tests/networks/nets/test_vnet.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import VNet -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -76,6 +76,11 @@ def test_script(self): test_data = torch.randn(1, 1, 32, 32, 32) test_script_save(net, test_data) + def test_export(self): + net = VNet(spatial_dims=3, in_channels=1, out_channels=3, dropout_dim=3) + test_data = torch.randn(1, 1, 32, 32, 32) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/nets/test_voxelmorph.py b/tests/networks/nets/test_voxelmorph.py index 1a04bab568..6322efe784 100644 --- a/tests/networks/nets/test_voxelmorph.py +++ b/tests/networks/nets/test_voxelmorph.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets import VoxelMorph, VoxelMorphUNet -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -263,6 +263,17 @@ def test_script(self): test_data = torch.randn(1, 2, 96, 96) test_script_save(net, test_data) + def test_export(self): + net = VoxelMorphUNet( + spatial_dims=2, + in_channels=2, + unet_out_channels=32, + channels=(16, 32, 32, 32, 32, 32), + final_conv_channels=(16, 16), + ) + test_data = torch.randn(1, 2, 96, 96) + test_export_save(net, test_data) + @parameterized.expand(ILL_CASES) def test_ill_input_hyper_params(self, input_param): with self.assertRaises(ValueError): diff --git a/tests/networks/nets/test_vqvae.py b/tests/networks/nets/test_vqvae.py index 13a34f9657..cf551fb5c1 100644 --- a/tests/networks/nets/test_vqvae.py +++ b/tests/networks/nets/test_vqvae.py @@ -18,6 +18,7 @@ from monai.networks import eval_mode from monai.networks.nets.vqvae import VQVAE +from tests.test_utils import test_export_save TEST_CASES = [ [ @@ -143,6 +144,23 @@ def test_shape_with_checkpoint(self, input_param, input_shape, expected_shape): # test_data = torch.randn(1, 1, 16, 16) # test_script_save(net, test_data) + def test_export(self): + net = VQVAE( + spatial_dims=2, + in_channels=1, + out_channels=1, + downsample_parameters=((2, 4, 1, 1),) * 2, + upsample_parameters=((2, 4, 1, 1, 0),) * 2, + num_res_layers=1, + channels=(8, 8), + num_res_channels=(8, 8), + num_embeddings=16, + embedding_dim=8, + ddp_sync=False, + ) + test_data = torch.randn(1, 1, 16, 16) + test_export_save(net, test_data) + def test_channels_not_same_size_of_num_res_channels(self): with self.assertRaises(ValueError): VQVAE( From 885ba7ab8f604c18065d8f494db8d003461bb41c Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:30:12 +0000 Subject: [PATCH 11/18] Add export tests for blocks Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- tests/networks/blocks/test_dynunet_block.py | 17 +++++++++++++- tests/networks/blocks/test_fpn_block.py | 18 ++++++++++++++- tests/networks/blocks/test_se_block.py | 8 ++++++- tests/networks/blocks/test_se_blocks.py | 14 +++++++++++- tests/networks/blocks/test_selfattention.py | 19 +++++++++++++++- .../networks/blocks/test_subpixel_upsample.py | 8 ++++++- tests/networks/blocks/test_unetr_block.py | 22 ++++++++++++++++++- 7 files changed, 99 insertions(+), 7 deletions(-) diff --git a/tests/networks/blocks/test_dynunet_block.py b/tests/networks/blocks/test_dynunet_block.py index 4cf68912be..c63401ea04 100644 --- a/tests/networks/blocks/test_dynunet_block.py +++ b/tests/networks/blocks/test_dynunet_block.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.blocks.dynunet_block import UnetBasicBlock, UnetResBlock, UnetUpBlock, get_padding -from tests.test_utils import dict_product, test_script_save +from tests.test_utils import dict_product, test_export_save, test_script_save TEST_CASE_RES_BASIC_BLOCK = [] for params in dict_product( @@ -91,6 +91,13 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand([[UnetResBlock], [UnetBasicBlock]]) + def test_export(self, net_type): + input_param, input_shape, _ = TEST_CASE_RES_BASIC_BLOCK[0] + net = net_type(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + class TestUpBlock(unittest.TestCase): @parameterized.expand(TEST_UP_BLOCK) @@ -108,6 +115,14 @@ def test_script(self): skip_data = torch.randn(skip_shape) test_script_save(net, test_data, skip_data) + def test_export(self): + input_param, input_shape, _, skip_shape = TEST_UP_BLOCK[0] + + net = UnetUpBlock(**input_param) + test_data = torch.randn(input_shape) + skip_data = torch.randn(skip_shape) + test_export_save(net, test_data, skip_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/blocks/test_fpn_block.py b/tests/networks/blocks/test_fpn_block.py index 4b45850159..b38036014c 100644 --- a/tests/networks/blocks/test_fpn_block.py +++ b/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_script_save +from tests.test_utils import test_export_save, test_script_save _, has_torchvision = optional_import("torchvision") @@ -63,6 +63,15 @@ def test_script(self, input_param, input_shape, expected_shape): data["feat1"] = torch.rand(input_shape[1]) test_script_save(net, data) + @parameterized.expand(TEST_CASES) + def test_export(self, input_param, input_shape, expected_shape): + # test whether support export + 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) + @unittest.skipUnless(has_torchvision, "Requires torchvision") class TestFPN(unittest.TestCase): @@ -81,6 +90,13 @@ def test_script(self, input_param, input_shape, expected_shape): data = torch.rand(input_shape) test_script_save(net, data) + @parameterized.expand(TEST_CASES2) + def test_export(self, input_param, input_shape, expected_shape): + # test whether support export + 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) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/blocks/test_se_block.py b/tests/networks/blocks/test_se_block.py index d799cd095c..715722edb6 100644 --- a/tests/networks/blocks/test_se_block.py +++ b/tests/networks/blocks/test_se_block.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.blocks import SEBlock from monai.networks.layers.factories import Act, Norm -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save device = "cuda" if torch.cuda.is_available() else "cpu" @@ -76,6 +76,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASES[0] + net = SEBlock(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + def test_ill_arg(self): with self.assertRaises(ValueError): SEBlock(spatial_dims=1, in_channels=4, n_chns_1=2, n_chns_2=3, n_chns_3=4, r=100) diff --git a/tests/networks/blocks/test_se_blocks.py b/tests/networks/blocks/test_se_blocks.py index b40f3a0955..7978d768a9 100644 --- a/tests/networks/blocks/test_se_blocks.py +++ b/tests/networks/blocks/test_se_blocks.py @@ -18,7 +18,7 @@ from monai.networks import eval_mode from monai.networks.blocks import ChannelSELayer, ResidualSELayer -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASES = [ # single channel 3D, batch 16 [{"spatial_dims": 2, "in_channels": 4, "r": 3}, (7, 4, 64, 48), (7, 4, 64, 48)], # 4-channel 2D, batch 7 @@ -54,6 +54,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASES[0] + net = ChannelSELayer(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + def test_ill_arg(self): with self.assertRaises(ValueError): ChannelSELayer(spatial_dims=1, in_channels=4, r=100) @@ -73,6 +79,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASES[0] + net = ResidualSELayer(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/blocks/test_selfattention.py b/tests/networks/blocks/test_selfattention.py index af52918612..b268c7b14f 100644 --- a/tests/networks/blocks/test_selfattention.py +++ b/tests/networks/blocks/test_selfattention.py @@ -22,7 +22,7 @@ from monai.networks.blocks.selfattention import SABlock from monai.networks.layers.factories import RelPosEmbedding 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, test_script_save einops, has_einops = optional_import("einops") @@ -207,6 +207,23 @@ def test_script(self, include_fc, use_combined_linear): test_data = torch.randn(input_shape) test_script_save(net, test_data) + @parameterized.expand([[True, False], [True, True], [False, True], [False, False]]) + @skipUnless(has_einops, "Requires einops") + def test_export(self, include_fc, use_combined_linear): + input_param = { + "hidden_size": 360, + "num_heads": 4, + "dropout_rate": 0.0, + "rel_pos_embedding": None, + "input_size": (16, 32), + "include_fc": include_fc, + "use_combined_linear": use_combined_linear, + } + net = SABlock(**input_param) + input_shape = (2, 512, 360) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + @skipUnless(has_einops, "Requires einops") def test_flash_attention(self): for causal in [True, False]: diff --git a/tests/networks/blocks/test_subpixel_upsample.py b/tests/networks/blocks/test_subpixel_upsample.py index f4075f5099..3d375d7189 100644 --- a/tests/networks/blocks/test_subpixel_upsample.py +++ b/tests/networks/blocks/test_subpixel_upsample.py @@ -20,7 +20,7 @@ from monai.networks import eval_mode from monai.networks.blocks import SubpixelUpsample from monai.networks.layers.factories import Conv -from tests.test_utils import test_script_save +from tests.test_utils import test_export_save, test_script_save TEST_CASE_SUBPIXEL = [] for inch in range(1, 5): @@ -81,6 +81,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASE_SUBPIXEL[0] + net = SubpixelUpsample(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/networks/blocks/test_unetr_block.py b/tests/networks/blocks/test_unetr_block.py index 0073efc609..50fe805480 100644 --- a/tests/networks/blocks/test_unetr_block.py +++ b/tests/networks/blocks/test_unetr_block.py @@ -19,7 +19,7 @@ from monai.networks import eval_mode from monai.networks.blocks.dynunet_block import get_padding from monai.networks.blocks.unetr_block import UnetrBasicBlock, UnetrPrUpBlock, UnetrUpBlock -from tests.test_utils import dict_product, test_script_save +from tests.test_utils import dict_product, test_export_save, test_script_save def _get_out_size(params): @@ -118,6 +118,13 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_CASE_UNETR_BASIC_BLOCK[0] + net = UnetrBasicBlock(**input_param) + with eval_mode(net): + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + class TestUpBlock(unittest.TestCase): @parameterized.expand(TEST_UP_BLOCK) @@ -134,6 +141,13 @@ def test_script(self): skip_data = torch.randn(skip_shape) test_script_save(net, test_data, skip_data) + def test_export(self): + input_param, input_shape, _, skip_shape = TEST_UP_BLOCK[0] + net = UnetrUpBlock(**input_param) + test_data = torch.randn(input_shape) + skip_data = torch.randn(skip_shape) + test_export_save(net, test_data, skip_data) + class TestPrUpBlock(unittest.TestCase): @parameterized.expand(TEST_PRUP_BLOCK) @@ -149,6 +163,12 @@ def test_script(self): test_data = torch.randn(input_shape) test_script_save(net, test_data) + def test_export(self): + input_param, input_shape, _ = TEST_PRUP_BLOCK[0] + net = UnetrPrUpBlock(**input_param) + test_data = torch.randn(input_shape) + test_export_save(net, test_data) + if __name__ == "__main__": unittest.main() From c8d0735f831eaf77a4a952d4f2f7dcdefe980bb9 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 19:30:14 +0000 Subject: [PATCH 12/18] Add export tests for apps Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- .../apps/detection/networks/test_retinanet.py | 35 ++++++++++++++++++- .../networks/test_retinanet_detector.py | 16 ++++++++- tests/apps/detection/utils/test_anchor_box.py | 18 +++++++++- 3 files changed, 66 insertions(+), 3 deletions(-) diff --git a/tests/apps/detection/networks/test_retinanet.py b/tests/apps/detection/networks/test_retinanet.py index eb491361ea..2a5e28becc 100644 --- a/tests/apps/detection/networks/test_retinanet.py +++ b/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, test_script_save _, has_torchvision = optional_import("torchvision") _, has_onnxruntime = optional_import("onnxruntime") @@ -169,6 +169,39 @@ def test_script(self, model, input_param, input_shape): if idx == 2: test_script_save(net, data) + @parameterized.expand(TEST_CASES_TS) + def test_export(self, model, input_param, input_shape): + try: + idx = int(self.id().split("test_export_")[-1]) + except Exception: + idx = 0 + idx %= 3 + # test whether support export + data = torch.randn(input_shape) + backbone = model(**input_param) + if idx == 0: + test_export_save(backbone, data) + return + feature_extractor = resnet_fpn_feature_extractor( + backbone=backbone, + spatial_dims=input_param["spatial_dims"], + pretrained_backbone=input_param["pretrained"], + trainable_backbone_layers=None, + returned_layers=[1, 2], + ) + if idx == 1: + test_export_save(feature_extractor, data) + return + net = RetinaNet( + spatial_dims=input_param["spatial_dims"], + num_classes=input_param["num_classes"], + num_anchors=num_anchors, + feature_extractor=feature_extractor, + size_divisible=32, + ) + if idx == 2: + test_export_save(net, data) + @parameterized.expand(TEST_CASES_TS) @unittest.skipUnless(has_onnxruntime, "onnxruntime not installed") def test_onnx(self, model, input_param, input_shape): diff --git a/tests/apps/detection/networks/test_retinanet_detector.py b/tests/apps/detection/networks/test_retinanet_detector.py index 6ac1efd734..9836de8484 100644 --- a/tests/apps/detection/networks/test_retinanet_detector.py +++ b/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, test_script_save _, has_torchvision = optional_import("torchvision") @@ -197,6 +197,20 @@ def test_script(self, input_param, input_shape): input_data = torch.randn(input_shape) test_script_save(detector.network, input_data) + @parameterized.expand(TEST_CASES_TS) + def test_export(self, input_param, input_shape): + # test whether support export + returned_layers = [1] + anchor_generator = AnchorGeneratorWithAnchorShape( + feature_map_scales=(1, 2), base_anchor_shapes=((8,) * input_param["spatial_dims"],) + ) + detector = retinanet_resnet50_fpn_detector( + **input_param, anchor_generator=anchor_generator, returned_layers=returned_layers + ) + with eval_mode(detector): + input_data = torch.randn(input_shape) + test_export_save(detector.network, input_data) + if __name__ == "__main__": unittest.main() diff --git a/tests/apps/detection/utils/test_anchor_box.py b/tests/apps/detection/utils/test_anchor_box.py index 4ae198e409..86fe8ec58d 100644 --- a/tests/apps/detection/utils/test_anchor_box.py +++ b/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, test_script_save _, has_torchvision = optional_import("torchvision") @@ -77,6 +77,14 @@ def test_script_2d(self, input_param, image_shape, feature_maps_shapes): feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes) test_script_save(anchor, images, feature_maps) + @parameterized.expand(TEST_CASES_2D) + def test_export_2d(self, input_param, image_shape, feature_maps_shapes): + # test whether support 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_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 @@ -85,6 +93,14 @@ def test_script_3d(self, input_param, image_shape, feature_maps_shapes): feature_maps = tuple(torch.rand(fs) for fs in feature_maps_shapes) test_script_save(anchor, images, feature_maps) + @parameterized.expand(TEST_CASES_SHAPE_3D) + def test_export_3d(self, input_param, image_shape, feature_maps_shapes): + # test whether support 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_export_save(anchor, images, feature_maps) + if __name__ == "__main__": unittest.main() From 70734777c57180fe631c0a8735b9dfdf3ebe30ce Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 21:22:50 +0000 Subject: [PATCH 13/18] Document .pt2 in the API docs Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- docs/source/bundle.rst | 1 + docs/source/data.rst | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/docs/source/bundle.rst b/docs/source/bundle.rst index fdf745e951..5fa8d70997 100644 --- a/docs/source/bundle.rst +++ b/docs/source/bundle.rst @@ -38,6 +38,7 @@ Model Bundle `Scripts` --------- .. autofunction:: ckpt_export +.. autofunction:: export_checkpoint .. autofunction:: trt_export .. autofunction:: onnx_export .. autofunction:: download diff --git a/docs/source/data.rst b/docs/source/data.rst index f53588f423..f4b8a72a65 100644 --- a/docs/source/data.rst +++ b/docs/source/data.rst @@ -227,6 +227,11 @@ Utilities .. automodule:: monai.data.utils :members: +Exported programs +~~~~~~~~~~~~~~~~~ +.. autofunction:: monai.data.save_exported_program +.. autofunction:: monai.data.load_exported_program + Partition Dataset ~~~~~~~~~~~~~~~~~ .. autofunction:: monai.data.partition_dataset From b64dbcf96f1fc8906385c8dff6de84b42d63b616 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sat, 7 Nov 2026 21:22:52 +0000 Subject: [PATCH 14/18] Add .pt2 to the bundle specification Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- docs/source/mb_specification.rst | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/source/mb_specification.rst b/docs/source/mb_specification.rst index 56d660e35c..9147ca8414 100644 --- a/docs/source/mb_specification.rst +++ b/docs/source/mb_specification.rst @@ -6,9 +6,9 @@ MONAI Bundle Specification Overview ======== -This is the specification for the MONAI Bundle (MB) format of portable described deep learning models. The objective of a MB is to define a packaged network or model which includes the critical information necessary to allow users and programs to understand how the model is used and for what purpose. A bundle includes the stored weights of a single network as a pickled state dictionary plus optionally a Torchscript object and/or an ONNX object. Additional JSON files are included to store metadata about the model, information for constructing training, inference, and post-processing transform sequences, plain-text description, legal information, and other data the model creator wishes to include. +This is the specification for the MONAI Bundle (MB) format of portable described deep learning models. The objective of a MB is to define a packaged network or model which includes the critical information necessary to allow users and programs to understand how the model is used and for what purpose. A bundle includes the stored weights of a single network as a pickled state dictionary plus optionally a Torchscript object, a `torch.export` program and/or an ONNX object. Additional JSON files are included to store metadata about the model, information for constructing training, inference, and post-processing transform sequences, plain-text description, legal information, and other data the model creator wishes to include. -This specification defines the directory structure a bundle must have and the necessary files it must contain. Additional files may be included and the directory packaged into a zip file or included as extra files directly in a Torchscript file. +This specification defines the directory structure a bundle must have and the necessary files it must contain. Additional files may be included and the directory packaged into a zip file or included as extra files directly in a Torchscript or `torch.export` file. Directory Structure =================== @@ -24,6 +24,7 @@ A MONAI Bundle is defined primarily as a directory with a set of specifically na ┣━ models ┃ ┣━ model.pt ┃ ┣━ *model.ts + ┃ ┣━ *model.pt2 ┃ ┗━ *model.onnx ┗━ docs ┣━ *README.md @@ -39,6 +40,7 @@ The following files are **required** to be present with the given filenames for The following files are optional but must have these names in the directory given above: * **model.ts**: the Torchscript saved model if the model is compatible with being saved correctly in this format. +* **model.pt2**: the `torch.export` program saved with `torch.export.save`, if the model can be exported. It only accepts the input shapes it was exported with, unless it was exported with dynamic shapes. Loading it uses pickle, so it should only be loaded from a trusted source. * **model.onnx**: the ONNX model if the model is compatible with being saved correctly in this format. * **README.md**: plain-language information on the model, how to use it, author information, etc. in Markdown format. * **license.txt**: software license attached to the data, can be left blank if no license needed. @@ -54,6 +56,8 @@ The Torchscript file format is also just a zip file with a specific structure. W The `bundle` submodule of MONAI contains a number of command line programs. To produce a Torchscript bundle use `ckpt_export` with a set of specified components such as the saved weights file and metadata file. Config files can be provided as JSON or YAML dictionaries defining Python constructs used by the `ConfigParser`, however regardless of format the produced bundle Torchscript object will store the files as JSON. +The `torch.export` file format is also a zip file. A MB-compliant `.pt2` file can be created with `save_exported_program` in the same way, by including the contents of `metadata.json` as the `meta_values` argument and other files as `more_extra_files` entries. These can be retrieved with `load_exported_program`. To produce such a file from a bundle use `export_checkpoint`, which takes the same components as `ckpt_export` and also stores config files as JSON. Since the program is traced with an example input, its shape is taken from the `input_shape` argument or, if that is not given, from `metadata.json`. + metadata.json File ================== From 74bab30785f03f6e717e5f539b5443ee04d5f77d Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sun, 8 Nov 2026 12:15:11 +0000 Subject: [PATCH 15/18] Fix device and keyword inputs in convert_to_export verification Verify on the model's device by default instead of CUDA when available, since the exported graph can keep the device it was exported on. Move tensors nested in dict, list and tuple inputs with copy_to_device, and pass the forward keyword inputs given through kwargs to both checks. Raised by CodeRabbit: https://github.com/Project-MONAI/MONAI/pull/9264#discussion_r4506119986 https://github.com/Project-MONAI/MONAI/pull/9264#discussion_r4506119992 https://github.com/Project-MONAI/MONAI/pull/9264#discussion_r4506120002 Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/networks/utils.py | 23 ++++++++------ tests/networks/test_convert_to_export.py | 40 ++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 10 deletions(-) diff --git a/monai/networks/utils.py b/monai/networks/utils.py index ed2189d87e..6633d8a59c 100644 --- a/monai/networks/utils.py +++ b/monai/networks/utils.py @@ -31,7 +31,7 @@ from monai.apps.utils import get_logger from monai.config import PathLike -from monai.utils.misc import ensure_tuple, save_obj, set_determinism +from monai.utils.misc import copy_to_device, ensure_tuple, save_obj, set_determinism from monai.utils.module import look_up_option, optional_import from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor @@ -893,12 +893,14 @@ def convert_to_export( verify: whether to verify that the exported program, reloaded from ``filename_or_obj`` if given, produces the same outputs as ``model``. inputs: input test data for export and verification. Should be a sequence of - tensors that map to positional arguments of ``model()``. + tensors, or of tuples, lists or dicts of tensors, that map to positional arguments of ``model()``. dynamic_shapes: dynamic shape specifications passed to :func:`torch.export.export`. - device: target device to verify the model. If None, uses CUDA if available. + device: target device to verify the model. If None, uses the device of the model's parameters, because + the exported program can keep the device it was exported on, for example from ``device=x.device``. rtol: the relative tolerance when comparing outputs. atol: the absolute tolerance when comparing outputs. - kwargs: additional keyword arguments for :func:`torch.export.export`. + kwargs: additional keyword arguments for :func:`torch.export.export`. Its ``kwargs`` argument, the keyword + inputs of ``model()``, is also used for verification. Returns: A :class:`torch.export.ExportedProgram` representing the exported model. @@ -918,25 +920,26 @@ def convert_to_export( torch.export.save(exported, filename_or_obj, extra_files=extra_files) if verify: + # the model and the returned program share weights, so the model is moved back after the check + model_device = next(itertools.chain(model.parameters(), model.buffers()), torch.empty(0)).device if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = model_device - verify_inputs = [i.to(device) if isinstance(i, torch.Tensor) else i for i in inputs] + verify_inputs = copy_to_device(list(inputs), device, non_blocking=False) + verify_kwargs = copy_to_device(dict(kwargs.get("kwargs") or {}), device, non_blocking=False) if filename_or_obj is not None: loaded_module = torch.export.load(filename_or_obj).module() else: loaded_module = deepcopy(exported).module() loaded_module.to(device) - # the model and the returned program share weights, so the model is moved back after the check - model_device = next(itertools.chain(model.parameters(), model.buffers()), torch.empty(0)).device model.to(device) try: with torch.no_grad(): set_determinism(seed=0) - torch_out = model(*verify_inputs) + torch_out = model(*verify_inputs, **verify_kwargs) set_determinism(seed=0) - export_out = loaded_module(*verify_inputs) + export_out = loaded_module(*verify_inputs, **verify_kwargs) set_determinism(seed=None) finally: model.to(model_device) diff --git a/tests/networks/test_convert_to_export.py b/tests/networks/test_convert_to_export.py index cb908993a6..05965eeacc 100644 --- a/tests/networks/test_convert_to_export.py +++ b/tests/networks/test_convert_to_export.py @@ -50,6 +50,31 @@ def forward(self, x): return out +class _AddsOnInputDevice(torch.nn.Module): + """Creates a tensor on the input's device, which the exported graph keeps as the export device.""" + + def forward(self, x): + return x + torch.ones(x.shape, device=x.device) + + +class _TakesDict(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(3, 3) + + def forward(self, inputs): + return self.linear(inputs["a"]) + + +class _TakesKeyword(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(3, 3) + + def forward(self, x, y): + return self.linear(x) + y + + def _build_unet(): return UNet(spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2), num_res_units=0) @@ -114,6 +139,21 @@ def test_verify_leaves_devices_alone(self): self.assertEqual(next(model.parameters()).device.type, "cpu") self.assertEqual(next(iter(exported.state_dict.values())).device.type, "cpu") + def test_verify_default_device_is_model_device(self): + """Without ``device``, verification runs where the model is, so a device kept in the graph still matches.""" + convert_to_export(_AddsOnInputDevice(), verify=True, inputs=[torch.randn(2, 3)]) + + @skip_if_no_cuda + def test_verify_moves_nested_inputs(self): + """Tensors inside dict inputs are moved to the verification device.""" + convert_to_export(_TakesDict(), verify=True, inputs=[{"a": torch.randn(2, 3)}], device="cuda:0") + + def test_verify_with_keyword_inputs(self): + """Keyword inputs passed to export through ``kwargs`` are also used for verification.""" + convert_to_export( + _TakesKeyword(), verify=True, inputs=[torch.randn(2, 3)], kwargs={"y": torch.randn(2, 3)}, device="cpu" + ) + def test_extra_files(self): """Extra files are stored in the saved archive.""" with tempfile.TemporaryDirectory() as tempdir: From 505c6c73c648668124fa7d327e667a94c57c2083 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sun, 8 Nov 2026 12:15:44 +0000 Subject: [PATCH 16/18] Format only the timestamp in save_exported_program With append_timestamp=True the whole path went through strftime, so a "%" in a directory or file name was substituted too. Raised by CodeRabbit: https://github.com/Project-MONAI/MONAI/pull/9264#discussion_r4506119968 Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/data/export_utils.py | 2 +- tests/data/test_export_utils.py | 11 +++++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/monai/data/export_utils.py b/monai/data/export_utils.py index 8e4162f6b5..1af94f7ea2 100644 --- a/monai/data/export_utils.py +++ b/monai/data/export_utils.py @@ -92,7 +92,7 @@ def save_exported_program( ext = ".pt2" if append_timestamp: - filename_prefix_or_stream = now.strftime(f"{filename_no_ext}_%Y%m%d%H%M%S{ext}") + filename_prefix_or_stream = f"{filename_no_ext}_{now.strftime('%Y%m%d%H%M%S')}{ext}" else: filename_prefix_or_stream = filename_no_ext + ext diff --git a/tests/data/test_export_utils.py b/tests/data/test_export_utils.py index 1edb1c1722..1773699f83 100644 --- a/tests/data/test_export_utils.py +++ b/tests/data/test_export_utils.py @@ -45,6 +45,17 @@ def test_save_exported_program_ext(self): save_exported_program(ep, f"{tempdir}/test.zip") self.assertTrue(os.path.isfile(f"{tempdir}/test.zip")) + def test_save_with_timestamp_keeps_path(self): + """Only the timestamp is formatted, so ``%`` in the path is kept as written.""" + ep = torch.export.export(TestModule(), args=(torch.tensor(1.0),)) + with tempfile.TemporaryDirectory() as tempdir: + folder = os.path.join(tempdir, "run_%Y") + os.makedirs(folder) + save_exported_program(ep, os.path.join(folder, "model"), append_timestamp=True) + saved = os.listdir(folder) + self.assertEqual(len(saved), 1) + self.assertRegex(saved[0], r"^model_\d{14}\.pt2$") + 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),)) From d9a5c0a3023be2a93bfef67ca61e28505c8aa88c Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sun, 8 Nov 2026 12:21:50 +0000 Subject: [PATCH 17/18] Load checkpoints onto CPU before bundle export export_checkpoint exports on CPU so the file loads anywhere, but a checkpoint saved from a GPU failed to load on a machine without one. Load it with map_location="cpu" and pass the loaded dict to Ignite, so both loading paths behave the same. ckpt_export shares this code. Raised by CodeRabbit: https://github.com/Project-MONAI/MONAI/pull/9264#discussion_r4506119964 Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/bundle/scripts.py | 5 +++-- tests/bundle/test_bundle_export_checkpoint.py | 22 +++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/monai/bundle/scripts.py b/monai/bundle/scripts.py index bd870ae2b8..1b3ad78c4a 100644 --- a/monai/bundle/scripts.py +++ b/monai/bundle/scripts.py @@ -1326,11 +1326,12 @@ def _export( """ net = parser.get_parsed_content(net_id) + # load onto CPU so that a checkpoint saved on a GPU can be converted on a machine without one + ckpt = torch.load(ckpt_file, map_location="cpu", weights_only=True) if has_ignite: # here we use ignite Checkpoint to support nested weights and be compatible with MONAI CheckpointSaver - Checkpoint.load_objects(to_load={key_in_ckpt: net}, checkpoint=ckpt_file) + Checkpoint.load_objects(to_load={key_in_ckpt: net}, checkpoint=ckpt) else: - ckpt = torch.load(ckpt_file, weights_only=True) copy_model_state(dst=net, src=ckpt if key_in_ckpt == "" else ckpt[key_in_ckpt]) # Use the given converter to convert a model and save with metadata, config content diff --git a/tests/bundle/test_bundle_export_checkpoint.py b/tests/bundle/test_bundle_export_checkpoint.py index 7306de8878..b3af5f73e1 100644 --- a/tests/bundle/test_bundle_export_checkpoint.py +++ b/tests/bundle/test_bundle_export_checkpoint.py @@ -17,6 +17,7 @@ import tempfile import unittest from pathlib import Path +from unittest.mock import patch import torch from parameterized import parameterized @@ -161,6 +162,27 @@ def test_network_on_gpu_exports_on_cpu(self): self.assertEqual(next(module.parameters()).device.type, "cpu") self.assertEqual(module(torch.rand(1, 1, 16, 16)).shape, torch.Size([1, 2, 16, 16])) + @skip_if_no_cuda + def test_gpu_checkpoint_exports_without_gpu(self): + """A checkpoint saved from a GPU can be exported where no GPU is visible.""" + with tempfile.TemporaryDirectory() as tempdir: + config_file = os.path.join(tempdir, "inference.json") + net_def = {"_target_": "UNet", "spatial_dims": 2, "in_channels": 1, "out_channels": 2} + net_def.update({"channels": [4, 8], "strides": [2]}) + ConfigParser.export_config_file(config={"network_def": net_def}, filepath=config_file) + ckpt_file = os.path.join(tempdir, "model.pt") + parser = ConfigParser() + parser.read_config(config_file) + save_state(src=parser.get_parsed_content("network_def").to("cuda:0"), path=ckpt_file) + pt2_file = os.path.join(tempdir, "model.pt2") + + cmd = ["coverage", "run", "-m", "monai.bundle", "export_checkpoint", "network_def"] + cmd += ["--filepath", pt2_file, "--ckpt_file", ckpt_file, "--config_file", config_file] + cmd += ["--input_shape", "[1, 1, 16, 16]"] + with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}): + command_line_tests(cmd) + self.assertTrue(os.path.exists(pt2_file)) + def test_meta_files_are_merged(self): config_file = os.path.join(TESTS_PATH, "testing_data", "inference.json") with tempfile.TemporaryDirectory() as tempdir: From b5ae6ecfdb28a05f4345e295a842ef035dec2459 Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Sun, 8 Nov 2026 13:03:42 +0000 Subject: [PATCH 18/18] Keep input container types when verifying an export copy_to_device turns an OrderedDict into a dict, which the exported program rejects because it checks inputs against the containers it traced. This broke the FeaturePyramidNetwork export test. Move tensors with torch.utils._pytree.tree_map_only, the same flattening torch.export uses, so container types are kept. Assisted-by: Claude Code Signed-off-by: Soumya Snigdha Kundu --- monai/networks/utils.py | 12 +++++++++--- tests/networks/test_convert_to_export.py | 5 +++++ 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/monai/networks/utils.py b/monai/networks/utils.py index 6633d8a59c..9542ffbc17 100644 --- a/monai/networks/utils.py +++ b/monai/networks/utils.py @@ -28,10 +28,11 @@ import numpy as np import torch import torch.nn as nn +from torch.utils._pytree import tree_map_only from monai.apps.utils import get_logger from monai.config import PathLike -from monai.utils.misc import copy_to_device, ensure_tuple, save_obj, set_determinism +from monai.utils.misc import ensure_tuple, save_obj, set_determinism from monai.utils.module import look_up_option, optional_import from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor @@ -925,8 +926,13 @@ def convert_to_export( if device is None: device = model_device - verify_inputs = copy_to_device(list(inputs), device, non_blocking=False) - verify_kwargs = copy_to_device(dict(kwargs.get("kwargs") or {}), device, non_blocking=False) + # torch.export checks inputs against the containers it traced, so move tensors without changing the + # container types, e.g. an OrderedDict must stay an OrderedDict + def _to_device(x: Any) -> Any: + return tree_map_only(torch.Tensor, lambda t: t.to(device), x) + + verify_inputs = _to_device(list(inputs)) + verify_kwargs = _to_device(dict(kwargs.get("kwargs") or {})) if filename_or_obj is not None: loaded_module = torch.export.load(filename_or_obj).module() else: diff --git a/tests/networks/test_convert_to_export.py b/tests/networks/test_convert_to_export.py index 05965eeacc..c039653d96 100644 --- a/tests/networks/test_convert_to_export.py +++ b/tests/networks/test_convert_to_export.py @@ -14,6 +14,7 @@ import os import tempfile import unittest +from collections import OrderedDict from unittest.mock import patch import torch @@ -148,6 +149,10 @@ def test_verify_moves_nested_inputs(self): """Tensors inside dict inputs are moved to the verification device.""" convert_to_export(_TakesDict(), verify=True, inputs=[{"a": torch.randn(2, 3)}], device="cuda:0") + def test_verify_keeps_input_containers(self): + """An OrderedDict input stays an OrderedDict, which the exported program checks against.""" + convert_to_export(_TakesDict(), verify=True, inputs=[OrderedDict(a=torch.randn(2, 3))], device="cpu") + def test_verify_with_keyword_inputs(self): """Keyword inputs passed to export through ``kwargs`` are also used for verification.""" convert_to_export(