Skip to content

learninginvision/ProjectedGAN-CLC

Repository files navigation

📖 Closed Loop Training for Projected GAN

This repository contains code for our paper

Closed-Loop Training for Projected GAN
Jiangwei Zhao, Liang Zhang, Lili Pan, Hongliang Li
IEEE Signal Processing Letters

Abstract:Projected GAN, a pre-trained GAN, has been found to perform well in generating images with only a few training samples. However, it struggles with extended training, which may lead to decreased performance over time. This is because the pre-trained discriminator consistently surpasses the generator, creating an unstable training environment. In this work, we propose a solution to this issue by introducing closed-loop control (CLC) into the dynamics of Projected GAN, stabilizing training, and improving generation performance. Our proposed method consistently reduces the Fréchet Inception Distance (FID) of the previous methods; for example, it reduces the FID of Projected GAN by 4.31 on the Obama dataset. Our finding is fundamental and can be used in other pre-trained GANs.

🔧 Dependencies and Installation

Dependencies

  • 64-bit Python 3.8
  • PyTorch 1.9.0 (or later). See https://pytorch.org for PyTorch install instructions.

Installation

First, you can clone this repo using the command:

git clone https://github.com/learninginvision/ProjectedGAN-CLC

Then, you can create a virtual environment using conda, as follows:

conda env create -f environment.yaml
conda activate pg-clc

🗺 Data Preparation

For a quick start, you can download the few-shot datasets provided by the authors of FastGAN. You can download them here. To prepare the dataset at the respective resolution, run for example

python dataset_tool.py --source=./data/pokemon --dest=./data/pokemon256.zip \
  --resolution=256x256 --transform=center-crop

You can get the datasets we used in our paper at their respective websites: AFHQ, Landscape.

💻 Training

Training your own PG-CLC on Pokemon using 2 GPUs:

python train.py --outdir=./training-runs/ --cfg=fastgan --data=./data/pokemon256.zip \
  --gpus=2 --batch=64 --mirror=1 --snap=50 --batch-gpu=16 --kimg=10000

--batch specifies the overall batch size, --batch-gpu specifies the batch size per GPU.

We use a lightweight version of FastGAN (--cfg=fastgan_lite). This backbone trains fast regarding wallclock time and yields better results on small datasets like Pokemon. Samples and metrics are saved in outdir. To monitor the training progress, you can inspect fid50k_full.json or run tensorboard in training-runs.

🔧 Config and Pretrained Models

You can change the config of clc on train.py#L240-L243 We provide the following pretrained models (pass the url as PATH_TO_NETWORK_PKL):

Dataset Loss Weight Queue Factor FID PATH
Pokemon 0.1 100 25.04 https://drive.google.com/file/d/18-678PSsr4sYX28qtIkdkOd3TtdpKCWf
Art-Paint 0.05 200 26.91 https://drive.google.com/file/d/1if_qohz0PYtSzuSlL72nE71oATxuSmVT
Flowers 0.05 200 12.82 https://drive.google.com/file/d/1B844ooziyOhk3dGbS389XWujIPjTpYbN
landscapes 0.05 100 6.55 https://drive.google.com/file/d/1RpDg4vRPgD6UXajzmWDNSuyxkS2F_pwK
Obama 0.05 100 20.12 https://drive.google.com/file/d/1A0SbqW3xvHMfWVs_Pp7nUs8Ih5Uj9aYL

📊 Quality Metrics

Per default, train.py tracks FID50k during training. To calculate metrics for a specific network snapshot, run

python calc_metrics.py --metrics=fid50k_full --network=PATH_TO_NETWORK_PKL

To see the available metrics, run

python calc_metrics.py --help

📈 Results

Results on Pokemon, Art-Paint, Flowers, Landscape and Obama.

📜Citation

@ARTICLE{10334000,
  author={Zhao, Jiangwei and Zhang, Liang and Pan, Lili and Li, Hongliang},
  journal={IEEE Signal Processing Letters}, 
  title={Closed-Loop Training for Projected GAN}, 
  year={2024},
  volume={31},
  number={},
  pages={106-110},
  doi={10.1109/LSP.2023.3337711}}

📜Acknowledgments

Our codebase build and extends the awesome StyleGAN2-ADA repo, ProjectedGAN repo and StyleGAN3 repo

Furthermore, we use parts of the code of FastGAN and MiDas.

About

No description, website, or topics provided.

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published