Spaces:
Running on Zero
Running on Zero
| name = 'gmmlp_k64_checkerboard' | |
| model = dict( | |
| type='Diffusion2D', | |
| diffusion=dict( | |
| type='GMFlow', | |
| denoising=dict( | |
| type='GMFlowMLP2DDenoiser', | |
| num_gaussians=64), | |
| flow_loss=dict( | |
| type='GMFlowNLLLoss', | |
| log_cfgs=dict(type='quartile', prefix_name='loss_trans', total_timesteps=1000), | |
| data_info=dict( | |
| pred_means='means', | |
| target='x_t_low', | |
| pred_logstds='logstds', | |
| pred_logweights='logweights'), | |
| rescale_mode='constant', | |
| rescale_cfg=dict(scale=2.0)), | |
| num_timesteps=1000, | |
| timestep_sampler=dict(type='ContinuousTimeStepSampler', shift=1.0, logit_normal_enable=False), | |
| denoising_mean_mode='U'), | |
| diffusion_use_ema=True, | |
| ) | |
| save_interval = 20000 | |
| must_save_interval = 40000 # interval to save regardless of max_keep_ckpts | |
| work_dir = f'work_dirs/{name}' | |
| train_cfg = dict( | |
| trans_ratio=0.9, | |
| ) | |
| test_cfg = dict( | |
| output_mode='sample', | |
| sampler='FlowSDE', | |
| num_timesteps=8, | |
| order=2 | |
| ) | |
| optimizer = { | |
| 'diffusion': dict(type='AdamW', lr=2e-4, weight_decay=0.0) | |
| } | |
| data = dict( | |
| workers_per_gpu=8, | |
| train=dict(type='CheckerboardData', scale=4), | |
| train_dataloader=dict(samples_per_gpu=4096), | |
| val=dict(type='CheckerboardData', scale=4), | |
| val_dataloader=dict(samples_per_gpu=4096), | |
| test_dataloader=dict(samples_per_gpu=4096), | |
| persistent_workers=True, | |
| prefetch_factor=1024) | |
| lr_config = dict( | |
| policy='fixed', | |
| warmup='linear', | |
| warmup_iters=1000, | |
| warmup_ratio=0.001) | |
| checkpoint_config = dict( | |
| interval=save_interval, | |
| must_save_interval=must_save_interval, | |
| by_epoch=False, | |
| max_keep_ckpts=1, | |
| out_dir='checkpoints/') | |
| evaluation = [] | |
| total_iters = 200000 | |
| log_config = dict( | |
| interval=1000, | |
| hooks=[ | |
| dict(type='TextLoggerHook'), | |
| dict(type='TensorboardLoggerHook'), | |
| ]) | |
| # yapf:enable | |
| custom_hooks = [ | |
| dict( | |
| type='ExponentialMovingAverageHookMod', | |
| module_keys=('diffusion_ema', ), | |
| interp_mode='lerp', | |
| interval=1, | |
| start_iter=0, | |
| momentum_policy='karras', | |
| momentum_cfg=dict(gamma=7.0), | |
| priority='VERY_HIGH'), | |
| ] | |
| # use dynamic runner | |
| runner = dict( | |
| type='DynamicIterBasedRunnerMod', | |
| pass_training_status=True, | |
| ckpt_trainable_only=True, | |
| ckpt_fp16_ema=True, | |
| gc_interval=20) | |
| dist_params = dict(backend='nccl') | |
| log_level = 'INFO' | |
| load_from = None | |
| resume_from = f'checkpoints/{name}/latest.pth' # resume by default | |
| workflow = [('train', save_interval)] | |
| module_wrapper = 'ddp' | |
| cudnn_benchmark = True | |
| mp_start_method = 'fork' | |