Shortcuts

mmedit.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 mmedit.models.editors.swinir.swinir_rstb.DropPath(drop_prob: float = 0.0, scale_by_keep: bool = True)[source]

Bases: torch.nn.Module

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

forward(x)[source]

Forward function.

Parameters

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

Returns

Forward results.

Return type

Tensor

extra_repr()[source]

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 mmedit.models.editors.swinir.swinir_rstb.Mlp(in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0)[source]

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)[source]

Forward function.

Parameters

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

Returns

Forward results.

Return type

Tensor

class mmedit.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)[source]

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)[source]
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[source]

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 mmedit.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)[source]

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)[source]

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)[source]

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[source]

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 mmedit.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)[source]

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)[source]

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[source]

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 mmedit.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')[source]

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)[source]

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

Read the Docs v: latest
Versions
master
latest
stable
zyh-re-docs
zyh-doc-notfound-extend
zyh-api-rendering
Downloads
pdf
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.