swinir_ffhq512.yaml

trainer:
  target: trainer.TrainerSR

model:
  target: models.swinir.SwinIR
  ckpt_path: ~
  params:
    img_size: 64
    patch_size: 1
    in_chans: 3
    embed_dim: 180
    depths: [6, 6, 6, 6, 6, 6, 6, 6]
    num_heads: [6, 6, 6, 6, 6, 6, 6, 6]
    window_size: 8
    mlp_ratio: 2
    sf: 8
    img_range: 1.0
    upsampler: "nearest+conv"
    resi_connection: "1conv"
    unshuffle: True
    unshuffle_scale: 8

data:
  train:
    type: gfpgan 
    params:
      files_txt: /mnt/lustre/share/zsyue/data/FFHQ/files_txt/ffhq512.txt
      io_backend:
        type: disk
   
      use_hflip: true
      mean: [0.0, 0.0, 0.0]
      std: [1.0, 1.0, 1.0]
      out_size: 512
   
      blur_kernel_size: 41
      kernel_list: ['iso', 'aniso']
      kernel_prob: [0.5, 0.5]
      blur_sigma: [0.1, 15]
      downsample_range: [0.8, 32]
      noise_range: [0, 20]
      jpeg_range: [30, 100]
   
      color_jitter_prob: ~
      color_jitter_pt_prob: ~
      gray_prob: 0.01
      gt_gray: True

      need_gt_path: False
  val:
    type: folder 
    params:
      dir_path: /mnt/lustre/zsyue/projects/DifFace/Restoration/testdata/CelebA-Test/lq
      dir_path_extra: /mnt/lustre/zsyue/projects/DifFace/Restoration/testdata/CelebA-Test/hq
      transform_type: default
      transform_kwargs:
        mean: 0.0
        std: 1.0
      im_exts: png
      length: 500
      need_path: False
      recursive: False

train:
  lr: 5e-5
  lr_min: 5e-6
  batch: [32, 16]   # batchsize for training and validation
  microbatch: 8
  num_workers: 4
  prefetch_factor: 2
  iterations: 300000
  weight_decay: 0
  save_freq: 10000
  val_freq: ${train.save_freq}
  log_freq: [100, 2000, 10]
  loss_type: L1
  tf_logging: True
  local_logging: True