Torchscript deprecation

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 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 an exported program (``.pt2``, via ``torch.export``) 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 the exported archive.
 
 Directory Structure
 ===================
@@ -23,7 +23,8 @@ A MONAI Bundle is defined primarily as a directory with a set of specifically na
   ┃  ┗━ metadata.json
   ┣━ models
   ┃  ┣━ model.pt
-  ┃  ┣━ *model.ts
+  ┃  ┣━ *model.pt2
+  ┃  ┣━ *model.ts (deprecated)
   ┃  ┗━ *model.onnx
   ┗━ docs
      ┣━ *README.md
@@ -38,7 +39,8 @@ 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`` exported program if the model is compatible with being exported in this format. This is the preferred format for model deployment.
+* **model.ts**: the TorchScript saved model (deprecated since v1.5, will be removed in v1.7; use ``model.pt2`` instead).
 * **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.
@@ -50,9 +52,13 @@ Archive Format
 
 The bundle directory and its contents can be compressed into a zip file to constitute a single file package. When unzipped into a directory this file will reproduce the above directory structure, and should itself also be named after the model it contains. For example, `ModelName.zip` would contain at least `ModelName/configs/metadata.json` and `ModelName/models/model.pt`, thus when unzipped would place files into the directory `ModelName` rather than into the current working directory.
 
-The Torchscript file format is also just a zip file with a specific structure. When creating such an archive with `save_net_with_metadata` a MB-compliant Torchscript file can be created by including the contents of `metadata.json` as the `meta_values` argument of the function, and other files included as `more_extra_files` entries. These will be stored in a `extras` directory in the zip file and can be retrieved with `load_net_with_metadata` or with any other library/tool that can read zip data. In this format the `model.*` files are obviously not needed but `README.md` and `license.txt` as well as any others provided can be added as more extra files.
+The ``.pt2`` file format (produced by ``torch.export``) is also a zip file with a specific structure. When creating such an archive with ``save_exported_program`` a MB-compliant exported program file can be created by including the contents of ``metadata.json`` as the ``meta_values`` argument of the function, and other files included as ``more_extra_files`` entries. These will be stored in the zip file and can be retrieved with ``load_exported_program`` or with any other library/tool that can read zip data. In this format the ``model.*`` files are obviously not needed but ``README.md`` and ``license.txt`` as well as any others provided can be added as more extra files.
 
-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 ``bundle`` submodule of MONAI contains a number of command line programs. To produce an exported bundle use ``export_checkpoint`` 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 archive will store the files as JSON.
+
+.. note::
+
+   The legacy TorchScript (``ckpt_export``, ``save_net_with_metadata``, ``load_net_with_metadata``) workflow is deprecated since v1.5 and will be removed in v1.7. Use ``export_checkpoint``, ``save_exported_program``, and ``load_exported_program`` instead.
 
 metadata.json File
 ==================

docs/source/modules.md

@@ -237,7 +237,7 @@ and [MLflow](https://github.com/Project-MONAI/tutorials/blob/main/experiment_man
 
 The objective of a MONAI bundle is to define a packaged 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
+single network as a pickled state dictionary plus optionally an exported program (`.pt2`, via `torch.export`) 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. More details are available at [bundle specification](https://monai.readthedocs.io/en/latest/mb_specification.html).

monai/apps/detection/networks/retinanet_network.py

@@ -332,7 +332,7 @@ def forward(self, images: Tensor) -> Any:
         features = self.feature_extractor(images)
         if isinstance(features, Tensor):
             feature_maps = [features]
-        elif torch.jit.isinstance(features, dict[str, Tensor]):
+        elif isinstance(features, dict):
             feature_maps = list(features.values())
         else:
             feature_maps = list(features)

monai/apps/detection/utils/anchor_utils.py

@@ -257,7 +257,7 @@ def grid_anchors(self, grid_sizes: list[list[int]], strides: list[list[Tensor]])
                 for axis in range(self.spatial_dims)
             ]
 
-            # to support torchscript, cannot directly use torch.meshgrid(shifts_centers).
+            # unpack before passing to torch.meshgrid for compatibility.
             shifts_centers = list(torch.meshgrid(shifts_centers[: self.spatial_dims], indexing="ij"))
 
             for axis in range(self.spatial_dims):

monai/apps/mmars/mmars.py

@@ -205,7 +205,8 @@ def load_from_mmar(
         mmar_dir: : target directory to store the MMAR, default is mmars subfolder under `torch.hub get_dir()`.
         progress: whether to display a progress bar when downloading the content.
         version: version number of the MMAR. Set it to `-1` to use `item[Keys.VERSION]`.
-        map_location: pytorch API parameter for `torch.load` or `torch.jit.load`.
+        map_location: pytorch API parameter for ``torch.load`` or ``torch.jit.load`` (legacy ``.ts`` files).
+            Ignored when loading ``.pt2`` (ExportedProgram) files.
         pretrained: whether to load the pretrained weights after initializing a network module.
         weights_only: whether to load only the weights instead of initializing the network module and assign weights.
         model_key: a key to search in the model file or config file for the model dictionary.
@@ -232,12 +233,26 @@ def load_from_mmar(
     _model_file = model_dir / item.get(Keys.MODEL_FILE, model_file)
     logger.info(f'\n*** "{item.get(Keys.NAME)}" available at {model_dir}.')
 
-    # loading with `torch.jit.load`
+    # loading with `torch.export.load` for .pt2 files
+    if _model_file.name.endswith(".pt2"):
+        if not pretrained:
+            warnings.warn("Loading an ExportedProgram, 'pretrained' option ignored.", stacklevel=2)
+        if weights_only:
+            warnings.warn("Loading an ExportedProgram, 'weights_only' option ignored.", stacklevel=2)
+        return torch.export.load(str(_model_file))
+
+    # loading with `torch.jit.load` for legacy .ts files
     if _model_file.name.endswith(".ts"):
+        warnings.warn(
+            "Loading TorchScript (.ts) models is deprecated since MONAI v1.5 and will be removed in v1.7. "
+            "Use torch.export (.pt2) format instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
         if not pretrained:
-            warnings.warn("Loading a ScriptModule, 'pretrained' option ignored.")
+            warnings.warn("Loading a ScriptModule, 'pretrained' option ignored.", stacklevel=2)
         if weights_only:
-            warnings.warn("Loading a ScriptModule, 'weights_only' option ignored.")
+            warnings.warn("Loading a ScriptModule, 'weights_only' option ignored.", stacklevel=2)
         return torch.jit.load(_model_file, map_location=map_location)
 
     # loading with `torch.load`

monai/apps/reconstruction/networks/nets/coil_sensitivity_model.py

@@ -85,32 +85,40 @@ def __init__(
         self.spatial_dims = spatial_dims
         self.coil_dim = coil_dim
 
-    def get_fully_sampled_region(self, mask: Tensor) -> tuple[int, int]:
+    def _compute_acr_mask(self, mask: Tensor) -> Tensor:
         """
-        Extracts the size of the fully-sampled part of the kspace. Note that when a kspace
-        is under-sampled, a part of its center is fully sampled. This part is called the Auto
-        Calibration Region (ACR). ACR is used for sensitivity map computation.
+        Compute a boolean mask for the Auto Calibration Region (ACR) — the contiguous
+        fully-sampled center of the k-space sampling mask.
+
+        Uses pure tensor operations (``cumprod``) instead of while-loops so that
+        the computation is compatible with ``torch.export``.
 
         Args:
-            mask: the under-sampling mask of shape (..., S, 1) where S denotes the sampling dimension
+            mask: the under-sampling mask of shape (..., S, 1) where S denotes the sampling dimension.
 
         Returns:
-            A tuple containing
-                (1) left index of the region
-                (2) right index of the region
-
-        Note:
-            Suppose the mask is of shape (1,1,20,1). If this function returns 8,12 as left and right
-                indices, then it means that the fully-sampled center region has size 4 starting from 8 to 12.
+            A boolean tensor broadcastable to ``masked_kspace`` that is True inside the ACR.
         """
-        left = right = mask.shape[-2] // 2
-        while mask[..., right, :]:
-            right += 1
+        s_len = mask.shape[-2]
+        center = s_len // 2
+
+        # Flatten to 1-D along the sampling axis
+        m = mask.reshape(-1)[:s_len].bool()
+
+        # Count consecutive True values from center going right
+        right_count = torch.cumprod(m[center:].int(), dim=0).sum()
+        # Count consecutive True values from center going left (including center)
+        left_count = torch.cumprod(m[: center + 1].flip(0).int(), dim=0).sum()
+        num_low_freqs = left_count + right_count - 1
 
-        while mask[..., left, :]:
-            left -= 1
+        # Build a boolean mask over the sampling dimension
+        start = (s_len - num_low_freqs + 1) // 2
+        freq_idx = torch.arange(s_len, device=mask.device)
+        acr_1d = (freq_idx >= start) & (freq_idx < start + num_low_freqs)
 
-        return left + 1, right
+        # Reshape to (..., S, 1) so it broadcasts against masked_kspace
+        result: Tensor = acr_1d.view(*([1] * (mask.ndim - 2)), s_len, 1)
+        return result
 
     def forward(self, masked_kspace: Tensor, mask: Tensor) -> Tensor:
         """
@@ -122,13 +130,10 @@ 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
+        acr_mask = self._compute_acr_mask(mask)
 
         # 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 * acr_mask
 
         # apply inverse fourier to the extracted fully-sampled data
         x = ifftn_centered_t(x, spatial_dims=self.spatial_dims, is_complex=True)

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,

monai/bundle/__main__.py

@@ -15,6 +15,7 @@
     ckpt_export,
     download,
     download_large_files,
+    export_checkpoint,
     init_bundle,
     onnx_export,
     run,