diff --git a/docs/source/bundle.rst b/docs/source/bundle.rst index fdf745e9513..5fa8d709979 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 f53588f4236..f4b8a72a652 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 diff --git a/docs/source/mb_specification.rst b/docs/source/mb_specification.rst index 56d660e35c2..9147ca8414a 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 ================== diff --git a/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py b/monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py index 91a9f3d8d3e..d2b48765a2e 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/monai/bundle/__init__.py b/monai/bundle/__init__.py index 3f3c8d545e4..194d3deacd2 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 778c9ef2f0c..95932969785 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 b2809191281..1b3ad78c4a1 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_exported_program, 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, @@ -647,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 @@ -695,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`, @@ -711,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 @@ -720,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: @@ -738,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) @@ -1295,7 +1314,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 @@ -1307,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 @@ -1453,6 +1473,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 +1601,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: @@ -1587,6 +1629,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/monai/data/__init__.py b/monai/data/__init__.py index 562289edcbc..6e38285520f 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 00000000000..1af94f7ea22 --- /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 = f"{filename_no_ext}_{now.strftime('%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/monai/losses/dice.py b/monai/losses/dice.py index 5561012f0ee..858d706caee 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 @@ -197,7 +200,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 @@ -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) @@ -446,7 +447,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/focal_loss.py b/monai/losses/focal_loss.py index 1e7063e06f0..05fa15dd080 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/monai/losses/mcc_loss.py b/monai/losses/mcc_loss.py index 17323f89417..0b57096feac 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 a76c54aa355..f54c86149d3 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/monai/networks/__init__.py b/monai/networks/__init__.py index 5a240021d6c..4271bbd8214 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 83815fc5436..9542ffbc17b 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 @@ -27,6 +28,7 @@ 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 @@ -54,6 +56,7 @@ "save_state", "convert_to_onnx", "convert_to_torchscript", + "convert_to_export", "convert_to_trt", "meshgrid_ij", "meshgrid_xy", @@ -865,6 +868,93 @@ 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, 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 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`. Its ``kwargs`` argument, the keyword + inputs of ``model()``, is also used for verification. + + 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: + # 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 = model_device + + # 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: + loaded_module = deepcopy(exported).module() + loaded_module.to(device) + + model.to(device) + try: + with torch.no_grad(): + set_determinism(seed=0) + torch_out = model(*verify_inputs, **verify_kwargs) + set_determinism(seed=0) + export_out = loaded_module(*verify_inputs, **verify_kwargs) + 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/apps/detection/networks/test_retinanet.py b/tests/apps/detection/networks/test_retinanet.py index eb491361ea6..2a5e28becc7 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 6ac1efd734c..9836de84845 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 4ae198e4091..86fe8ec58d6 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() 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 00000000000..8e792c445d1 --- /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/bundle/test_bundle_export_checkpoint.py b/tests/bundle/test_bundle_export_checkpoint.py new file mode 100644 index 00000000000..b3af5f73e1d --- /dev/null +++ b/tests/bundle/test_bundle_export_checkpoint.py @@ -0,0 +1,243 @@ +# 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 inspect +import json +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import torch +from parameterized import parameterized + +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_no_cuda, 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)) + + _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") + 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_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])) + + @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: + 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}) + + +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() diff --git a/tests/data/test_export_utils.py b/tests/data/test_export_utils.py new file mode 100644 index 00000000000..1773699f831 --- /dev/null +++ b/tests/data/test_export_utils.py @@ -0,0 +1,144 @@ +# 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_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),)) + 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() diff --git a/tests/losses/test_calibration_loss.py b/tests/losses/test_calibration_loss.py index 1504a0029ec..528b469a069 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 97c7ae5050d..761ab0db78f 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 98ea475dedd..d78ef9ecae6 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 d8fb5e11957..28222d99eab 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) @@ -212,11 +212,31 @@ 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) test_script_save(loss, test_input, test_input) + def test_export(self): + loss = DiceLoss() + 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_ds_loss.py b/tests/losses/test_ds_loss.py index 586826aafe5..376eeb1669d 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_focal_loss.py b/tests/losses/test_focal_loss.py index 35017ec8982..21580bbb07f 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) diff --git a/tests/losses/test_generalized_dice_focal_loss.py b/tests/losses/test_generalized_dice_focal_loss.py index 2af4aa68db1..df0cb13ad50 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 7d60b99932e..223a4209fea 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 1ed05c137de..73bc2ffd5c8 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 2451cf76d48..38a5ae3ddc7 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 87ccca7676f..0d0260e69c1 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 8a4988a30d7..035ac996251 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 32303434ca2..a97c000b656 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() diff --git a/tests/min_tests.py b/tests/min_tests.py index 25a42fe4b5d..b153b6aa555 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", diff --git a/tests/networks/blocks/test_dynunet_block.py b/tests/networks/blocks/test_dynunet_block.py index 4cf68912be6..c63401ea04f 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 4b45850159f..b38036014c9 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 d799cd095cc..715722edb62 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 b40f3a0955a..7978d768a98 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 af529186128..b268c7b14f8 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 f4075f5099d..3d375d7189d 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 0073efc609d..50fe805480f 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() diff --git a/tests/networks/nets/dints/test_dints_mixop.py b/tests/networks/nets/dints/test_dints_mixop.py index ea78514fa58..74a66871468 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 d255934495c..d5a0936e6ed 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 152a05fc623..d0e1397f38c 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 7facf9af242..ef28ab8e53f 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 dcf90b809ab..879d1f02a82 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 edd7b599d4c..eee8c23b277 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 37a2813f2ec..1c5699a626a 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 fe7b9bc7408..977f29a4c9f 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 bc6d6d00948..dbc49ef2dff 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 fad32c03506..cfff81807de 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 8f460a2450f..b5457c5a167 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 a573956296e..ebaa5124d77 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 e76d5a6d5ab..dd08160b4f1 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 9ec08194e99..50ad0513c08 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 84e3d466dec..5d504f5f554 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 1384dfaefff..3ca769b8f66 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 58657e6deaa..7064dc516c7 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 15fda15a118..3d7e4014a59 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 08344900e4f..0c1d275f1a6 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 184f76fa9df..60985a580ab 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 913f30b6616..c195a4c846f 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 241f57c78dc..52e2d19f0b5 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 1536d338537..c587f1e0879 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 4a2cf40e6f3..48d8365c9e1 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 90d711d0d98..11c04ca4110 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 7a6d0e98bb5..6402534b39f 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 5d4faa39794..bf11dc3d5aa 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 a8dc715301d..377819149ba 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 949b0f4409d..430b8411ebc 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 6c938934809..a6b3fb087d2 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 1a04bab5689..6322efe7845 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 13a34f96572..cf551fb5c13 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( diff --git a/tests/networks/test_convert_to_export.py b/tests/networks/test_convert_to_export.py new file mode 100644 index 00000000000..c039653d96e --- /dev/null +++ b/tests/networks/test_convert_to_export.py @@ -0,0 +1,202 @@ +# 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 collections import OrderedDict +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 + + +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) + + +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_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_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( + _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: + 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() diff --git a/tests/networks/test_varnet.py b/tests/networks/test_varnet.py index 9cadf15ce55..4d1f56d0542 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() diff --git a/tests/test_utils.py b/tests/test_utils.py index 9d245f9778a..c7006750114 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.