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?

[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:

[ericspod]

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

[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?