SwinIR#

SwinIR derived from Liang et al., 2021..

from pssr.models import SwinIR
pssr.models.SwinIR.__init__(self, image_size: int = 128, channels: list[int] = 1, scale: int = 4, embed_dim: int = 96, mlp_ratio: int = 2, depths: list[int] = [4, 4, 4, 4], num_heads: list[int] = [6, 6, 6, 6], window_size: int = 8, patch_size: int = 1, upsampler: str = 'pixelshuffle', qkv_bias: bool = True, qk_scale: float = None, drop_rate: float = 0, attn_drop_rate: float = 0, drop_path_rate: float = 0.1, norm_layer: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.normalization.LayerNorm'>, ape: bool = False, patch_norm: bool = True, use_checkpoint: bool = False, resi_connection: str = '1conv')#

SwinIR as detailed in Liang et al., 2021.

Parameters:

  • image_size (int) – Input image size.

  • channels (list[int]) – Number of channels in image data. Can also be a list of in channels (low-resolution) and out channels (high-resolution) respectively.

  • scale (int) – Upscaling factor for predictions. Choose a power of 2 for best results. Default is 4.

  • embed_dim (int) – Number of channels in patch embedding dimension. Default is 96.

  • mlp_ratio (int) – Ratio of hidden channels to embedding channels. Default is 2.

  • depths (list[int]) – Depths for each Swin Transformer layer. Default is 4 by 4 blocks.

  • num_heads (list[int]) – Number of attention heads for each Swin Transformer layer. Default is 6 by 4 blocks.

  • window_size (int) – Window size for Swin Transformer layers. Default is 8.

  • patch_size (int) – Patch size for embedding and Swin Transformer layers. Default is 1.

  • upsampler (str) – Upsamling algorithm for outputs. Options are “pixelshuffle”, “pixelshuffledirect”, “nearest+conv”, and None. Default is “pixelshuffle”.

  • qkv_bias (bool) – Whether to add learnable bias to query. Default is True.

  • qk_scale (float) – Override default qk scale of head_dim ** -0.5 if set. Default is None.

  • drop_rate (float) – Dropout rate during training. Default is 0.

  • attn_drop_rate (float) – Attention dropout rate during training. Default is 0.

  • drop_path_rate (float) – Stochastic depth rate. Default is 0.1.

  • norm_layer (nn.Module) – Data normalization layer. Default is nn.LayerNorm.

  • ape (bool) – Whether to add absolute position embedding to patch embedding. Default is False.

  • patch_norm (bool) – Whether to use normalization after patch embedding. Default is False.

  • use_checkpoint (bool) – Whether to use checkpointing to save memory. Default is False.

  • resi_connection (str) – Which convolutional block to use before residual connection. Options are “1conv” or “3conv”. Default is “1conv”.