Pytorch implementation of the NeurIPS 2022 paper Poisson Flow Generative Models,
by Yilun Xu*, Ziming Liu*, Max Tegmark, Tommi S. Jaakkola
We propose a new Poisson flow generative model (PFGM) that maps a uniform distribution on a high-dimensional hemisphere into any data distribution. We interpret the data points as electrical charges on the
Experimentally, PFGM achieves current state-of-the-art performance among the normalizing flow models on CIFAR-10, with an Inception score of 9.68 and a FID score of 2.48. It also performs on par with the state-of-the-art SDE approaches (e.g., score-based SDEs or Diffusion models) while offering 10x to 20x acceleration on image generation tasks. Additionally, PFGM appears more tolerant of estimation errors on a weaker network architecture and robust to the step size in the Euler method, and capable of scale-up to higher resolution datasets.
Acknowledgement: Our implementation heavily relies on the repo https://github.com/yang-song/score_sde_pytorch.
Run the following to install a subset of necessary python packages for our code
pip install -r requirements.txt
Train and evaluate our models through main.py
.
python3 main.py:
--config: Training configuration.
--eval_folder: The folder name for storing evaluation results
(default: 'eval')
--mode: <train|eval>: Running mode: train or eval
--workdir: Working directory
For example, to train a new PFGM w/ DDPM++ model on CIFAR-10 dataset, one could execute
python3 main.py --config ./configs/poisson/cifar10_ddpmpp.py --mode train \
--workdir poisson_ddpmpp
-
config
is the path to the config file. The prescribed config files are provided inconfigs/
. They are formatted according toml_collections
and should be quite self-explanatory.Naming conventions of config files: the path of a config file is a combination of the following dimensions:
- Method: One of
poisson
(:star2:PFGM),ve
,vp
,sub_vp
- dataset: One of
cifar10
,celeba
,celebahq
,celebahq_256
,ffhq_256
,celebahq
,ffhq
. - model: One of
ncsn
,ncsnv2
,ncsnpp
,ddpm
,ddpmpp
. - continuous: train the model with continuously sampled time steps.
๐Important Notes : We use a large batch (e.g. current
training.batch_size=4096
for CIFAR-10, ~25G GPU memory usage) to calculate the Poisson field for each mini-batch samples (e.g.training.small_batch_size=128
for CIFAR-10). To adjust GPU memory cost, please modify thetraining.batch_size
parameter in the config files. - Method: One of
-
workdir
is the path that stores all artifacts of one experiment, like checkpoints, samples, and evaluation results. -
eval_folder
is the name of a subfolder inworkdir
that stores all artifacts of the evaluation process, like meta checkpoints for pre-emption prevention, image samples, and numpy dumps of quantitative results. -
mode
is either "train" or "eval". When set to "train", it starts the training of a new model, or resumes the training of an old model if its meta-checkpoints (for resuming running after pre-emption in a cloud environment) exist inworkdir/checkpoints-meta
. -
Below are the list of evalutation command-line flags:
--config.eval.enable_sampling
: Generate samples and evaluate sample quality, measured by FID and Inception score.--config.eval.enable_bpd
: Compute log-likelihoods--config.eval.dataset=train/test
: Indicate whether to compute the likelihoods on the training or test dataset.--config.eval.enable_interpolate
: Image Interpolation--config.eval.enable_rescale
: Temperature scaling
Please place the pretrained checkpoints under the directory workdir/checkpoints
, e.g., cifar10_ddpmpp/checkpoints
. To generate and evaluate the samples of the PFGM w/ DDPM++ model, you could execute:
python3 main.py --config ./configs/poisson/cifar10_ddpmpp.py --mode eval \
--workdir cifar10_ddpmpp --config.eval.enable_sampling
All checkpoints are provided in this Google drive folder.
Checkpoint path | Invertible? | IS | FID | NFE |
---|---|---|---|---|
poisson/cifar10_ddpmpp/ |
โ๏ธ | 9.62 | 2.54 | ~110 |
poisson/cifar10_ddpmpp_deep/ |
โ๏ธ | 9.68 | 2.48 | ~110 |
poisson/bedroom_ddpmpp/ |
โ๏ธ | - | 13.66 | ~122 |
poisson/celeba_ddpmpp/ (TODO) |
Please find the statistics for FID scores in the following links:
CIFAR-10, CelebA 64, LSUN bedroom 256
-
๐Important : We use a large batch (e.g. current
training.batch_size=4096
for CIFAR-10, ~25G GPU memory usage) to calculate the Poisson field for each mini-batch samples (e.g.training.small_batch_size=128
for CIFAR-10). To adjust GPU memory cost, please modify thetraining.batch_size
parameter in the config files. -
The prior distribution on the
$z=z_{max}$ hyperplane is a long-tail distribution. We recommend clipping the sample norm by the hyper-parameterssampling.upper_norm
-
To use the Euler forward method for solving PFGM ODE, please uncomment the Euler method code block in config files.
-
TODO