(
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,
)
| 104 | """ |
| 105 | |
| 106 | def __init__( |
| 107 | self, |
| 108 | spatial_dims: int = 2, |
| 109 | in_channels: int = 3, |
| 110 | out_channels: int = 3, |
| 111 | dim: int = 48, |
| 112 | num_blocks: tuple[int, ...] = (1, 1, 1, 1), |
| 113 | heads: tuple[int, ...] = (1, 1, 1, 1), |
| 114 | num_refinement_blocks: int = 4, |
| 115 | ffn_expansion_factor: float = 2.66, |
| 116 | bias: bool = False, |
| 117 | layer_norm_use_bias: bool = True, |
| 118 | dual_pixel_task: bool = False, |
| 119 | flash_attention: bool = False, |
| 120 | ) -> None: |
| 121 | super().__init__() |
| 122 | """Initialize Restormer model. |
| 123 | |
| 124 | Args: |
| 125 | spatial_dims: Number of spatial dimensions (2D or 3D) |
| 126 | in_channels: Number of input image channels |
| 127 | out_channels: Number of output image channels |
| 128 | dim: Base feature dimension. Defaults to 48. |
| 129 | num_blocks: Number of transformer blocks at each scale. Defaults to (1,1,1,1). |
| 130 | heads: Number of attention heads at each scale. Defaults to (1,1,1,1). |
| 131 | num_refinement_blocks: Number of final refinement blocks. Defaults to 4. |
| 132 | ffn_expansion_factor: Expansion factor for feed-forward network. Defaults to 2.66. |
| 133 | bias: Whether to use bias in convolutions. Defaults to False. |
| 134 | layer_norm_use_bias: Whether to use bias in layer normalization. Defaults to True. |
| 135 | dual_pixel_task: Enable dual-pixel specific processing. Defaults to False. |
| 136 | flash_attention: Use flash attention if available. Defaults to False. |
| 137 | |
| 138 | Note: |
| 139 | The number of blocks must be greater than 1 |
| 140 | The length of num_blocks and heads must be equal |
| 141 | All values in num_blocks must be greater than 0 |
| 142 | """ |
| 143 | # Check input parameters |
| 144 | assert len(num_blocks) > 1, "Number of blocks must be greater than 1" |
| 145 | assert len(num_blocks) == len(heads), "Number of blocks and heads must be equal" |
| 146 | assert all(n > 0 for n in num_blocks), "Number of blocks must be greater than 0" |
| 147 | |
| 148 | # Initial feature extraction |
| 149 | self.patch_embed = OverlapPatchEmbed(spatial_dims, in_channels, dim) |
| 150 | self.encoder_levels = nn.ModuleList() |
| 151 | self.downsamples = nn.ModuleList() |
| 152 | self.decoder_levels = nn.ModuleList() |
| 153 | self.upsamples = nn.ModuleList() |
| 154 | self.reduce_channels = nn.ModuleList() |
| 155 | num_steps = len(num_blocks) - 1 |
| 156 | self.num_steps = num_steps |
| 157 | self.spatial_dims = spatial_dims |
| 158 | spatial_multiplier = 2 ** (spatial_dims - 1) |
| 159 | |
| 160 | # Define encoder levels |
| 161 | for n in range(num_steps): |
| 162 | current_dim = dim * (2) ** (n) |
| 163 | next_dim = current_dim // spatial_multiplier |
no test coverage detected