train_unet_ddim.yaml

 1# train ae or vae
 2rand_seed: 2024
 3model:
 4  # can be AE, VAE and DDPM
 5  name: DDIM
 6  num_steps: 1000
 7  sample_num_steps: 100
 8  beta_schedule: consine
 9  eps_model:
10    # encoder config
11    time_channel: 128
12    encoder:
13      name: UNet
14      image_size: [28, 28]
15      in_channel: 1
16      patch_size: 1
17      down_channels: [32, 64]
18      middle_channels: [128, 128]
19      building_block: res
20      activation: silu
21      num_blocks: 2
22      normalization: group
23      num_norm_groups: 16
24      dropout: 0.1
25    # condition_embedding:
26    #   combine_condition: add
27    #   merge_timestep_and_condition: true
28    #   encoder:
29    #     name: OneHot
30    #     type: categories
31    #     out_channel: 128
32    #     num_classes: 10
33  # classifier_free_guidance:
34  #   condition_dropout: 0.2
35  #   guidance_weight: 2.0
36  ## 2D ddpm
37  eps_loss:
38    name: MSE
39loader:
40  dataset:
41    name: MNIST
42  data_path_list:
43    - /media/wlsdzyzl/DATA1/datasets/img/MNIST
44  batch_size: 128
45  num_workers: 8
46  shuffle: true
47  data_transforms:
48    - name: Resize
49      ## the value must be list
50      size: [28, 28]
51    - name: ToTensor
52    - name: Normalize
53      mean: [0.5]
54      std: [0.5]
55check_point_dir: path/to/checkpoint/MNIST/DDPM
56sampler:
57  num_sample_steps: 100
58  rand_seed: 2024
59  clipped: true
60  clip_range: [-1.0, 1.0]
61### parameter for optimizer
62optimizer:
63  name: AdamW
64  lr: 0.0003
65  weight_decay: 0.00000001
66### scheduler for learning rate
67lr_scheduler:
68  name: LinearLR
69  start_factor: 1.0
70  end_factor: 0.01
71max_epoch: 1000
72write_after_iters: 50
73save_after_epochs: 5