[Feature] Support NAFNet model

configs/nafnet/nafnet_c64eb11128mb1db1111_lr1e-3_400k_gopro.py

@@ -0,0 +1,126 @@
+_base_ = '../_base_/default_runtime.py'
+
+experiment_name = 'nafnet_c64eb11128mb1db1111_lr1e-3_400k_gopro'
+work_dir = f'./work_dirs/{experiment_name}'
+save_dir = './work_dirs/'
+
+# DistributedDataParallel
+model_wrapper_cfg = dict(type='MMSeparateDistributedDataParallel')
+
+# model settings
+model = dict(
+    type='BaseEditModel',
+    generator=dict(
+        type='NAFNetLocal',
+        img_channel=3,
+        mid_channels=64,
+        enc_blk_nums=[1, 1, 1, 28],
+        middle_blk_num=1,
+        dec_blk_nums=[1, 1, 1, 1],
+    ),
+    pixel_loss=dict(type='PSNRLoss'),
+    train_cfg=dict(),
+    test_cfg=dict(),
+    data_preprocessor=dict(
+        type='EditDataPreprocessor',
+        mean=[0., 0., 0.],
+        std=[255., 255., 255.],
+    ))
+
+train_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='SetValues', dictionary=dict(scale=1)),
+    dict(
+        type='Flip',
+        keys=['img', 'gt'],
+        flip_ratio=0.5,
+        direction='horizontal'),
+    dict(
+        type='Flip', keys=['img', 'gt'], flip_ratio=0.5, direction='vertical'),
+    dict(type='RandomTransposeHW', keys=['img', 'gt'], transpose_ratio=0.5),
+    dict(type='PairedRandomCrop', gt_patch_size=256),
+    dict(type='PackEditInputs')
+]
+
+val_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='PackEditInputs')
+]
+
+# dataset settings
+dataset_type = 'BasicImageDataset'
+
+train_dataloader = dict(
+    num_workers=8,
+    batch_size=8,  # gpus 4
+    persistent_workers=False,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='gopro', task_name='deblur'),
+        data_root='../datasets/gopro/train',
+        data_prefix=dict(gt='sharp', img='blur'),
+        ann_file='meta_info_gopro_train.txt',
+        pipeline=train_pipeline))
+
+val_dataloader = dict(
+    num_workers=4,
+    persistent_workers=False,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='gopro', task_name='deblur'),
+        data_root='../datasets/gopro/test',
+        ann_file='meta_info_gopro_test.txt',
+        data_prefix=dict(gt='sharp', img='blur'),
+        pipeline=val_pipeline))
+
+test_dataloader = val_dataloader
+
+val_evaluator = [
+    dict(type='MAE'),
+    dict(type='PSNR'),
+    dict(type='SSIM'),
+]
+test_evaluator = val_evaluator
+
+train_cfg = dict(
+    type='IterBasedTrainLoop', max_iters=400_000, val_interval=20000)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+# optimizer
+# optim_wrapper = dict(
+#     constructor='MultiOptimWrapperConstructor',
+#     generator=dict(
+#         type='OptimWrapper',
+#         optimizer=dict(type='Adam', lr=1e-3, betas=(0.9, 0.9))))
+optim_wrapper = dict(
+    constructor='DefaultOptimWrapperConstructor',
+    type='OptimWrapper',
+    optimizer=dict(type='AdamW', lr=1e-3, weight_decay=1e-3, betas=(0.9, 0.9)))
+
+# learning policy
+param_scheduler = dict(
+    type='CosineAnnealingLR', by_epoch=False, T_max=400_000, eta_min=1e-7)
+
+default_hooks = dict(
+    checkpoint=dict(
+        type='CheckpointHook',
+        interval=5000,
+        save_optimizer=True,
+        by_epoch=False,
+        out_dir=save_dir,
+    ),
+    timer=dict(type='IterTimerHook'),
+    logger=dict(type='LoggerHook', interval=100),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+)
+
+visualizer = dict(bgr2rgb=False)
+
+randomness = dict(seed=10, diff_rank_seed=True)

configs/nafnet/nafnet_c64eb2248mb12db2222_lr1e-3_400k_sidd.py

@@ -0,0 +1,120 @@
+_base_ = '../_base_/default_runtime.py'
+
+experiment_name = 'nafnet_c64eb2248mb12db2222_lr1e-3_400k_sidd'
+work_dir = f'./work_dirs/{experiment_name}'
+save_dir = './work_dirs/'
+
+# DistributedDataParallel
+model_wrapper_cfg = dict(type='MMSeparateDistributedDataParallel')
+
+# model settings
+model = dict(
+    type='BaseEditModel',
+    generator=dict(
+        type='NAFNet',
+        img_channel=3,
+        mid_channels=64,
+        enc_blk_nums=[2, 2, 4, 8],
+        middle_blk_num=12,
+        dec_blk_nums=[2, 2, 2, 2],
+    ),
+    pixel_loss=dict(type='PSNRLoss'),
+    train_cfg=dict(),
+    test_cfg=dict(),
+    data_preprocessor=dict(
+        type='EditDataPreprocessor',
+        mean=[0.0, 0.0, 0.0],
+        std=[255.0, 255.0, 255.0],
+    ))
+
+train_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='SetValues', dictionary=dict(scale=1)),
+    dict(
+        type='Flip',
+        keys=['img', 'gt'],
+        flip_ratio=0.5,
+        direction='horizontal'),
+    dict(
+        type='Flip', keys=['img', 'gt'], flip_ratio=0.5, direction='vertical'),
+    dict(type='RandomTransposeHW', keys=['img', 'gt'], transpose_ratio=0.5),
+    dict(type='PairedRandomCrop', gt_patch_size=256),
+    dict(type='PackEditInputs')
+]
+
+val_pipeline = [
+    dict(type='LoadImageFromFile', key='img'),
+    dict(type='LoadImageFromFile', key='gt'),
+    dict(type='PackEditInputs')
+]
+
+# dataset settings
+dataset_type = 'BasicImageDataset'
+
+train_dataloader = dict(
+    num_workers=8,
+    batch_size=8,  # gpus 4
+    persistent_workers=False,
+    sampler=dict(type='InfiniteSampler', shuffle=True),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='sidd', task_name='denoising'),
+        data_root='../datasets/SIDD/train',
+        data_prefix=dict(gt='gt', img='noisy'),
+        filename_tmpl=dict(img='{}_NOISY', gt='{}_GT'),
+        pipeline=train_pipeline))
+
+val_dataloader = dict(
+    num_workers=4,
+    persistent_workers=False,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False),
+    dataset=dict(
+        type=dataset_type,
+        metainfo=dict(dataset_type='sidd', task_name='denoising'),
+        data_root='../datasets/SIDD/val/',
+        data_prefix=dict(gt='gt', img='noisy'),
+        filename_tmpl=dict(gt='{}_GT', img='{}_NOISY'),
+        pipeline=val_pipeline))
+
+test_dataloader = val_dataloader
+
+val_evaluator = [
+    dict(type='MAE'),
+    dict(type='PSNR'),
+    dict(type='SSIM'),
+]
+test_evaluator = val_evaluator
+
+train_cfg = dict(
+    type='IterBasedTrainLoop', max_iters=400_000, val_interval=20000)
+val_cfg = dict(type='ValLoop')
+test_cfg = dict(type='TestLoop')
+
+# optimizer
+optim_wrapper = dict(
+    constructor='MultiOptimWrapperConstructor',
+    generator=dict(
+        type='OptimWrapper',
+        optimizer=dict(type='AdamW', lr=1e-3, betas=(0.9, 0.9))))
+
+# learning policy
+param_scheduler = dict(
+    type='CosineAnnealingLR', by_epoch=False, T_max=400_000, eta_min=1e-7)
+
+default_hooks = dict(
+    checkpoint=dict(
+        type='CheckpointHook',
+        interval=5000,
+        save_optimizer=True,
+        by_epoch=False,
+        out_dir=save_dir,
+    ),
+    timer=dict(type='IterTimerHook'),
+    logger=dict(type='LoggerHook', interval=100),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+)
+
+visualizer = dict(bgr2rgb=True)

mmedit/models/editors/__init__.py

@@ -31,6 +31,7 @@
 from .liif import LIIF, MLPRefiner
 from .lsgan import LSGAN
 from .mspie import MSPIEStyleGAN2, PESinGAN
+from .nafnet import Baseline, BaselineLocal, NAFNet, NAFNetLocal
 from .pconv import (MaskConvModule, PartialConv2d, PConvDecoder, PConvEncoder,
                     PConvEncoderDecoder, PConvInpaintor)
 from .pggan import ProgressiveGrowingGAN
@@ -73,5 +74,6 @@
     'FBADecoder', 'WGANGP', 'CycleGAN', 'SAGAN', 'LSGAN', 'GGAN', 'Pix2Pix',
     'StyleGAN1', 'StyleGAN2', 'StyleGAN3', 'BigGAN', 'DCGAN',
     'ProgressiveGrowingGAN', 'SinGAN', 'IDLossModel', 'PESinGAN',
-    'MSPIEStyleGAN2', 'StyleGAN3Generator'
+    'MSPIEStyleGAN2', 'StyleGAN3Generator', 'Baseline', 'BaselineLocal',
+    'NAFNet', 'NAFNetLocal'
 ]

mmedit/models/editors/nafnet/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from .baseline_net import Baseline, BaselineLocal
+from .nafnet_net import NAFNet, NAFNetLocal
+
+__all__ = [
+    'NAFNet',
+    'NAFNetLocal',
+    'Baseline',
+    'BaselineLocal',
+]