Remove torch.jit constructs from network architectures

Replace @torch.jit.script, torch.jit.is_scripting() guards,
@torch.jit.unused, and @torch.jit.export annotations with plain
Python across all network modules.  Convert torch.jit-unfriendly
patterns (ModuleList iteration, Python containers) back to idiomatic
code.

Also adds .pt2 loading support to MMARS, removes outdated TorchScript
notes from MetaTensor, and refactors _compute_acr_mask in the coil
sensitivity model.

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

Author: aymuos15

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,7 @@ 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).
         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 +232,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,39 @@ 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
+        return acr_1d.view(*([1] * (mask.ndim - 2)), s_len, 1)
 
     def forward(self, masked_kspace: Tensor, mask: Tensor) -> Tensor:
         """
@@ -122,13 +129,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/data/meta_tensor.py

@@ -84,13 +84,6 @@ class MetaTensor(MetaObj, torch.Tensor):
             assert torch.all(m2.affine == affine)
 
     Notes:
-        - Requires pytorch 1.9 or newer for full compatibility.
-        - Older versions of pytorch (<=1.8), `torch.jit.trace(net, im)` may
-          not work if `im` is of type `MetaTensor`. This can be resolved with
-          `torch.jit.trace(net, im.as_tensor())`.
-        - For pytorch < 1.8, sharing `MetaTensor` instances across processes may not be supported.
-        - For pytorch < 1.9, next(iter(meta_tensor)) returns a torch.Tensor.
-          see: https://github.com/pytorch/pytorch/issues/54457
         - A warning will be raised if in the constructor `affine` is not `None` and
           `meta` already contains the key `affine`.
         - You can query whether the `MetaTensor` is a batch with the `is_batch` attribute.

monai/networks/blocks/feature_pyramid_network.py

@@ -206,7 +206,7 @@ def __init__(
     def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor:
         """
         This is equivalent to self.inner_blocks[idx](x),
-        but torchscript doesn't support this yet
+        but module indexing with a variable is used for compatibility
         """
         num_blocks = len(self.inner_blocks)
         if idx < 0:
@@ -220,7 +220,7 @@ def get_result_from_inner_blocks(self, x: Tensor, idx: int) -> Tensor:
     def get_result_from_layer_blocks(self, x: Tensor, idx: int) -> Tensor:
         """
         This is equivalent to self.layer_blocks[idx](x),
-        but torchscript doesn't support this yet
+        but module indexing with a variable is used for compatibility
         """
         num_blocks = len(self.layer_blocks)
         if idx < 0:

monai/networks/blocks/selfattention.py

@@ -112,13 +112,13 @@ def __init__(
 
         if use_combined_linear:
             self.qkv = nn.Linear(self.hidden_input_size, self.inner_dim * 3, bias=qkv_bias)
-            self.to_q = self.to_k = self.to_v = nn.Identity()  # add to enable torchscript
+            self.to_q = self.to_k = self.to_v = nn.Identity()  # placeholder for unused code path
             self.input_rearrange = Rearrange("b h (qkv l d) -> qkv b l h d", qkv=3, l=num_heads)
         else:
             self.to_q = nn.Linear(self.hidden_input_size, self.inner_dim, bias=qkv_bias)
             self.to_k = nn.Linear(self.hidden_input_size, self.inner_dim, bias=qkv_bias)
             self.to_v = nn.Linear(self.hidden_input_size, self.inner_dim, bias=qkv_bias)
-            self.qkv = nn.Identity()  # add to enable torchscript
+            self.qkv = nn.Identity()  # placeholder for unused code path
             self.input_rearrange = Rearrange("b h (l d) -> b l h d", l=num_heads)
         self.out_rearrange = Rearrange("b l h d -> b h (l d)")
         self.drop_output = nn.Dropout(dropout_rate)

monai/networks/blocks/squeeze_and_excitation.py

@@ -81,8 +81,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         y = self.fc(y).view([b, c] + [1] * (x.ndim - 2))
         result = x * y
 
-        # Residual connection is moved here instead of providing an override of forward in ResidualSELayer since
-        # Torchscript has an issue with using super().
+        # Residual connection is applied here rather than in a forward override in ResidualSELayer.
         if self.add_residual:
             result += x
 

monai/networks/layers/factories.py

@@ -262,8 +262,7 @@ def instance_nvfuser_factory(dim):
     It only supports 3d tensors as the input. It also requires to use with CUDA and non-Windows OS.
     In this function, if the required library `apex.normalization.InstanceNorm3dNVFuser` does not exist,
     `nn.InstanceNorm3d` will be returned instead.
-    This layer is based on a customized autograd function, which is not supported in TorchScript currently.
-    Please switch to use `nn.InstanceNorm3d` if TorchScript is necessary.
+    This layer is based on a customized autograd function.
 
     Please check the following link for more details about how to install `apex`:
     https://github.com/NVIDIA/apex#installation

monai/networks/layers/simplelayers.py

@@ -163,7 +163,7 @@ def __init__(self, *shape: int) -> None:
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         shape = list(self.shape)
-        shape[0] = x.shape[0]  # done this way for Torchscript
+        shape[0] = x.shape[0]
         return x.reshape(shape)