mmagic.models.editors.nafnet.nafnet_net
¶
Module Contents¶
Classes¶
NAFNet. |
|
The original version of NAFNetLocal in "Simple Baseline for Image |
|
NAFNet's Block in paper. |
|
The Simple Gate in "Simple Baseline for Image Restoration". |
- class mmagic.models.editors.nafnet.nafnet_net.NAFNet(img_channels=3, mid_channels=16, middle_blk_num=1, enc_blk_nums=[], dec_blk_nums=[])[source]¶
Bases:
mmengine.model.BaseModule
NAFNet.
The original version of NAFNet in “Simple Baseline for Image Restoration”.
- Parameters
img_channels (int) – Channel number of inputs.
mid_channels (int) – Channel number of intermediate features.
middle_blk_num (int) – Number of middle blocks.
enc_blk_nums (List of int) – Number of blocks for each encoder.
dec_blk_nums (List of int) – Number of blocks for each decoder.
- class mmagic.models.editors.nafnet.nafnet_net.NAFNetLocal(*args, train_size=(1, 3, 256, 256), fast_imp=False, **kwargs)[source]¶
Bases:
mmagic.models.editors.nafnet.naf_avgpool2d.Local_Base
,NAFNet
The original version of NAFNetLocal in “Simple Baseline for Image Restoration”.
NAFNetLocal uses local average pooling modules than NAFNet.
- Parameters
img_channels (int) – Channel number of inputs.
mid_channels (int) – Channel number of intermediate features.
middle_blk_num (int) – Number of middle blocks.
enc_blk_nums (List of int) – Number of blocks for each encoder.
dec_blk_nums (List of int) – Number of blocks for each decoder.
- class mmagic.models.editors.nafnet.nafnet_net.NAFBlock(in_channels, DW_Expand=2, FFN_Expand=2, drop_out_rate=0.0)[source]¶
Bases:
mmengine.model.BaseModule
NAFNet’s Block in paper.
Simple gate will shrink the channel to a half. To keep the number of channels, it expands the channels first.
- Parameters
in_channels (int) – number of channels
DW_Expand (int) – channel expansion factor for part 1
FFN_Expand (int) – channel expansion factor for part 2
drop_out_rate (float) – drop out ratio
- class mmagic.models.editors.nafnet.nafnet_net.SimpleGate(init_cfg: Union[dict, List[dict], None] = None)[source]¶
Bases:
mmengine.model.BaseModule
The Simple Gate in “Simple Baseline for Image Restoration”.
- Parameters
x – input tensor feature map with (B, 2 * C, H, W)
- Returns
x1 * x2 (where x1, x2 are two separate parts by simple split x to [B, C, H, W])