Restormer Implementation

[phisanti]

Fixes # .

Description

This PR implements the Restormer architecture for high-resolution image restoration in MONAI following the discussion in issue #8261. The implementation supports both 2D and 3D images using MONAI's convolution as the base. Key additions include:

• Downsample class for efficient downsampling operations
• pixel_unshuffle operation complementing existing pixel_shuffle
• Channel Attention Block (CABlock) with FeedForward layer
• Multi-DConv Head Transposed Self-Attention (MDTA)
• OverlapPatchEmbed class
• Comprehensive unit tests for all new components

The implementation follows MONAI's coding patterns and includes performance validations against native PyTorch operations where applicable.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[ericspod]

Looks good overall but I had a few inline comments, and we should have full docstrings everywhere appropriate. For any classes meant for general purpose use (ie. not just by Restormer) please ensure they have docstring descriptions for the arguments (at the very least for constructor args). Thanks!

[ericspod]

Here we should specifically state that x has shape BCHW[D].

[ericspod]

Suggested change

	f"All spatial dimensions must be divisible by factor {factor}. " f"Got spatial dimensions: {input_size[2:]}"
	f"All spatial dimensions must be divisible by {factor}, spatial shape is: {input_size[2:]}"

Maybe a little shorter?

[phisanti]

Done

[ericspod]

Suggested change

	if down_mode == "conv":
	if down_mode == DownsampleMode.CONV:

[ericspod]

Suggested change

	elif down_mode == "convgroup":
	elif down_mode == DownsampleMode.CONVGROUP:

[ericspod]

Suggested change

	elif down_mode == "pixelunshuffle":
	elif down_mode == DownsampleMode.PIXELSHUFFLE:

[phisanti]

done, but I used DownsampleMode.PIXELUNSHUFFLE as in downsampling the restormer uses pixel_unshuffling while pixel_shuffling is reserved for upsampling.

[ericspod]

We should have a full docstring here describing the arguments for the constructor, and in the previous class.

[phisanti]

Done

[ericspod]

Suggested change

	class OverlapPatchEmbed(nn.Module):
	"""Initial feature extraction using overlapped convolutions.
	Unlike standard patch embeddings that use non-overlapping patches,
	this approach maintains spatial continuity through 3x3 convolutions."""
	def __init__(self, spatial_dims: int, in_c: int = 3, embed_dim: int = 48, bias: bool = False):
	super().__init__()
	self.proj = Convolution(
	spatial_dims=spatial_dims,
	in_channels=in_c,
	out_channels=embed_dim,
	kernel_size=3,
	strides=1,
	padding=1,
	bias=bias,
	conv_only=True,
	)
	def forward(self, x: torch.Tensor) -> torch.Tensor:
	return self.proj(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.
	"""
	def __init__(self, spatial_dims: int, in_c: int = 3, embed_dim: int = 48, bias: bool = False):
	super().__init__(
	spatial_dims=spatial_dims,
	in_channels=in_c,
	out_channels=embed_dim,
	kernel_size=3,
	strides=1,
	padding=1,
	bias=bias,
	conv_only=True,
	)

Would it work to inherit directly from Convolution?

[phisanti]

Works! very elegant suggestion btw!

[aylward]

Thank you for this outstanding contribution!

[aylward]

@ericspod - do we normally provide alternative capitalizations to functions? "Downsample" is the generally accepted term (vs "Down Sample" which is less common).

I suggest using "Downsample" throughout, unless we offer alternative usage elsewhere that I haven't encountered. IDE auto-complete can help folks get the right capitalization.

Looking at your enums, you use "Downsample" which confirms (in my mind) that we should be using "Downsample" throughout.

[ericspod]

We have had this mechanism for other classes so other capitalisation could be used in scripts, such as "Transformd" and "TransformD". Here I'd say we don't need it though.

[phisanti]

Here, I mirrored the naming style of the Upsample class. Happy either way to remove it or keep it:

monai/networks/blocks/upsample.py

Upsample = UpSample
Subpixelupsample = SubpixelUpSample = SubpixelUpsample

[aylward]

Spell-out in_c to in_channels

[phisanti]

Done

[aylward]

in_channels

[phisanti]

Done

[aylward]

Make enum or convert to Bool (e.g., layer_norm_use_bias).

[aylward]

Looks good .... as long as the runtime is reasonable.