File size: 2,107 Bytes
3ce19a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | # 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](https://github.com/POSTECH-CVLab/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):
```bash
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
```bash
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:
```bash
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
```
|