imageEnhencer / options /CodeFormer_stage3.yml

Adding

184401d over 1 year ago

3.61 kB

# general settings
name: CodeFormer_stage3
model_type: CodeFormerJointModel
num_gpu: 8
manual_seed: 0

# dataset and data loader settings
datasets:
  train:
    name: FFHQ
    type: FFHQBlindJointDataset
    dataroot_gt: datasets/ffhq/ffhq_512
    filename_tmpl: '{}'
    io_backend:
      type: disk

    in_size: 512
    gt_size: 512
    mean: [0.5, 0.5, 0.5]
    std: [0.5, 0.5, 0.5]
    use_hflip: true
    use_corrupt: true

    blur_kernel_size: 41
    use_motion_kernel: false
    motion_kernel_prob: 0.001
    kernel_list: ['iso', 'aniso']
    kernel_prob: [0.5, 0.5]
    # small degradation in stageIII
    blur_sigma: [0.1, 10]
    downsample_range: [1, 12]
    noise_range: [0, 15]
    jpeg_range: [60, 100]
    # large degradation in stageII
    blur_sigma_large: [1, 15]
    downsample_range_large: [4, 30]
    noise_range_large: [0, 20]
    jpeg_range_large: [30, 80]

    latent_gt_path: ~ # without pre-calculated latent code
    # latent_gt_path: './experiments/pretrained_models/VQGAN/latent_gt_code1024.pth'

    # data loader
    num_worker_per_gpu: 1
    batch_size_per_gpu: 3
    dataset_enlarge_ratio: 100
    prefetch_mode: ~

  # val:
  #   name: CelebA-HQ-512
  #   type: PairedImageDataset
  #   dataroot_lq: datasets/faces/validation/lq
  #   dataroot_gt: datasets/faces/validation/gt
  #   io_backend:
  #     type: disk
  #   mean: [0.5, 0.5, 0.5]
  #   std: [0.5, 0.5, 0.5]
  #   scale: 1
    
# network structures
network_g:
  type: CodeFormer
  dim_embd: 512
  n_head: 8
  n_layers: 9
  codebook_size: 1024
  connect_list: ['32', '64', '128', '256']
  fix_modules: ['quantize','generator']

network_vqgan: # this config is needed if no pre-calculated latent
  type: VQAutoEncoder
  img_size: 512
  nf: 64
  ch_mult: [1, 2, 2, 4, 4, 8]
  quantizer: 'nearest'
  codebook_size: 1024

network_d:
  type: VQGANDiscriminator
  nc: 3
  ndf: 64
  n_layers: 4

# path
path:
  pretrain_network_g: './experiments/pretrained_models/CodeFormer_stage2/net_g_latest.pth' # pretrained G model in StageII 
  param_key_g: params_ema
  strict_load_g: false
  pretrain_network_d: './experiments/pretrained_models/CodeFormer_stage2/net_d_latest.pth' # pretrained D model in StageII
  resume_state: ~

# base_lr(4.5e-6)*bach_size(4)
train:
  use_hq_feat_loss: true
  feat_loss_weight: 1.0
  cross_entropy_loss: true
  entropy_loss_weight: 0.5
  scale_adaptive_gan_weight: 0.1

  optim_g:
    type: Adam
    lr: !!float 5e-5
    weight_decay: 0
    betas: [0.9, 0.99]
  optim_d:
    type: Adam
    lr: !!float 5e-5
    weight_decay: 0
    betas: [0.9, 0.99]

  scheduler:
    type: CosineAnnealingRestartLR
    periods: [150000]
    restart_weights: [1]
    eta_min: !!float 2e-5


  total_iter: 150000

  warmup_iter: -1  # no warm up
  ema_decay: 0.997

  pixel_opt:
    type: L1Loss
    loss_weight: 1.0
    reduction: mean

  perceptual_opt:
    type: LPIPSLoss
    loss_weight: 1.0
    use_input_norm: true
    range_norm: true

  gan_opt:
    type: GANLoss
    gan_type: hinge
    loss_weight: !!float 1.0 # adaptive_weighting

  use_adaptive_weight: true

  net_g_start_iter: 0
  net_d_iters: 1
  net_d_start_iter: 5001
  manual_seed: 0

# validation settings
val:
  val_freq: !!float 5e10 # no validation
  save_img: true

  metrics:
    psnr: # metric name, can be arbitrary
      type: calculate_psnr
      crop_border: 4
      test_y_channel: false

# 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: 29413

find_unused_parameters: true

	# general settings
	name: CodeFormer_stage3
	model_type: CodeFormerJointModel
	num_gpu: 8
	manual_seed: 0

	# dataset and data loader settings
	datasets:
	train:
	name: FFHQ
	type: FFHQBlindJointDataset
	dataroot_gt: datasets/ffhq/ffhq_512
	filename_tmpl: '{}'
	io_backend:
	type: disk

	in_size: 512
	gt_size: 512
	mean: [0.5, 0.5, 0.5]
	std: [0.5, 0.5, 0.5]
	use_hflip: true
	use_corrupt: true

	blur_kernel_size: 41
	use_motion_kernel: false
	motion_kernel_prob: 0.001
	kernel_list: ['iso', 'aniso']
	kernel_prob: [0.5, 0.5]
	# small degradation in stageIII
	blur_sigma: [0.1, 10]
	downsample_range: [1, 12]
	noise_range: [0, 15]
	jpeg_range: [60, 100]
	# large degradation in stageII
	blur_sigma_large: [1, 15]
	downsample_range_large: [4, 30]
	noise_range_large: [0, 20]
	jpeg_range_large: [30, 80]

	latent_gt_path: ~ # without pre-calculated latent code
	# latent_gt_path: './experiments/pretrained_models/VQGAN/latent_gt_code1024.pth'

	# data loader
	num_worker_per_gpu: 1
	batch_size_per_gpu: 3
	dataset_enlarge_ratio: 100
	prefetch_mode: ~

	# val:
	# name: CelebA-HQ-512
	# type: PairedImageDataset
	# dataroot_lq: datasets/faces/validation/lq
	# dataroot_gt: datasets/faces/validation/gt
	# io_backend:
	# type: disk
	# mean: [0.5, 0.5, 0.5]
	# std: [0.5, 0.5, 0.5]
	# scale: 1

	# network structures
	network_g:
	type: CodeFormer
	dim_embd: 512
	n_head: 8
	n_layers: 9
	codebook_size: 1024
	connect_list: ['32', '64', '128', '256']
	fix_modules: ['quantize','generator']

	network_vqgan: # this config is needed if no pre-calculated latent
	type: VQAutoEncoder
	img_size: 512
	nf: 64
	ch_mult: [1, 2, 2, 4, 4, 8]
	quantizer: 'nearest'
	codebook_size: 1024

	network_d:
	type: VQGANDiscriminator
	nc: 3
	ndf: 64
	n_layers: 4

	# path
	path:
	pretrain_network_g: './experiments/pretrained_models/CodeFormer_stage2/net_g_latest.pth' # pretrained G model in StageII
	param_key_g: params_ema
	strict_load_g: false
	pretrain_network_d: './experiments/pretrained_models/CodeFormer_stage2/net_d_latest.pth' # pretrained D model in StageII
	resume_state: ~

	# base_lr(4.5e-6)*bach_size(4)
	train:
	use_hq_feat_loss: true
	feat_loss_weight: 1.0
	cross_entropy_loss: true
	entropy_loss_weight: 0.5
	scale_adaptive_gan_weight: 0.1

	optim_g:
	type: Adam
	lr: !!float 5e-5
	weight_decay: 0
	betas: [0.9, 0.99]
	optim_d:
	type: Adam
	lr: !!float 5e-5
	weight_decay: 0
	betas: [0.9, 0.99]

	scheduler:
	type: CosineAnnealingRestartLR
	periods: [150000]
	restart_weights: [1]
	eta_min: !!float 2e-5


	total_iter: 150000

	warmup_iter: -1 # no warm up
	ema_decay: 0.997

	pixel_opt:
	type: L1Loss
	loss_weight: 1.0
	reduction: mean

	perceptual_opt:
	type: LPIPSLoss
	loss_weight: 1.0
	use_input_norm: true
	range_norm: true

	gan_opt:
	type: GANLoss
	gan_type: hinge
	loss_weight: !!float 1.0 # adaptive_weighting

	use_adaptive_weight: true

	net_g_start_iter: 0
	net_d_iters: 1
	net_d_start_iter: 5001
	manual_seed: 0

	# validation settings
	val:
	val_freq: !!float 5e10 # no validation
	save_img: true

	metrics:
	psnr: # metric name, can be arbitrary
	type: calculate_psnr
	crop_border: 4
	test_y_channel: false

	# 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: 29413

	find_unused_parameters: true