nafnet_denoising.yaml

total_iters: 400000
output_dir: output_dir

model:
  name: NAFNetModel
  generator:
    name: NAFNet
    img_channel: 3
    width: 64
    enc_blk_nums: [2, 2, 4, 8]
    middle_blk_num: 12
    dec_blk_nums: [2, 2, 2, 2]
  psnr_criterion:
    name: PSNRLoss

dataset:
  train:
    name: NAFNetTrain
    rgb_dir: data/SIDD/train
    # TODO fix out of memory for val while training
    num_workers: 0
    batch_size: 8 # 8GPU
    img_options:
      patch_size: 256
  test:
    name: NAFNetVal
    rgb_dir: data/SIDD/val
    # TODO fix out of memory for val while training
    num_workers: 0
    batch_size: 1
    img_options:
      patch_size: 256

export_model:
  - {name: 'generator', inputs_num: 1}

lr_scheduler:
  name: CosineAnnealingRestartLR
  learning_rate: 0.001
  periods: [400000]
  restart_weights: [1]
  eta_min: !!float 8e-7

validate:
  interval: 5000
  save_img: false

  metrics:
    psnr: # metric name, can be arbitrary
      name: PSNR
      crop_border: 4
      test_y_channel: True
    ssim:
      name: SSIM
      crop_border: 4
      test_y_channel: True

optimizer:
  name: AdamW
  # add parameters of net_name to optim
  # name should in self.nets
  net_names:
    - generator
  weight_decay: 0.0
  beta1: 0.9
  beta2: 0.9
  epsilon: 1e-8

log_config:
  interval: 10
  visiual_interval: 5000

snapshot_config:
  interval: 5000