cifar10.yaml

# @package _global_

# to execute this experiment run:
# python train.py experiment=generation/diffusion/train/cifar10

defaults:
  - override /data: cifar10.yaml
  - override /model: diffusion/diffusion_module.yaml
  - override /callbacks: generation_diffusionn.yaml
  - override /trainer: ddp.yaml
  - override /logger: wandb.yaml

# all parameters below will be merged with parameters from default configurations set above
# this allows you to overwrite only specified parameters

seed: 12345

task_name: generation/diffusion/train

trainer:
  max_epochs: 2000
  devices: 1
  check_val_every_n_epoch: 5
  gradient_clip_val: 1.0

data:
  batch_size: 128 # for 1 GPU 16G
  num_workers: 10

model:
  net:
    img_dims:
      - 3
      - ${data.image_size}
      - ${data.image_size}
    denoise_net:
      base_channels: 128
      n_layer_blocks: 2
      channel_multipliers: [1, 2, 2, 2]
      attention: "Attn" # for cifar10 dataset
      n_attention_heads: null
      n_attention_layers: null
      attention_levels: [1] # resolution 16
      drop_rate: 0.1

logger:                                                                                                                                                                                                                                                                                                                                         
  wandb:
    name: ${now:%Y-%m-%d}_${now:%H-%M-%S}
    project: generation
    tags: ${tags}
    group: diffusion
    id: null # pass correct id to resume experiment! (if ckpt_path is not null)

ckpt_path: null # simply provide checkpoint path to resume training

callbacks:
  metrics:
    metric_list:
      - fid
      - is