mmagic.models.editors.swinir.swinir_net

Module Contents

Classes

SwinIRNet

SwinIR

class mmagic.models.editors.swinir.swinir_net.SwinIRNet(img_size=64, patch_size=1, in_chans=3, embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6], window_size=7, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.1, norm_layer=nn.LayerNorm, ape=False, patch_norm=True, use_checkpoint=False, upscale=2, img_range=1.0, upsampler='', resi_connection='1conv', **kwargs)

Bases: mmengine.model.BaseModule

SwinIR

A PyTorch impl of: SwinIR: Image Restoration Using Swin Transformer, based on Swin Transformer. Ref repo: https://github.com/JingyunLiang/SwinIR

Parameters

• img_size (int | tuple(int)) – Input image size. Default 64

• patch_size (int | tuple(int)) – Patch size. Default: 1

• in_chans (int) – Number of input image channels. Default: 3

• embed_dim (int) – Patch embedding dimension. Default: 96

• depths (tuple(int)) – Depth of each Swin Transformer layer. Default: [6, 6, 6, 6]

• num_heads (tuple(int)) – Number of attention heads in different layers. Default: [6, 6, 6, 6]

• window_size (int) – Window size. Default: 7

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

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

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

• drop_rate (float) – Dropout rate. Default: 0

• attn_drop_rate (float) – Attention dropout rate. Default: 0

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

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

• ape (bool) – If True, add absolute position embedding to the patch embedding. Default: False

• patch_norm (bool) – If True, add normalization after patch embedding. Default: True

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

• upscale (int) – Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction. Default: 2

• img_range (float) – Image range. 1. or 255. Default: 1.0

• upsampler (string, optional) – The reconstruction module. ‘pixelshuffle’ / ‘pixelshuffledirect’ /’nearest+conv’/None. Default: ‘’

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

_init_weights(m)

no_weight_decay()

no_weight_decay_keywords()

check_image_size(x)

Check image size and pad images so that it has enough dimension do window size.

Parameters

x – input tensor image with (B, C, H, W) shape.

forward_features(x)

Forward function of Deep Feature Extraction.

Parameters

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

Returns

Forward results.

Return type

Tensor

forward(x)

Forward function.

Parameters

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

Returns

Forward results.

Return type

Tensor