ControlNet / model_config /mnist.yaml
YashNagraj75's picture
Add dataloader and add logging for training script
9774d79
raw
history blame contribute delete
746 Bytes
dataset_params:
im_path: 'data/mnist/train/images'
im_test_path: 'data/mnist/test/images'
canny_im_size: 28
diffusion_params:
num_timesteps : 1000
beta_start : 0.0001
beta_end : 0.02
model_params:
im_channels : 1
im_size : 28
hint_channels : 3
down_channels : [32, 64, 128, 256]
mid_channels : [256, 256, 128]
down_sample : [True, True, False]
time_emb_dim : 128
num_down_layers : 2
num_mid_layers : 2
num_up_layers : 2
num_heads : 4
train_params:
task_name: 'mnist'
batch_size: 64
num_epochs: 40
controlnet_epochs : 1
num_samples : 25
num_grid_rows : 5
save_epoch: 2
ddpm_lr: 0.0001
controlnet_lr: 0.0001
ddpm_ckpt_name: 'ddpm_ckpt.pth'
controlnet_ckpt_name: 'ddpm_controlnet_ckpt.pth'