seed: 8888
data:
  target: datasets.celeba.CelebA
  params:
    root: ~/data/CelebA/
    img_size: 64
    split: train
dataloader:
  num_workers: 16
  pin_memory: true
  prefetch_factor: 2
model:
  target: models.transformer.Transformer
  params:
    codebook_num: 512
    embed_dim: 512
    n_heads: 8
    n_layers: 12
    max_tokens: 256
vqvae:
  target: models.vqvae.VQVAE
  params:
    img_channels: 3
    hidden_dim: 256
    n_resblocks: 2
    codebook_num: 512
    codebook_dim: 64
    codebook_update: kmeans
    ema_decay: 0.99
  pretrained: ./runs/vqvae-kmeans-celeba/ckpt/step499999/model.pt
train:
  n_steps: 500000
  batch_size: 64
  print_freq: 100
  sample_freq: 1000
  save_freq: 10000
  optim:
    target: torch.optim.Adam
    params:
      lr: 0.0001