123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145 |
- # 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
|