修正流

这是ICLR2023聚焦论文的官方实现

流动直接且快速：学习生成和传输数据的修正流

作者：来自德克萨斯大学奥斯汀分校的Xingchao Liu, Chengyue Gong, Qiang Liu

InstaFlow

修正流可以应用于稳定扩散,使其成为一步生成器。详见此处

简介

修正流是一种新颖的方法,用于学习两个分布$\pi_0$和$\pi_1$之间的传输映射,通过连接样本之间的直线路径并学习ODE模型来实现。

然后,通过重新流动操作,我们迭代地拉直ODE轨迹,最终实现一步生成,比GAN具有更高的多样性,比快速扩散模型具有更好的FID。

这里可以找到介绍性网站,主要思想如下图所示:

修正流可以应用于生成建模和无监督域转移,如下图所示:

要更全面地了解其理论性质及与最优传输的关系,请参阅修正流:保持边际的最优传输方法

交互式Colab笔记本

我们提供交互式教程和Colab笔记本,引导您完成修正流的整个流程。我们提供两个版本,采用不同的速度模型,神经网络版本和非参数版本

图像生成

图像生成的代码在./ImageGeneration中。首先运行以下命令

cd ./ImageGeneration

依赖项

以下说明已在Lambda Labs的"1x A100 (40 GB SXM4)"实例上测试,即gpu_1x_a100_sxm4和Ubuntu 20.04.2(驱动版本:495.29.05,CUDA版本:11.5,CuDNN版本:8.1.0)。我们建议使用Anaconda。

运行以下命令安装依赖项:

conda create -n rectflow python=3.8
conda activate rectflow
pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install tensorflow==2.9.0 tensorflow-probability==0.12.2 tensorflow-gan==2.0.0 tensorflow-datasets==4.6.0
pip install jax==0.3.4 jaxlib==0.3.2
pip install numpy==1.21.6 ninja==1.11.1 matplotlib==3.7.0 ml_collections==0.1.1

训练1-修正流

运行以下命令从头开始训练1-修正流

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_ddpmpp.py --eval_folder eval --mode train --workdir ./logs/1_rectified_flow

--config 此运行的配置文件。
--eval_folder 训练期间每次评估生成的图像和其他文件将存储在./workdir/eval_folder中。在此命令中,为./logs/1_rectified_flow/eval/
---mode main.py的模式选择。从train、eval和reflow中选择。

采样和评估

我们遵循Score SDE中的评估流程。您可以下载cifar10_stats.npz并将其保存到assets/stats/。然后运行

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_ddpmpp.py --eval_folder eval --mode eval --workdir ./logs/1_rectified_flow --config.eval.enable_sampling  --config.eval.batch_size 1024 --config.eval.num_samples 50000 --config.eval.begin_ckpt 2

这使用1024的批量大小采样50000张图像,从checkpoint-2.pth开始,并计算FID和IS。

使用1-修正流生成数据对$(Z_0, Z_1)$

要准备重新流动的数据,运行以下命令

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_generate_data.py  --eval_folder eval --mode reflow --workdir ./logs/tmp --config.reflow.last_flow_ckpt "./logs/1_rectified_flow/checkpoints/checkpoint-10.pth" --config.reflow.data_root "./assets/reflow_data/1_rectified_flow/" --config.reflow.total_number_of_samples 100000 --config.seed 0

--config.reflow.last_flow_ckpt 用于数据生成的检查点。
--config.reflow.data_root 您希望保存生成对的位置。$(Z_0, Z_1)$对将保存到./data_root/seed/
--config.reflow.total_number_of_samples 您想生成的对的总数
--config.seed 随机种子。更改随机种子以并行执行数据生成。

对于CIFAR10,我们建议为重新流动生成至少1M对。

使用生成的数据对重新流动以获得2-修正流

生成数据对后,运行以下命令训练2-修正流

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_train.py  --eval_folder eval --mode reflow --workdir ./logs/2_rectified_flow --config.reflow.last_flow_ckpt "./logs/1_rectified_flow/checkpoints/checkpoint-10.pth" --config.reflow.data_root "./assets/reflow_data/1_rectified_flow/"

此命令使用上一步生成的数据对微调1-修正流的检查点,并将2-修正流的日志保存到./logs/2_rectified_flow。 2-修正流在使用一步生成$z_1=z_0 + v(z_0, 0)$时应该有更好的性能,如下图所示:

要使用步骤N=1进行评估,运行

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_ddpmpp.py --eval_folder eval --mode eval --workdir ./logs/2_rectified_flow --config.sampling.use_ode_sampler "euler" --config.sampling.sample_N 1 --config.eval.enable_sampling  --config.eval.batch_size 1024 --config.eval.num_samples 50000 --config.eval.begin_ckpt 2

其中sample_N指采样步骤的数量。

我们可以通过蒸馏进一步提高2-修正流在一步生成中的质量。

蒸馏以获得一步2-修正流

在蒸馏之前,我们需要来自2-修正流的新数据对。这可以简单地通过以下方式完成

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_generate_data.py  --eval_folder eval --mode reflow --workdir ./logs/tmp --config.reflow.last_flow_ckpt "./logs/2_rectified_flow/checkpoints/checkpoint-10.pth" --config.reflow.data_root "./assets/reflow_data/2_rectified_flow/" --config.reflow.total_number_of_samples 100000 --config.seed 0

然后我们可以通过以下方式蒸馏2-修正流

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_distill_k=1.py  --eval_folder eval --mode reflow --workdir ./logs/2_rectified_flow_k=1_distill --config.reflow.last_flow_ckpt "./logs/2_rectified_flow/checkpoints/checkpoint-10.pth" --config.reflow.data_root "./assets/reflow_data/2_rectified_flow/"

蒸馏以获得k步2-修正流(k>1)

同样,我们可以通过以下方式蒸馏2-修正流以进行k步生成(k>1)

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_distill_k_g_1.py  --eval_folder eval --mode reflow --config.reflow.reflow_t_schedule 2 --workdir ./logs/2_rectified_flow_k=2_distill --config.reflow.last_flow_ckpt "./logs/2_rectified_flow/checkpoints/checkpoint-10.pth" --config.reflow.data_root "./assets/reflow_data/2_rectified_flow/"

这里,我们以k=2为例。更改--config.reflow.reflow_t_schedule以适应不同的k。

使用在线数据生成进行重新流动和蒸馏

为了节省存储空间并简化重新流动中的训练流程,数据对也可以在训练期间由教师模型生成,这需要更长的训练时间。要使用在线数据生成执行重新流动,运行

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_train_online.py --eval_folder eval --mode reflow --workdir ./logs/2_rectified_flow --config.reflow.last_flow_ckpt "./logs/1_rectified_flow/checkpoints/checkpoint-10.pth"

然后，要使用在线数据生成蒸馏1-Rectified Flow，运行以下命令：

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_reflow_distill_k=1_online.py  --eval_folder eval --mode reflow --workdir ./logs/1_rectified_flow_k=1_distill --config.reflow.last_flow_ckpt "./logs/1_rectified_flow/checkpoints/checkpoint-10.pth"

高分辨率生成

我们提供了在四个256×256数据集（LSUN卧室、LSUN教堂、CelebA-HQ和AFHQ-猫）上进行高分辨率生成的代码和预训练检查点。

以CelebA-HQ为例。要在CelebA-HQ上训练1-rectified flow，运行：

python ./main.py --config ./configs/rectified_flow/celeba_hq_pytorch_rf_gaussian.py --eval_folder eval --mode train --workdir ./logs/celebahq

要从预训练的rectified flow中采样图像，运行：

python ./main.py --config ./configs/rectified_flow/celeba_hq_pytorch_rf_gaussian.py --eval_folder eval --mode eval --workdir ./logs/celebahq --config.eval.enable_figures_only --config.eval.begin_ckpt 10 --config.eval.end_ckpt 10 --config.training.data_dir YOUR_DATA_DIR

生成的图像将保存在./logs/celebahq/eval/ckpt/figs中。

预训练检查点

以使用预训练检查点为例，从CIFAR10 1-Rectified Flow下载checkpoint_8.pth，将其放入./logs/1_rectified_flow/checkpoints/，然后运行：

python ./main.py --config ./configs/rectified_flow/cifar10_rf_gaussian_ddpmpp.py --eval_folder eval --mode eval --workdir ./logs/1_rectified_flow --config.eval.enable_sampling  --config.eval.batch_size 1024 --config.eval.num_samples 50000 --config.eval.begin_ckpt 8

预训练检查点列表如下：

引用

如果您使用了我们的代码或者我们的工作与您的研究相关，请引用我们的论文：

@article{liu2022flow,
  title={Flow straight and fast: Learning to generate and transfer data with rectified flow},
  author={Liu, Xingchao and Gong, Chengyue and Liu, Qiang},
  journal={arXiv preprint arXiv:2209.03003},
  year={2022}
}