# general settings name: CodeFormer_colorization model_type: CodeFormerIdxModel num_gpu: 8 manual_seed: 0 # dataset and data loader settings datasets: train: name: FFHQ type: FFHQBlindDataset 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 # large degradation in stageII blur_kernel_size: 41 use_motion_kernel: false motion_kernel_prob: 0.001 kernel_list: ['iso', 'aniso'] kernel_prob: [0.5, 0.5] blur_sigma: [1, 15] downsample_range: [4, 30] noise_range: [0, 20] jpeg_range: [30, 80] # color jitter and gray color_jitter_prob: 0.3 color_jitter_shift: 20 color_jitter_pt_prob: 0.3 gray_prob: 0.01 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: 2 batch_size_per_gpu: 4 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'] vqgan_path: './experiments/pretrained_models/vqgan/vqgan_code1024.pth' # pretrained VQGAN 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 # path path: pretrain_network_g: ~ param_key_g: params_ema strict_load_g: false pretrain_network_d: ~ strict_load_d: true 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 fidelity_weight: 0 optim_g: type: Adam lr: !!float 1e-4 weight_decay: 0 betas: [0.9, 0.99] scheduler: type: MultiStepLR milestones: [400000, 450000] gamma: 0.5 total_iter: 500000 warmup_iter: -1 # no warm up ema_decay: 0.995 use_adaptive_weight: true net_g_start_iter: 0 net_d_iters: 1 net_d_start_iter: 0 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 1e4 use_tb_logger: true wandb: project: ~ resume_id: ~ # dist training settings dist_params: backend: nccl port: 29419 find_unused_parameters: true