mmagic.models.editors.swinir.swinir_rstb

Module Contents

Classes

DropPath	Drop paths (Stochastic Depth) per sample (when applied in main path of
Mlp	Multilayer Perceptron layer.
WindowAttention	Window based multi-head self attention (W-MSA)
SwinTransformerBlock	Swin Transformer Block.
BasicLayer	A basic Swin Transformer layer for one stage.
RSTB	Residual Swin Transformer Block (RSTB).

class mmagic.models.editors.swinir.swinir_rstb.DropPath(drop_prob: float = 0.0, scale_by_keep: bool = True)

Bases: torch.nn.Module

Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).

forward(x)

Forward function.

Parameters

x (Tensor) – Input tensor with shape (B, L, C).

Returns

Forward results.

Return type

Tensor

extra_repr()

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class mmagic.models.editors.swinir.swinir_rstb.Mlp(in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0)

Bases: torch.nn.Module

Multilayer Perceptron layer.

Parameters

• in_features (int) – Number of input channels.

• hidden_features (int | None, optional) – Number of hidden layer channels. Default: None

• out_features (int | None, optional) – Number of output channels. Default: None

• act_layer (nn.Module, optional) – Activation layer. Default: nn.GELU

• drop (float, optional) – Dropout ratio of attention weight. Default: 0.0

forward(x)

Forward function.

Parameters

x (Tensor) – Input tensor with shape (B, L, C).

Returns

Forward results.

Return type

Tensor

class mmagic.models.editors.swinir.swinir_rstb.WindowAttention(dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0.0, proj_drop=0.0)

Bases: torch.nn.Module

Window based multi-head self attention (W-MSA) module with relative position bias. It supports both of shifted and non-shifted window. :param dim: Number of input channels. :type dim: int :param window_size: The height and width of the window. :type window_size: tuple[int] :param num_heads: Number of attention heads. :type num_heads: int :param qkv_bias: If True, add a learnable bias to

query, key, value. Default: True

Parameters

• qk_scale (float | None, optional) – Override default qk scale of head_dim ** -0.5 if set

• attn_drop (float, optional) – Dropout ratio of attention weight. Default: 0.0

• proj_drop (float, optional) – Dropout ratio of output. Default: 0.0

forward(x, mask=None)

Parameters

• x – input features with shape of (num_windows*B, N, C)

• mask – (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None

extra_repr() → str

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class mmagic.models.editors.swinir.swinir_rstb.SwinTransformerBlock(dim, input_resolution, num_heads, window_size=7, shift_size=0, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0, act_layer=nn.GELU, norm_layer=nn.LayerNorm)

Bases: torch.nn.Module

Swin Transformer Block. :param dim: Number of input channels. :type dim: int :param input_resolution: Input resolution. :type input_resolution: tuple[int] :param num_heads: Number of attention heads. :type num_heads: int :param window_size: Window size. :type window_size: int :param shift_size: Shift size for SW-MSA. :type shift_size: int :param mlp_ratio: Ratio of mlp hidden dim to embedding dim. :type mlp_ratio: float :param qkv_bias: If True, add a learnable bias

to query, key, value. Default: True

Parameters

• qk_scale (float | None, optional) – Override default qk scale of head_dim ** -0.5 if set.

• drop (float, optional) – Dropout rate. Default: 0.0

• attn_drop (float, optional) – Attention dropout rate. Default: 0.0

• drop_path (float, optional) – Stochastic depth rate. Default: 0.0

• act_layer (nn.Module, optional) – Activation layer. Default: nn.GELU

• norm_layer (nn.Module, optional) – Normalization layer. Default: nn.LayerNorm

calculate_mask(x_size)

Calculate attention mask for SW-MSA.

Parameters

x_size (tuple[int]) – Resolution of input feature.

Returns

Attention mask

Return type

Tensor

forward(x, x_size)

Forward function.

Parameters

• x (Tensor) – Input tensor with shape (B, L, C).

• x_size (tuple[int]) – Resolution of input feature.

Returns

Forward results.

Return type

Tensor

extra_repr() → str

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class mmagic.models.editors.swinir.swinir_rstb.BasicLayer(dim, input_resolution, depth, num_heads, window_size, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0, norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False)

Bases: torch.nn.Module

A basic Swin Transformer layer for one stage.

Parameters

• dim (int) – Number of input channels.

• input_resolution (tuple[int]) – Input resolution.

• depth (int) – Number of blocks.

• num_heads (int) – Number of attention heads.

• window_size (int) – Local window size.

• mlp_ratio (float) – Ratio of mlp hidden dim to embedding dim.

• qkv_bias (bool, optional) – If True, add a learnable bias to query, key, value. Default: True

• qk_scale (float | None, optional) – Override default qk scale of head_dim ** -0.5 if set.

• drop (float, optional) – Dropout rate. Default: 0.0

• attn_drop (float, optional) – Attention dropout rate. Default: 0.0

• drop_path (float | tuple[float], optional) – Stochastic depth rate. Default: 0.0

• norm_layer (nn.Module, optional) – Normalization layer. Default: nn.LayerNorm

• downsample (nn.Module | None, optional) – Downsample layer at the end of the layer. Default: None

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

forward(x, x_size)

Forward function.

Parameters

• x (Tensor) – Input tensor with shape (B, L, C).

• x_size (tuple[int]) – Resolution of input feature.

Returns

Forward results.

Return type

Tensor

extra_repr() → str

Set the extra representation of the module

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

class mmagic.models.editors.swinir.swinir_rstb.RSTB(dim, input_resolution, depth, num_heads, window_size, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0, norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False, img_size=224, patch_size=4, resi_connection='1conv')

Bases: torch.nn.Module

Residual Swin Transformer Block (RSTB).

Parameters

• dim (int) – Number of input channels.

• input_resolution (tuple[int]) – Input resolution.

• depth (int) – Number of blocks.

• num_heads (int) – Number of attention heads.

• window_size (int) – Local window size.

• mlp_ratio (float) – Ratio of mlp hidden dim to embedding dim. Default: 4.0

• qkv_bias (bool, optional) – If True, add a learnable bias to query, key, value. Default: True

• qk_scale (float | None, optional) – Override default qk scale of head_dim ** -0.5 if set.

• drop (float, optional) – Dropout rate. Default: 0.0

• attn_drop (float, optional) – Attention dropout rate. Default: 0.0

• drop_path (float | tuple[float], optional) – Stochastic depth rate. Default: 0.0

• norm_layer (nn.Module, optional) – Normalization layer. Default: nn.LayerNorm

• downsample (nn.Module | None, optional) – Downsample layer at the end of the layer. Default: None

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

• img_size (int) – Input image size. Default: 224

• patch_size (int) – Patch size. Default: 4

• resi_connection (string) – The convolutional block before residual connection. Default: ‘1conv’

forward(x, x_size)

Forward function.

Parameters

• x (Tensor) – Input tensor with shape (B, L, C).

• x_size (tuple[int]) – Resolution of input feature.

Returns

Forward results.

Return type

Tensor