StudioGAN + Recursive Token Mapper (RTM)
StyleGAN2-ADA + RTM training and evaluation pipeline used for the AFHQ-v1 (512x512) and CIFAR-10 (32x32) StyleGAN baselines in the paper.
The whole PyTorch-StudioGAN
codebase is vendored under src/ (see LICENSE-StudioGAN and
LICENSE-NVIDIA for the upstream licenses);
Setup
Use the same Python environment as the rest of this repository (PyTorch ≥ 2.1, CUDA ≥ 12, plus StudioGAN's extra deps below):
pip install h5py kornia==0.6.5 prdc seaborn imageio scikit-learn
Datasets are loaded with StudioGAN's stock loader. For AFHQ-v1 set
DATA_DIR to a folder with train/{cat,dog,wild} and valid/{cat,dog,wild}
sub-directories. For CIFAR-10 set DATA_DIR to the standard CIFAR-10 folder.
Train
DATA_DIR=/path/to/AFHQ bash scripts/train_afhq.sh
DATA_DIR=/path/to/cifar-10 bash scripts/train_cifar10.sh
Both scripts read RTM hyper-parameters from the YAML config (see the
STYLEGAN.use_rtm_mapper block in each config). The headline numbers in the
paper use the values shipped in those YAMLs:
| Knob | AFHQ (512) | CIFAR-10 (32) |
|---|---|---|
rtm_num_tokens |
4 (a 2x2 grid) | 4 (a 2x2 grid) |
rtm_H_cycles |
4 | 4 |
rtm_L_cycles |
1 | 1 |
rtm_H_layers |
2 | 2 |
rtm_L_layers |
2 | 2 |
rtm_hidden_size |
256 | 256 |
rtm_refinement_steps |
4 | 4 |
use_rtm_equalized |
True | True |
rtm_lr_multiplier |
0.01 | 0.01 |
Evaluate
Point CKPT_DIR at the StudioGAN-format checkpoints/<exp>/ folder produced
by the training run:
CKPT_DIR=runs/afhq_rtm/checkpoints/<exp> DATA_DIR=/path/to/AFHQ bash scripts/eval_afhq.sh
CKPT_DIR=runs/cifar10_rtm/checkpoints/<exp> DATA_DIR=/path/to/cifar-10 bash scripts/eval_cifar10.sh