Skip to content

Configuration guide

Everything in GAN-Engine is controlled by YAML files. Configuration files define the dataset, normalisation pipeline, architectures, loss weights, schedulers, logging, and runtime options. This page summarises each section so you can design experiments for any image modality.

File layout

Configuration files live in gan_engine/configs/. Copy an existing template and edit it, or create a new file and run python -m gan_engine.train --config path/to/your.yaml.

Config
├── Data
├── Normalisation
├── Model
│   ├── Generator
│   └── Discriminator
├── Training
│   ├── Losses
│   ├── Optimisers & schedulers
│   ├── Stability controls
│   └── Logging

Use this block to set seeds, output paths, and naming. task influences default callbacks (e.g. image logging).

Data

Define your own dataloader and provide it to the model. Then, you need to register it with the data/dataset_selector.py function and select it in the YAML config.

Data:
  dataset: "CV"
  batch_size: 8
  num_workers: 8

Normalisation

Choose any of the pre-inmplemented Normalization options or augment it with your own.

Normalisation: "CV"

Model

Model:
  Generator:
    model_type: rrdb
    in_channels: 1
    out_channels: 1
    n_blocks: 23
    n_channels: 64
    growth_channels: 32
    scale: 4
  Discriminator:
    model_type: esrgan
    in_channels: 1
    n_blocks: 7
    base_channels: 64
    feature_matching: true

Change model_type to switch architectures. Additional parameters depend on the chosen type (e.g. window_size for SwinIR variants, kernel_size for LKA). Setting in_channels greater than 3 is fully supported.

Training & losses

Training:
  precision: 16
  max_steps: 500_000
  accumulate_grad_batches: 2
  ema:
    enabled: true
    decay: 0.999
  Losses:
    l1_weight: 1.0
    perceptual_weight: 0.2
    perceptual_metric: lpips
    perceptual_channels: [0]  # choose specific bands
    sam_weight: 0.05
    adv_loss_beta: 1e-3
    adv_warmup:
      steps: 20_000
      mode: cosine
  Optimizers:
    generator:
      name: adam
      lr: 2e-4
      betas: [0.9, 0.99]
    discriminator:
      name: adam
      lr: 1e-4
      betas: [0.9, 0.99]
  Schedulers:
    generator:
      warmup_steps: 5_000
      warmup_type: linear
      plateau_patience: 10_000
      factor: 0.5
    discriminator:
      warmup_steps: 0
      plateau_patience: 15_000
      factor: 0.5
  Stability:
    pretrain_g_only: true
    g_pretrain_steps: 100_000
    d_update_interval: 1
    gradient_clip_val: 1.0

Highlights:

  • ema toggles exponential moving averages for the generator.
  • Losses accepts any combination; omit weights to disable terms. Upcoming releases add text/image alignment losses here.
  • adv_warmup ramps in adversarial pressure gradually.
  • Optimizers and Schedulers are defined per network. You can reference cosine annealing, OneCycle, or custom schedulers.

Logging & callbacks

Logging:
  logger: wandb  # wandb | tensorboard | csv
  project: gan-engine-experiments
  entity: gan-engine
  log_images_every_n_steps: 2_000

Pick a logger, control how often image panels are captured, and configure checkpoint/export behaviour. Additional callbacks include early stopping, LR logging, and evaluation hooks that run custom metrics.

Armed with these configuration knobs, you can adapt GAN-Engine to any dataset, architecture, or training philosophy.