train_bthnrm.yml

name: bth_normalmap_colab
use_tb_logger: true
model: srragan
scale: 1
gpu_ids: [0]
use_amp: true
use_swa: false

# Dataset options:
datasets:
  train:
    name: DIV2K
    mode: LRHROTF
    dataroot_HR: '/content/BasicSR/bthnrm/hr'
    dataroot_LR: '/content/BasicSR/bthnrm/lr'
    subset_file: null
    use_shuffle: true
    znorm: false
    n_workers: 4
    batch_size: 4
    virtual_batch_size: 4
    HR_size: 128
    image_channels: 3

    lr_downscale: false
    lr_downscale_types: ["nearest"]

    hr_downscale: false
    hr_downscale_types: ["nearest"]

    # Rotations augmentations:
    use_flip: true
    use_rot: true
    hr_rrot: false

    # Noise and blur augmentations:
    lr_blur: false # true | false
    lr_blur_types: {gaussian: 1, clean: 3} # select from: "average","box","gaussian","bilateral","clean" ##blur options #median and motion aren't working yet
    noise_data: ../noise_patches/normal/ # location of the noise patches extracted from real images to use for noise injection with noise option "patches"
    lr_noise: false # true | false
    lr_noise_types: {gaussian: 1, JPEG: 1, clean: 4} # select from: "gaussian", "JPEG", "quantize", "poisson", "dither", "s&p", "speckle", "patches", "clean"
    lr_noise2: false # true | false
    lr_noise_types2: {dither: 2, clean: 2} # select from: "gaussian", "JPEG", "quantize", "poisson", "dither", "s&p", "speckle", "patches", "clean"
    hr_noise: false # true | false
    hr_noise_types: {gaussian: 1, clean: 4} # select from: "gaussian", "JPEG", "quantize", "poisson", "dither", "s&p", "speckle", "clean"

  val:
    name: val_set14_part
    mode: LRHROTF
    dataroot_HR: '/content/BasicSR/bthnrm/hr_val'
    dataroot_LR: '/content/BasicSR/bthnrm/lr_val'

path:
    strict: false
    root: '/content/drive/My Drive/Models/'
    pretrain_model_G: '/content/BasicSR/experiments/pretrained_models/1xESRGAN.pth'
    resume_state: '/content/drive/My Drive/Models/experiments/bth_normalmap_colab/training_state/latest.state'

# Generator:
network_G:
    strict: false # true | false # whether to load the model in strict mode or not
    # ESRGAN:
    which_model_G: RRDB_net # RRDB_net (original ESRGAN arch) | MRRDB_net (modified/"new" arch) | sr_resnet
    norm_type: null
    mode: CNA
    nf: 64 # number of discrim filters in the first conv layer
    nb: 23 # number of RRDB blocks
    nr: 3 #  number of residual layers in each RRDB block
    in_nc: 3 # of input image channels: 3 for RGB and 1 for grayscale
    out_nc: 3 # of output image channels: 3 for RGB and 1 for grayscale
    gc: 32
    group: 1
    convtype: Conv2D # Conv2D | PartialConv2D | DeformConv2D | Conv3D
    net_act: leakyrelu # swish | leakyrelu
    gaussian: true # true | false
    plus: false # true | false

# Discriminator:
network_D:
    strict: true # true | false # whether to load the model in strict mode or not
    # ESRGAN (default)| PPON:
    which_model_D: discriminator_vgg # discriminator_vgg_128 | discriminator_vgg | discriminator_vgg_128_fea (feature extraction) | discriminator_vgg_fea (feature extraction) | patchgan | multiscale
    norm_type: batch
    act_type: leakyrelu
    mode: CNA # CNA | NAC
    nf: 64
    in_nc: 3
    nlayer: 3 # only for patchgan and multiscale
    num_D: 3 # only for multiscale

# Training options:
train:
    lr_G: 0.0001 # 2e-4 # starting lr_g #Test, default: 1e-4
    weight_decay_G: 0
    beta1_G: 0.9
    lr_D: 0.0001 # 2e-4 # starting lr_d #Test, default: 1e-4
    weight_decay_D: 0
    beta1_D: 0.9

    # For MultiStepLR (ESRGAN, default):
    lr_scheme: MultiStepLR
    lr_steps_rel: [0.1, 0.2, 0.4, 0.6] # to use lr steps relative to % of training niter instead of fixed lr_steps
    lr_gamma: 0.5 # lr change at every step (multiplied by)

    # For SWA scheduler
    swa_start_iter: 375000 #Just reference: 75% of 500000. Can be any value, including 0 to start right away with a pretrained model.
    # swa_start_iter_rel: 0.75 # to use swa_start_iter relative to % of training niter instead of fixed swa_start_iter
    swa_lr: 1e-4 #Has to be ~order of magnitude of a stable lr for the regular scheduler
    swa_anneal_epochs: 10
    swa_anneal_strategy: "cos"

    # Losses:
    pixel_criterion: l1 # "l1" | "l2" | "cb" | "elastic" | "relativel1" | "l1cosinesim" | "clipl1" #pixel loss
    pixel_weight: 1e-2 # 1e-2 | 1
    feature_criterion: l1 # "l1" | "l2" | "cb" | "elastic" #feature loss (VGG feature network)
    feature_weight: 1
    cx_weight: 0.5
    cx_type: contextual
    cx_vgg_layers: {conv_3_2: 1, conv_4_2: 1}
    gan_type: 'vanilla'
    gan_weight: 5e-3

# Other training options:
    manual_seed: 0
    niter: 5e5
    # warmup_iter: -1  # number of warm up iterations, -1 for no warm up
    val_freq: 500 # 5e3
    # overwrite_val_imgs: true
    # val_comparison: true
    metrics: 'psnr,ssim,lpips' # select from: "psnr,ssim,lpips" or a combination separated by comma ","

logger:
    print_freq: 200
    save_checkpoint_freq: !!float 5e3
    overwrite_chkp: false