# general settings
name: Allweather_Histotormer
model_type: ImageCleanModel
scale: 1
num_gpu: 4  # set num_gpu: 0 for cpu mode
manual_seed: 100

# dataset and data loader settings
datasets:
  train:
    name: TrainSet
    type: Dataset_PairedImage
    dataroot_gt: /home1/ssq/data/allweather/gt/
    dataroot_lq: /home1/ssq/data/allweather/input/
    geometric_augs: true

    filename_tmpl: '{}'
    io_backend:
      type: disk

    # data loader
    use_shuffle: true
    num_worker_per_gpu: 8
    batch_size_per_gpu: 8

    ### -------------Progressive training--------------------------
    mini_batch_sizes: [8,5,2,1,1]             # Batch size per gpu   
    iters: [92000,84000,56000,36000,32000]
    gt_size: 362   # Max patch size for progressive training
    gt_sizes: [128,160,256,320,362]  # Patch sizes for progressive training.
    ### ------------------------------------------------------------

    ### ------- Training on single fixed-patch size 128x128---------
    # mini_batch_sizes: [8]   
    # iters: [300000]
    # gt_size: 128   
    # gt_sizes: [128]
    ### ------------------------------------------------------------

    dataset_enlarge_ratio: 1
    prefetch_mode: ~

  val_snow_s:
    name: ValSet_Snow100K-S
    type: Dataset_PairedImage
    dataroot_gt: /home1/ssq/data/allweather/test/Snow100K-S/gt/
    dataroot_lq: /home1/ssq/data/allweather/test/Snow100K-S/synthetic/
    io_backend:
      type: disk
  val_snow_l:
    name: ValSet_Snow100K-L
    type: Dataset_PairedImage
    dataroot_gt: /home1/ssq/data/allweather/test/Snow100K-L/gt/
    dataroot_lq: /home1/ssq/data/allweather/test/Snow100K-L/synthetic/
    io_backend:
      type: disk
  val_test1:
    name: ValSet_Test1
    type: Dataset_PairedImage
    dataroot_gt: /home1/ssq/data/allweather/test/Test1/gt/
    dataroot_lq: /home1/ssq/data/allweather/test/Test1/input/
    io_backend:
      type: disk
  val_raindrop:
    name: ValSet_RainDrop
    type: Dataset_PairedImage
    dataroot_gt: /home1/ssq/data/allweather/test/RainDrop/gt/
    dataroot_lq: /home1/ssq/data/allweather/test/RainDrop/input/
    io_backend:
      type: disk


# network structures
network_g:
  type: Histoformer
  inp_channels: 3
  out_channels: 3
  dim: 36
  num_blocks: [4,4,6,8]
  num_refinement_blocks: 4
  heads: [1,2,4,8]
  ffn_expansion_factor: 2.667
  bias: False
  LayerNorm_type: WithBias
  dual_pixel_task: False


# path
path:
  pretrain_network_g: ~
  strict_load_g: true
  resume_state: ~

# training settings
train:
  total_iter: 300000
  warmup_iter: -1 # no warm up
  use_grad_clip: true

  # Split 300k iterations into two cycles. 
  # 1st cycle: fixed 3e-4 LR for 92k iters. 
  # 2nd cycle: cosine annealing (3e-4 to 1e-6) for 208k iters.
  scheduler:
    type: CosineAnnealingRestartCyclicLR # ReduceLROnPlateau
    periods: [92000, 208000]       
    restart_weights: [1,1]
    eta_mins: [0.0003,0.000001]   
  
  mixing_augs:
    mixup: false
    mixup_beta: 1.2
    use_identity: true

  optim_g:
    type: AdamW
    lr: !!float 3e-4
    weight_decay: !!float 1e-4
    betas: [0.9, 0.999]
  
  # losses
  pixel_opt:
    type: L1Loss
    loss_weight: 1
    reduction: mean
  seq_opt:
    type: Pearson

# validation settings
val:
  window_size: 8
  val_freq: !!float 1e3
  save_img: true
  rgb2bgr: true
  use_image: true
  max_minibatch: 8

  metrics:
    psnr: # metric name, can be arbitrary
      type: calculate_psnr
      crop_border: 0
      test_y_channel: true

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

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