| # 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 |
| ``` |
|
|
