Spaces:
Running on Zero
Running on Zero
File size: 2,688 Bytes
e5a560a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | 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'
|