Add torch.export save/load infrastructure and deprecate TorchScript utils

Introduce save_exported_program() and load_exported_program() in a new
monai/data/export_utils.py module for the torch.export-based .pt2
serialisation format.  Mark the legacy save_net_with_metadata() and
load_net_with_metadata() as deprecated.

Also adds ExportMetadataKeys enum and inlines warn_deprecated() into
deprecate_utils to remove a circular-import hazard.

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

Author: aymuos15

monai/data/__init__.py

@@ -77,6 +77,7 @@
 from .test_time_augmentation import TestTimeAugmentation
 from .thread_buffer import ThreadBuffer, ThreadDataLoader
 from .torchscript_utils import load_net_with_metadata, save_net_with_metadata
+from .export_utils import load_exported_program, save_exported_program
 from .utils import (
     affine_to_spacing,
     compute_importance_map,

monai/data/export_utils.py

@@ -0,0 +1,129 @@
+# 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 logging
+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 ExportMetadataKeys
+
+__all__ = ["save_exported_program", "load_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.
+    """
+    now = datetime.datetime.now()
+    metadict: dict[str, Any] = {}
+
+    if include_config_vals:
+        metadict.update(get_config_values())
+        metadict[ExportMetadataKeys.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:
+        extra_files.update(more_extra_files)
+
+    # torch.export.save requires str values; decode bytes from legacy callers (e.g. _export helper)
+    extra_files = {k: v.decode() if isinstance(v, bytes) else v for k, v in extra_files.items()}
+
+    if isinstance(filename_prefix_or_stream, (str, os.PathLike)):
+        filename_prefix_or_stream = str(filename_prefix_or_stream)
+        filename_no_ext, ext = os.path.splitext(filename_prefix_or_stream)
+        if ext == "":
+            ext = ".pt2"
+
+        if append_timestamp:
+            filename_prefix_or_stream = now.strftime(f"{filename_no_ext}_%Y%m%d%H%M%S{ext}")
+        else:
+            filename_prefix_or_stream = filename_no_ext + ext
+
+    torch.export.save(exported_program, filename_prefix_or_stream, extra_files=extra_files)
+
+
+def load_exported_program(
+    filename_prefix_or_stream: str | os.PathLike | IO[bytes],
+    more_extra_files: Sequence[str] = (),
+) -> tuple[torch.export.ExportedProgram, dict, dict]:
+    """
+    Load an ``ExportedProgram`` from a ``.pt2`` file and extract stored JSON metadata.
+
+    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).
+    """
+    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)
+
+    extra_files = dict(extra_files)
+
+    json_data = extra_files.pop(METADATA_FILENAME, "{}")
+
+    try:
+        json_data_dict = json.loads(json_data)
+    except json.JSONDecodeError:
+        logging.getLogger(__name__).warning(
+            "Failed to parse metadata JSON from exported program, returning empty metadata."
+        )
+        json_data_dict = {}
+
+    return exported_program, json_data_dict, extra_files

monai/data/torchscript_utils.py

@@ -21,10 +21,12 @@
 
 from monai.config import get_config_values
 from monai.utils import JITMetadataKeys
+from monai.utils.deprecate_utils import deprecated
 
 METADATA_FILENAME = "metadata.json"
 
 
+@deprecated(since="1.5", removed="1.7", msg_suffix="Use monai.data.save_exported_program() instead.")
 def save_net_with_metadata(
     jit_obj: torch.nn.Module,
     filename_prefix_or_stream: str | IO[Any],
@@ -100,6 +102,7 @@ def save_net_with_metadata(
     torch.jit.save(jit_obj, filename_prefix_or_stream, extra_files)
 
 
+@deprecated(since="1.5", removed="1.7", msg_suffix="Use monai.data.load_exported_program() instead.")
 def load_net_with_metadata(
     filename_prefix_or_stream: str | IO[Any],
     map_location: torch.device | None = None,

monai/utils/__init__.py

@@ -31,6 +31,7 @@
     DiceCEReduction,
     DownsampleMode,
     EngineStatsKeys,
+    ExportMetadataKeys,
     FastMRIKeys,
     ForwardMode,
     GanKeys,

monai/utils/deprecate_utils.py

@@ -30,12 +30,6 @@ class DeprecatedError(Exception):
     pass
 
 
-def warn_deprecated(obj, msg, warning_category=FutureWarning):
-    """
-    Issue the warning message `msg`.
-    """
-    warnings.warn(f"{obj}: {msg}", category=warning_category, stacklevel=2)
-
 
 def deprecated(
     since: str | None = None,
@@ -107,7 +101,7 @@ def _wrapper(*args, **kwargs):
             if is_removed:
                 raise DeprecatedError(msg)
             if is_deprecated:
-                warn_deprecated(obj, msg, warning_category)
+                warnings.warn(f"{obj}: {msg}", category=warning_category, stacklevel=2)
 
             return call_obj(*args, **kwargs)
 
@@ -217,7 +211,7 @@ def _wrapper(*args, **kwargs):
                 if is_removed:
                     raise DeprecatedError(msg)
                 if is_deprecated:
-                    warn_deprecated(argname, msg, warning_category)
+                    warnings.warn(f"{argname}: {msg}", category=warning_category, stacklevel=2)
 
             return func(*args, **kwargs)
 
@@ -317,7 +311,7 @@ def _decorator(func):
         def _wrapper(*args, **kwargs):
             if name not in sig.bind(*args, **kwargs).arguments and is_deprecated:
                 # arg was not found so the default value is used
-                warn_deprecated(argname, msg, warning_category)
+                warnings.warn(f"{argname}: {msg}", category=warning_category, stacklevel=2)
 
             return func(*args, **kwargs)
 

monai/utils/enums.py

@@ -62,6 +62,8 @@
     "BundlePropertyConfig",
     "AlgoKeys",
     "IgniteInfo",
+    "JITMetadataKeys",
+    "ExportMetadataKeys",
 ]
 
 
@@ -423,6 +425,9 @@ class JITMetadataKeys(StrEnum):
     """
     Keys stored in the metadata file for saved Torchscript models. Some of these are generated by the routines
     and others are optionally provided by users.
+
+    .. deprecated:: 1.5
+        Use :class:`ExportMetadataKeys` instead.
     """
 
     NAME = "name"
@@ -431,6 +436,10 @@ class JITMetadataKeys(StrEnum):
     DESCRIPTION = "description"
 
 
+# ExportMetadataKeys shares the same members as JITMetadataKeys; alias to avoid duplication.
+ExportMetadataKeys = JITMetadataKeys
+
+
 class BoxModeName(StrEnum):
     """
     Box mode names.