total_iters: 800000
output_dir: output_dir

model:
  name: StyleGAN2Model
  generator:
    name: StyleGANv2Generator
    size: 256
    style_dim: 512
    n_mlp: 8
  discriminator:
    name: StyleGANv2Discriminator
    size: 256
  gan_criterion:
    name: GANLoss
    gan_mode: logistic
    loss_weight: !!float 1
  # r1 regularization for discriminator
  r1_reg_weight: 10.
  # path length regularization for generator
  path_batch_shrink: 2.
  path_reg_weight: 2.
  params:
    gen_iters: 4
    disc_iters: 16
  max_eval_steps: 50000

export_model:
  - {name: 'gen', inputs_num: 2}

dataset:
  train:
    name: SingleDataset
    dataroot: data/ffhq/images256x256/
    num_workers: 3
    batch_size: 3
    preprocess:
      - name: LoadImageFromFile
        key: A
      - name: Transforms
        input_keys: [A]
        pipeline:
          - name: RandomHorizontalFlip
          - name: Transpose
          - name: Normalize
            mean: [127.5, 127.5, 127.5]
            std: [127.5, 127.5, 127.5]
  test:
    name: SingleDataset
    dataroot: data/ffhq/images256x256/
    num_workers: 3
    batch_size: 3
    preprocess:
      - name: LoadImageFromFile
        key: A
      - name: Transforms
        input_keys: [A]
        pipeline:
          - name: Transpose
          - name: Normalize
            mean: [127.5, 127.5, 127.5]
            std: [127.5, 127.5, 127.5]

lr_scheduler:
  name: MultiStepDecay
  learning_rate: 0.002
  milestones: [600000]
  gamma: 0.5

optimizer:
  optimG:
    name: Adam
    beta1: 0.0
    beta2: 0.792
    net_names:
      - gen
  optimD:
    name: Adam
    net_names:
      - disc
    beta1: 0.0
    beta2: 0.9317647058823529


log_config:
  interval: 50
  visiual_interval: 500

snapshot_config:
  interval: 5000

validate:
  interval: 50000
  save_imig: False
  metrics:
    fid: # metric name, can be arbitrary
      name: FID
      batch_size: 4