monai-weekly 1.5.dev2511__py3-none-any.whl → 1.5.dev2513__py3-none-any.whl

monai/networks/blocks/downsample.py

@@ -16,8 +16,11 @@ from collections.abc import Sequence
 import torch
 import torch.nn as nn
 
-from monai.networks.layers.factories import Pool
-from monai.utils import ensure_tuple_rep
+from monai.networks.layers.factories import Conv, Pool
+from monai.networks.utils import pixelunshuffle
+from monai.utils import DownsampleMode, ensure_tuple_rep, look_up_option
+
+__all__ = ["MaxAvgPool", "DownSample", "Downsample", "SubpixelDownsample", "SubpixelDownSample", "Subpixeldownsample"]
 
 
 class MaxAvgPool(nn.Module):
@@ -61,3 +64,239 @@ class MaxAvgPool(nn.Module):
             Tensor in shape (batch, 2*channel, spatial_1[, spatial_2, ...]).
         """
         return torch.cat([self.max_pool(x), self.avg_pool(x)], dim=1)
+
+
+class DownSample(nn.Sequential):
+    """
+    Downsamples data by `scale_factor`.
+
+    Supported modes are:
+
+    - "conv": uses a strided convolution for learnable downsampling.
+    - "convgroup": uses a grouped strided convolution for efficient feature reduction.
+    - "nontrainable": uses :py:class:`torch.nn.Upsample` with inverse scale factor.
+    - "pixelunshuffle": uses :py:class:`monai.networks.blocks.PixelUnshuffle` for channel-space rearrangement.
+
+    This operation will cause non-deterministic behavior when ``mode`` is ``DownsampleMode.NONTRAINABLE``.
+    Please check the link below for more details:
+    https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms
+
+    This module can optionally take a pre-convolution
+    (often used to map the number of features from `in_channels` to `out_channels`).
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int | None = None,
+        out_channels: int | None = None,
+        scale_factor: Sequence[float] | float = 2,
+        kernel_size: Sequence[float] | float | None = None,
+        mode: DownsampleMode | str = DownsampleMode.CONV,
+        pre_conv: nn.Module | str | None = "default",
+        post_conv: nn.Module | None = None,
+        bias: bool = True,
+    ) -> None:
+        """
+        Downsamples data by `scale_factor`.
+        Supported modes are:
+
+            - DownsampleMode.CONV: uses a strided convolution for learnable downsampling.
+            - DownsampleMode.CONVGROUP: uses a grouped strided convolution for efficient feature reduction.
+            - DownsampleMode.MAXPOOL: uses maxpooling for non-learnable downsampling.
+            - DownsampleMode.AVGPOOL: uses average pooling for non-learnable downsampling.
+            - DownsampleMode.PIXELUNSHUFFLE: uses :py:class:`monai.networks.blocks.SubpixelDownsample`.
+
+        This operation will cause non-deterministic behavior when ``mode`` is ``DownsampleMode.NONTRAINABLE``.
+        Please check the link below for more details:
+        https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms
+
+        This module can optionally take a pre-convolution and post-convolution
+        (often used to map the number of features from `in_channels` to `out_channels`).
+
+        Args:
+            spatial_dims: number of spatial dimensions of the input image.
+            in_channels: number of channels of the input image.
+            out_channels: number of channels of the output image. Defaults to `in_channels`.
+            scale_factor: multiplier for spatial size reduction. Has to match input size if it is a tuple. Defaults to 2.
+            kernel_size: kernel size used during convolutions. Defaults to `scale_factor`.
+            mode: {``DownsampleMode.CONV``, ``DownsampleMode.CONVGROUP``, ``DownsampleMode.MAXPOOL``, ``DownsampleMode.AVGPOOL``,
+                ``DownsampleMode.PIXELUNSHUFFLE``}. Defaults to ``DownsampleMode.CONV``.
+            pre_conv: a conv block applied before downsampling. Defaults to "default".
+                When ``conv_block`` is ``"default"``, one reserved conv layer will be utilized.
+                Only used in the "maxpool", "avgpool" or "pixelunshuffle" modes.
+            post_conv: a conv block applied after downsampling. Defaults to None. Only used in the "maxpool" and "avgpool" modes.
+            bias: whether to have a bias term in the default preconv and conv layers. Defaults to True.
+        """
+        super().__init__()
+
+        scale_factor_ = ensure_tuple_rep(scale_factor, spatial_dims)
+        down_mode = look_up_option(mode, DownsampleMode)
+
+        if not kernel_size:
+            kernel_size_ = scale_factor_
+            padding = ensure_tuple_rep(0, spatial_dims)
+        else:
+            kernel_size_ = ensure_tuple_rep(kernel_size, spatial_dims)
+            padding = tuple((k - 1) // 2 for k in kernel_size_)
+
+        if down_mode == DownsampleMode.CONV:
+            if not in_channels:
+                raise ValueError("in_channels needs to be specified in conv mode")
+            self.add_module(
+                "conv",
+                Conv[Conv.CONV, spatial_dims](
+                    in_channels=in_channels,
+                    out_channels=out_channels or in_channels,
+                    kernel_size=kernel_size_,
+                    stride=scale_factor_,
+                    padding=padding,
+                    bias=bias,
+                ),
+            )
+        elif down_mode == DownsampleMode.CONVGROUP:
+            if not in_channels:
+                raise ValueError("in_channels needs to be specified")
+            if out_channels is None:
+                out_channels = in_channels
+            groups = in_channels if out_channels % in_channels == 0 else 1
+            self.add_module(
+                "convgroup",
+                Conv[Conv.CONV, spatial_dims](
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    kernel_size=kernel_size_,
+                    stride=scale_factor_,
+                    padding=padding,
+                    groups=groups,
+                    bias=bias,
+                ),
+            )
+        elif down_mode == DownsampleMode.MAXPOOL:
+            if pre_conv == "default" and (out_channels != in_channels):
+                if not in_channels:
+                    raise ValueError("in_channels needs to be specified")
+                self.add_module(
+                    "preconv",
+                    Conv[Conv.CONV, spatial_dims](
+                        in_channels=in_channels, out_channels=out_channels or in_channels, kernel_size=1, bias=bias
+                    ),
+                )
+            self.add_module(
+                "maxpool", Pool[Pool.MAX, spatial_dims](kernel_size=kernel_size_, stride=scale_factor_, padding=padding)
+            )
+            if post_conv:
+                self.add_module("postconv", post_conv)
+
+        elif down_mode == DownsampleMode.AVGPOOL:
+            if pre_conv == "default" and (out_channels != in_channels):
+                if not in_channels:
+                    raise ValueError("in_channels needs to be specified")
+                self.add_module(
+                    "preconv",
+                    Conv[Conv.CONV, spatial_dims](
+                        in_channels=in_channels, out_channels=out_channels or in_channels, kernel_size=1, bias=bias
+                    ),
+                )
+            self.add_module(
+                "avgpool", Pool[Pool.AVG, spatial_dims](kernel_size=kernel_size_, stride=scale_factor_, padding=padding)
+            )
+            if post_conv:
+                self.add_module("postconv", post_conv)
+
+        elif down_mode == DownsampleMode.PIXELUNSHUFFLE:
+            self.add_module(
+                "pixelunshuffle",
+                SubpixelDownsample(
+                    spatial_dims=spatial_dims,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    scale_factor=scale_factor_[0],
+                    conv_block=pre_conv,
+                    bias=bias,
+                ),
+            )
+
+
+class SubpixelDownsample(nn.Module):
+    """
+    Downsample via using a subpixel CNN. This module supports 1D, 2D and 3D input images.
+    The module consists of two parts. First, a convolutional layer is employed
+    to adjust the number of channels. Secondly, a pixel unshuffle manipulation
+    rearranges the spatial information into channel space, effectively reducing
+    spatial dimensions while increasing channel depth.
+
+    The pixel unshuffle operation is the inverse of pixel shuffle, rearranging dimensions
+    from (B, C, H*r, W*r) to (B, C*r², H, W) for 2D images or from (B, C, H*r, W*r, D*r) to (B, C*r³, H, W, D) in 3D case.
+
+    Example: (1, 1, 4, 4) with r=2 becomes (1, 4, 2, 2).
+
+    See: Shi et al., 2016, "Real-Time Single Image and Video Super-Resolution
+    Using a nEfficient Sub-Pixel Convolutional Neural Network."
+
+    The pixel unshuffle mechanism is the inverse operation of:
+    https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/blocks/upsample.py
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        in_channels: int | None,
+        out_channels: int | None = None,
+        scale_factor: int = 2,
+        conv_block: nn.Module | str | None = "default",
+        bias: bool = True,
+    ) -> None:
+        """
+        Downsamples data by rearranging spatial information into channel space.
+        This reduces spatial dimensions while increasing channel depth.
+
+        Args:
+            spatial_dims: number of spatial dimensions of the input image.
+            in_channels: number of channels of the input image.
+            out_channels: optional number of channels of the output image.
+            scale_factor: factor to reduce the spatial dimensions by. Defaults to 2.
+            conv_block: a conv block to adjust channels before downsampling. Defaults to None.
+                When ``conv_block`` is ``"default"``, one reserved conv layer will be utilized.
+                When ``conv_block`` is an ``nn.module``,
+                please ensure the input number of channels matches requirements.
+            bias: whether to have a bias term in the default conv_block. Defaults to True.
+        """
+        super().__init__()
+
+        if scale_factor <= 0:
+            raise ValueError(f"The `scale_factor` multiplier must be an integer greater than 0, got {scale_factor}.")
+
+        self.dimensions = spatial_dims
+        self.scale_factor = scale_factor
+
+        if conv_block == "default":
+            if not in_channels:
+                raise ValueError("in_channels need to be specified.")
+            out_channels = out_channels or in_channels
+            self.conv_block = Conv[Conv.CONV, self.dimensions](
+                in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=bias
+            )
+        elif conv_block is None:
+            self.conv_block = nn.Identity()
+        else:
+            self.conv_block = conv_block
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor in shape (batch, channel, spatial_1[, spatial_2, ...).
+        Returns:
+            Tensor with reduced spatial dimensions and increased channel depth.
+        """
+        x = self.conv_block(x)
+        if not all(d % self.scale_factor == 0 for d in x.shape[2:]):
+            raise ValueError(
+                f"All spatial dimensions {x.shape[2:]} must be evenly " f"divisible by scale_factor {self.scale_factor}"
+            )
+        x = pixelunshuffle(x, self.dimensions, self.scale_factor)
+        return x
+
+
+Downsample = DownSample
+SubpixelDownSample = Subpixeldownsample = SubpixelDownsample

monai/networks/nets/restormer.py

@@ -0,0 +1,337 @@
+# 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 torch
+import torch.nn as nn
+
+from monai.networks.blocks.cablock import CABlock, FeedForward
+from monai.networks.blocks.convolutions import Convolution
+from monai.networks.blocks.downsample import DownSample
+from monai.networks.blocks.upsample import UpSample
+from monai.networks.layers.factories import Norm
+from monai.utils.enums import DownsampleMode, UpsampleMode
+
+
+class MDTATransformerBlock(nn.Module):
+    """Basic transformer unit combining MDTA and GDFN with skip connections.
+    Unlike standard transformers that use LayerNorm, this block uses Instance Norm
+    for better adaptation to image restoration tasks.
+
+    Args:
+        spatial_dims: Number of spatial dimensions (2D or 3D)
+        dim: Number of input channels
+        num_heads: Number of attention heads
+        ffn_expansion_factor: Expansion factor for feed-forward network
+        bias: Whether to use bias in attention layers
+        layer_norm_use_bias: Whether to use bias in layer normalization. Defaults to False.
+        flash_attention: Whether to use flash attention optimization. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int,
+        dim: int,
+        num_heads: int,
+        ffn_expansion_factor: float,
+        bias: bool,
+        layer_norm_use_bias: bool = False,
+        flash_attention: bool = False,
+    ):
+        super().__init__()
+        self.norm1 = Norm[Norm.INSTANCE, spatial_dims](dim, affine=layer_norm_use_bias)
+        self.attn = CABlock(spatial_dims, dim, num_heads, bias, flash_attention)
+        self.norm2 = Norm[Norm.INSTANCE, spatial_dims](dim, affine=layer_norm_use_bias)
+        self.ffn = FeedForward(spatial_dims, dim, ffn_expansion_factor, bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class OverlapPatchEmbed(Convolution):
+    """Initial feature extraction using overlapped convolutions.
+    Unlike standard patch embeddings that use non-overlapping patches,
+    this approach maintains spatial continuity through 3x3 convolutions.
+
+    Args:
+        spatial_dims: Number of spatial dimensions (2D or 3D)
+        in_channels: Number of input channels
+        embed_dim: Dimension of embedded features. Defaults to 48.
+        bias: Whether to use bias in convolution layer. Defaults to False.
+    """
+
+    def __init__(self, spatial_dims: int, in_channels: int = 3, embed_dim: int = 48, bias: bool = False):
+        super().__init__(
+            spatial_dims=spatial_dims,
+            in_channels=in_channels,
+            out_channels=embed_dim,
+            kernel_size=3,
+            strides=1,
+            padding=1,
+            bias=bias,
+            conv_only=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = super().forward(x)
+        return x
+
+
+class Restormer(nn.Module):
+    """Restormer: Efficient Transformer for High-Resolution Image Restoration.
+
+    Implements a U-Net style architecture with transformer blocks, combining:
+    - Multi-scale feature processing through progressive down/upsampling
+    - Efficient attention via MDTA blocks
+    - Local feature mixing through GDFN
+    - Skip connections for preserving spatial details
+
+    Architecture:
+        - Encoder: Progressive feature downsampling with increasing channels
+        - Latent: Deep feature processing at lowest resolution
+        - Decoder: Progressive upsampling with skip connections
+        - Refinement: Final feature enhancement
+    """
+
+    def __init__(
+        self,
+        spatial_dims: int = 2,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        dim: int = 48,
+        num_blocks: tuple[int, ...] = (1, 1, 1, 1),
+        heads: tuple[int, ...] = (1, 1, 1, 1),
+        num_refinement_blocks: int = 4,
+        ffn_expansion_factor: float = 2.66,
+        bias: bool = False,
+        layer_norm_use_bias: bool = True,
+        dual_pixel_task: bool = False,
+        flash_attention: bool = False,
+    ) -> None:
+        super().__init__()
+        """Initialize Restormer model.
+
+        Args:
+            spatial_dims: Number of spatial dimensions (2D or 3D)
+            in_channels: Number of input image channels
+            out_channels: Number of output image channels
+            dim: Base feature dimension. Defaults to 48.
+            num_blocks: Number of transformer blocks at each scale. Defaults to (1,1,1,1).
+            heads: Number of attention heads at each scale. Defaults to (1,1,1,1).
+            num_refinement_blocks: Number of final refinement blocks. Defaults to 4.
+            ffn_expansion_factor: Expansion factor for feed-forward network. Defaults to 2.66.
+            bias: Whether to use bias in convolutions. Defaults to False.
+            layer_norm_use_bias: Whether to use bias in layer normalization. Defaults to True.
+            dual_pixel_task: Enable dual-pixel specific processing. Defaults to False.
+            flash_attention: Use flash attention if available. Defaults to False.
+
+        Note:
+            The number of blocks must be greater than 1
+            The length of num_blocks and heads must be equal
+            All values in num_blocks must be greater than 0
+        """
+        # Check input parameters
+        assert len(num_blocks) > 1, "Number of blocks must be greater than 1"
+        assert len(num_blocks) == len(heads), "Number of blocks and heads must be equal"
+        assert all(n > 0 for n in num_blocks), "Number of blocks must be greater than 0"
+
+        # Initial feature extraction
+        self.patch_embed = OverlapPatchEmbed(spatial_dims, in_channels, dim)
+        self.encoder_levels = nn.ModuleList()
+        self.downsamples = nn.ModuleList()
+        self.decoder_levels = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        self.reduce_channels = nn.ModuleList()
+        num_steps = len(num_blocks) - 1
+        self.num_steps = num_steps
+        self.spatial_dims = spatial_dims
+        spatial_multiplier = 2 ** (spatial_dims - 1)
+
+        # Define encoder levels
+        for n in range(num_steps):
+            current_dim = dim * (2) ** (n)
+            next_dim = current_dim // spatial_multiplier
+            self.encoder_levels.append(
+                nn.Sequential(
+                    *[
+                        MDTATransformerBlock(
+                            spatial_dims=spatial_dims,
+                            dim=current_dim,
+                            num_heads=heads[n],
+                            ffn_expansion_factor=ffn_expansion_factor,
+                            bias=bias,
+                            layer_norm_use_bias=layer_norm_use_bias,
+                            flash_attention=flash_attention,
+                        )
+                        for _ in range(num_blocks[n])
+                    ]
+                )
+            )
+
+            self.downsamples.append(
+                DownSample(
+                    spatial_dims=self.spatial_dims,
+                    in_channels=current_dim,
+                    out_channels=next_dim,
+                    mode=DownsampleMode.PIXELUNSHUFFLE,
+                    scale_factor=2,
+                    bias=bias,
+                )
+            )
+
+        # Define latent space
+        latent_dim = dim * (2) ** (num_steps)
+        self.latent = nn.Sequential(
+            *[
+                MDTATransformerBlock(
+                    spatial_dims=spatial_dims,
+                    dim=latent_dim,
+                    num_heads=heads[num_steps],
+                    ffn_expansion_factor=ffn_expansion_factor,
+                    bias=bias,
+                    layer_norm_use_bias=layer_norm_use_bias,
+                    flash_attention=flash_attention,
+                )
+                for _ in range(num_blocks[num_steps])
+            ]
+        )
+
+        # Define decoder levels
+        for n in reversed(range(num_steps)):
+            current_dim = dim * (2) ** (n)
+            next_dim = dim * (2) ** (n + 1)
+            self.upsamples.append(
+                UpSample(
+                    spatial_dims=self.spatial_dims,
+                    in_channels=next_dim,
+                    out_channels=(current_dim),
+                    mode=UpsampleMode.PIXELSHUFFLE,
+                    scale_factor=2,
+                    bias=bias,
+                    apply_pad_pool=False,
+                )
+            )
+
+            # Reduce channel layers to deal with skip connections
+            if n != 0:
+                self.reduce_channels.append(
+                    Convolution(
+                        spatial_dims=self.spatial_dims,
+                        in_channels=next_dim,
+                        out_channels=current_dim,
+                        kernel_size=1,
+                        bias=bias,
+                        conv_only=True,
+                    )
+                )
+                decoder_dim = current_dim
+            else:
+                decoder_dim = next_dim
+
+            self.decoder_levels.append(
+                nn.Sequential(
+                    *[
+                        MDTATransformerBlock(
+                            spatial_dims=spatial_dims,
+                            dim=decoder_dim,
+                            num_heads=heads[n],
+                            ffn_expansion_factor=ffn_expansion_factor,
+                            bias=bias,
+                            layer_norm_use_bias=layer_norm_use_bias,
+                            flash_attention=flash_attention,
+                        )
+                        for _ in range(num_blocks[n])
+                    ]
+                )
+            )
+
+        # Final refinement and output
+        self.refinement = nn.Sequential(
+            *[
+                MDTATransformerBlock(
+                    spatial_dims=spatial_dims,
+                    dim=decoder_dim,
+                    num_heads=heads[0],
+                    ffn_expansion_factor=ffn_expansion_factor,
+                    bias=bias,
+                    layer_norm_use_bias=layer_norm_use_bias,
+                    flash_attention=flash_attention,
+                )
+                for _ in range(num_refinement_blocks)
+            ]
+        )
+        self.dual_pixel_task = dual_pixel_task
+        if self.dual_pixel_task:
+            self.skip_conv = Convolution(
+                spatial_dims=self.spatial_dims,
+                in_channels=dim,
+                out_channels=dim * 2,
+                kernel_size=1,
+                bias=bias,
+                conv_only=True,
+            )
+        self.output = Convolution(
+            spatial_dims=self.spatial_dims,
+            in_channels=dim * 2,
+            out_channels=out_channels,
+            kernel_size=3,
+            strides=1,
+            padding=1,
+            bias=bias,
+            conv_only=True,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of Restormer.
+        Processes input through encoder-decoder architecture with skip connections.
+        Args:
+            inp_img: Input image tensor of shape (B, C, H, W, [D])
+
+        Returns:
+            Restored image tensor of shape (B, C, H, W, [D])
+        """
+        assert all(
+            x.shape[-i] > 2**self.num_steps for i in range(1, self.spatial_dims + 1)
+        ), "All spatial dimensions should be larger than 2^number_of_step"
+
+        # Patch embedding
+        x = self.patch_embed(x)
+        skip_connections = []
+
+        # Encoding path
+        for _idx, (encoder, downsample) in enumerate(zip(self.encoder_levels, self.downsamples)):
+            x = encoder(x)
+            skip_connections.append(x)
+            x = downsample(x)
+
+        # Latent space
+        x = self.latent(x)
+
+        # Decoding path
+        for idx in range(len(self.decoder_levels)):
+            x = self.upsamples[idx](x)
+            x = torch.concat([x, skip_connections[-(idx + 1)]], 1)
+            if idx < len(self.decoder_levels) - 1:
+                x = self.reduce_channels[idx](x)
+            x = self.decoder_levels[idx](x)
+
+        # Final refinement
+        x = self.refinement(x)
+
+        if self.dual_pixel_task:
+            x = x + self.skip_conv(skip_connections[0])
+            x = self.output(x)
+        else:
+            x = self.output(x)
+
+        return x

monai/networks/utils.py

@@ -49,6 +49,7 @@ __all__ = [
     "normal_init",
     "icnr_init",
     "pixelshuffle",
+    "pixelunshuffle",
     "eval_mode",
     "train_mode",
     "get_state_dict",
@@ -376,7 +377,7 @@ def pixelshuffle(x: torch.Tensor, spatial_dims: int, scale_factor: int) -> torch
     See: Aitken et al., 2017, "Checkerboard artifact free sub-pixel convolution".
 
     Args:
-        x: Input tensor
+        x: Input tensor with shape BCHW[D]
         spatial_dims: number of spatial dimensions, typically 2 or 3 for 2D or 3D
         scale_factor: factor to rescale the spatial dimensions by, must be >=1
 
@@ -411,6 +412,48 @@ def pixelshuffle(x: torch.Tensor, spatial_dims: int, scale_factor: int) -> torch
     return x
 
 
+def pixelunshuffle(x: torch.Tensor, spatial_dims: int, scale_factor: int) -> torch.Tensor:
+    """
+    Apply pixel unshuffle to the tensor `x` with spatial dimensions `spatial_dims` and scaling factor `scale_factor`.
+    Inverse operation of pixelshuffle.
+
+    See: Shi et al., 2016, "Real-Time Single Image and Video Super-Resolution
+    Using an Efficient Sub-Pixel Convolutional Neural Network."
+
+    See: Aitken et al., 2017, "Checkerboard artifact free sub-pixel convolution".
+
+    Args:
+        x: Input tensor with shape BCHW[D]
+        spatial_dims: number of spatial dimensions, typically 2 or 3 for 2D or 3D
+        scale_factor: factor to reduce the spatial dimensions by, must be >=1
+
+    Returns:
+        Unshuffled version of `x` with shape (B, C*(r**d), H/r, W/r) for 2D
+        or (B, C*(r**d), D/r, H/r, W/r) for 3D, where r is the scale_factor
+        and d is spatial_dims.
+
+    Raises:
+        ValueError: When spatial dimensions are not divisible by scale_factor
+    """
+    dim, factor = spatial_dims, scale_factor
+    input_size = list(x.size())
+    batch_size, channels = input_size[:2]
+    scale_factor_mult = factor**dim
+    new_channels = channels * scale_factor_mult
+
+    if any(d % factor != 0 for d in input_size[2:]):
+        raise ValueError(
+            f"All spatial dimensions must be divisible by factor {factor}. " f", spatial shape is: {input_size[2:]}"
+        )
+    output_size = [batch_size, new_channels] + [d // factor for d in input_size[2:]]
+    reshaped_size = [batch_size, channels] + sum([[d // factor, factor] for d in input_size[2:]], [])
+
+    permute_indices = [0, 1] + [(2 * i + 3) for i in range(spatial_dims)] + [(2 * i + 2) for i in range(spatial_dims)]
+    x = x.reshape(reshaped_size).permute(permute_indices)
+    x = x.reshape(output_size)
+    return x
+
+
 @contextmanager
 def eval_mode(*nets: nn.Module):
     """

monai/utils/__init__.py

@@ -29,6 +29,7 @@ from .enums import (
     CommonKeys,
     CompInitMode,
     DiceCEReduction,
+    DownsampleMode,
     EngineStatsKeys,
     FastMRIKeys,
     ForwardMode,

monai/utils/enums.py

@@ -24,6 +24,7 @@ __all__ = [
     "SplineMode",
     "InterpolateMode",
     "UpsampleMode",
+    "DownsampleMode",
     "BlendMode",
     "PytorchPadMode",
     "NdimageMode",
@@ -181,6 +182,18 @@ class UpsampleMode(StrEnum):
     PIXELSHUFFLE = "pixelshuffle"
 
 
+class DownsampleMode(StrEnum):
+    """
+    See also: :py:class:`monai.networks.blocks.UpSample`
+    """
+
+    CONV = "conv"  # e.g. using strided convolution
+    CONVGROUP = "convgroup"  # e.g. using grouped strided convolution
+    PIXELUNSHUFFLE = "pixelunshuffle"
+    MAXPOOL = "maxpool"
+    AVGPOOL = "avgpool"
+
+
 class BlendMode(StrEnum):
     """
     See also: :py:class:`monai.data.utils.compute_importance_map`