# general settings
name: train_HAT-L_SRx2_finetune_from_ImageNet_pretrain
model_type: HATModel
scale: 2
num_gpu: auto
manual_seed: 0

# dataset and data loader settings
datasets:
  train:
    name: DF2K
    type: PairedImageDataset
    dataroot_gt: datasets/DF2K/DF2K_HR_sub
    dataroot_lq: datasets/DF2K/DF2K_bicx2_sub
    meta_info_file: hat/data/meta_info/meta_info_DF2Ksub_GT.txt
    io_backend:
      type: disk

    gt_size: 128
    use_hflip: true
    use_rot: true

    # data loader
    use_shuffle: true
    num_worker_per_gpu: 6
    batch_size_per_gpu: 4
    dataset_enlarge_ratio: 1
    prefetch_mode: ~

  val_1:
    name: Set5
    type: PairedImageDataset
    dataroot_gt: ./datasets/Set5/GTmod2
    dataroot_lq: ./datasets/Set5/LRbicx2
    io_backend:
      type: disk

  val_2:
    name: Set14
    type: PairedImageDataset
    dataroot_gt: ./datasets/Set14/GTmod2
    dataroot_lq: ./datasets/Set14/LRbicx2
    io_backend:
      type: disk

  # val_3:
  #   name: Urban100
  #   type: PairedImageDataset
  #   dataroot_gt: ./datasets/urban100/GTmod2
  #   dataroot_lq: ./datasets/urban100/LRbicx2
  #   io_backend:
  #     type: disk


# network structures
network_g:
  type: HAT
  upscale: 2
  in_chans: 3
  img_size: 64
  window_size: 16
  compress_ratio: 3
  squeeze_factor: 30
  conv_scale: 0.01
  overlap_ratio: 0.5
  img_range: 1.
  depths: [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]
  embed_dim: 180
  num_heads: [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]
  mlp_ratio: 2
  upsampler: 'pixelshuffle'
  resi_connection: '1conv'

# path
path:
  pretrain_network_g: ./experiments/train_HAT-L_SRx2_ImageNet_from_scratch/models/net_g_latest.pth
  param_key_g: 'params_ema'
  strict_load_g: true
  resume_state: ~

# training settings
train:
  ema_decay: 0.999
  optim_g:
    type: Adam
    lr: !!float 1e-5
    weight_decay: 0
    betas: [0.9, 0.99]

  scheduler:
    type: MultiStepLR
    milestones: [125000, 200000, 225000, 240000]
    gamma: 0.5

  total_iter: 250000
  warmup_iter: -1  # no warm up

  # losses
  pixel_opt:
    type: L1Loss
    loss_weight: 1.0
    reduction: mean

# validation settings
val:
  val_freq: !!float 5e3
  save_img: false
  pbar: False

  metrics:
    psnr:
      type: calculate_psnr
      crop_border: 2
      test_y_channel: true
      better: higher  # the higher, the better. Default: higher
    ssim:
      type: calculate_ssim
      crop_border: 2
      test_y_channel: true
      better: higher  # the higher, the better. Default: higher

# logging settings
logger:
  print_freq: 100
  save_checkpoint_freq: !!float 5e3
  use_tb_logger: true
  wandb:
    project: ~
    resume_id: ~

# dist training settings
dist_params:
  backend: nccl
  port: 29500
