目录

1. 安装
2. 数据集准备
3. 训练
4. 测试
5. 致谢
6. 联系方式

"潜空间流匹配"的官方PyTorch实现

<div align="center"> <a href="https://quandao10.github.io/" target="_blank">Quan&nbsp;Dao</a> &emsp; <b>&middot;</b> &emsp; <a href="https://hao-pt.github.io/" target="_blank">Hao&nbsp;Phung</a> &emsp; <b>&middot;</b> &emsp; <a href="https://tbng.github.io/" target="_blank">Binh&nbsp;Nguyen</a> &emsp; <b>&middot;</b> &emsp; <a href="https://sites.google.com/site/anhttranusc/" target="_blank">Anh&nbsp;Tran</a> <br> <br> <a href="https://www.vinai.io/">VinAI Research</a> <br> <br> &emsp;&emsp;<a href="https://vinairesearch.github.io/LFM/">[项目主页]</a> &emsp;&emsp; <a href="https://arxiv.org/abs/2307.08698">[论文]</a> &emsp;&emsp; </div> <br> <div align="center"> <img width="1000" alt="teaser" src="https://yellow-cdn.veclightyear.com/835a84d5/61530812-9c17-4c55-8cb5-66df953ce5f7.svg"/> </div>

摘要：流匹配是一种新兴的生成模型训练框架，相比基于扩散的模型，它展现出令人印象深刻的经验性能，同时训练难度相对较低。尽管具有这些优势，先前的方法仍然面临着像素空间中现成求解器的高昂计算成本和大量函数评估的挑战。此外，尽管基于潜空间的生成方法近年来取得了巨大成功，但在这一领域仍未得到充分探索。在本研究中，我们提出在预训练自编码器的潜空间中应用流匹配，这提高了高分辨率图像合成的计算效率和可扩展性。这使得在有限的计算资源下进行流匹配训练成为可能，同时保持其质量和灵活性。此外，我们的工作是首次将各种条件整合到流匹配中用于条件生成任务的开创性贡献，包括标签条件下的图像生成、图像修复和语义到图像的生成。通过大量实验，我们的方法在CelebA-HQ、FFHQ、LSUN Church & Bedroom和ImageNet等各种数据集上展示了其在定量和定性结果方面的有效性。我们还提供了重构潜流分布与真实数据分布之间Wasserstein-2距离的理论控制，表明它受到潜流匹配目标的上界限制。

模型架构和实验结果的详细信息可以在我们的论文中找到：

@article{dao2023lfm,
    author    = {Quan Dao and Hao Phung and Binh Nguyen and Anh Tran},
    title     = {Flow Matching in Latent Space},
    journal   = {arXiv preprint arXiv:2307.08698},
    year      = {2023}
}

请在使用本代码库产生发表结果或将其纳入其他软件时引用我们的论文。

安装

本实现使用Python 3.10和PyTorch 1.13.1/2.0.0。 请安装所需的库：

pip install -r requirements.txt

数据集准备

对于CelebA HQ 256、FFHQ 256和LSUN，请参考NVAE的说明。

对于更高分辨率的数据集（CelebA HQ 512和1024），请参考WaveDiff的文档。

对于ImageNet数据集，请直接从官方网站下载。

训练

所有训练脚本都封装在run.sh中。只需注释/取消注释相关命令并运行bash run.sh即可。

测试

采样

运行run_test.sh / run_test_cls.sh，并带上相应的参数文件。

bash run_test.sh <参数文件路径>

只需要1个GPU。

<details> <summary>这些参数指定如下：</summary>

MODEL_TYPE=DiT-L/2
EPOCH_ID=475
DATASET=celeba_256
EXP=celeb_f8_dit
METHOD=dopri5
STEPS=0
USE_ORIGIN_ADM=False
IMG_SIZE=256

</details>

参数文件和检查点提供如下：

<table> <tr> <th>实验</th> <th>参数</th> <th>FID</th> <th>检查点</th> </tr> <tr> <td> celeb_f8_dit </td> <td><a href="test_args/celeb256_dit.txt"> test_args/celeb256_dit.txt</a></td> <td>5.26</td> <td><a href="https://drive.google.com/drive/folders/1tbd1t0Yt3ix1v_OCGWJ7xyeubhCi99ql?usp=share_link">model_475.pth</a></td> </tr> <tr> <td> ffhq_f8_dit </td> <td><a href="test_args/ffhq_dit.txt"> test_args/ffhq_dit.txt</a></td> <td>4.55</td> <td><a href="https://drive.google.com/drive/folders/1jn6xHlaQ72hKk9RtJKo5lvr7SvYMCobU?usp=share_link">model_475.pth</a></td> </tr> <tr> <td> bed_f8_dit </td> <td><a href="test_args/bed_dit.txt"> test_args/bed_dit.txt</a></td> <td>4.92</td> <td><a href="https://drive.google.com/drive/folders/1o1uDrTAPIENHRh56CdVdGiEHGNqKcaC8?usp=share_link">model_550.pth</a></td> </tr> <tr> <td> church_f8_dit </td> <td><a href="test_args/church_dit.txt"> test_args/church_dit.txt</a></td> <td>5.54</td> <td><a href="https://drive.google.com/drive/folders/15ONlqM2eNbA91j7BikWPQG_6RH80NUwz?usp=share_link">model_575.pth</a></td> </tr> <tr> <td> imnet_f8_ditb2 </td> <td><a href="test_args/imnet_dit.txt"> test_args/imnet_dit.txt</a></td> <td>4.46</td> <td><a href="https://drive.google.com/file/d/1SUc9M85WOfi4CZfRJkpf2tDu3g9dC3YI/view?usp=sharing">model_875.pth</a></td> </tr> <tr> <td> celeb512_f8_adm </td> <td><a href="test_args/celeb256_adm.txt"> test_args/celeb512_adm.txt</a></td> <td>6.35</td> <td><a href="https://drive.google.com/drive/folders/1lWE9hCqzZ2Q1mS2BmTsA3nYWB_T25wqV?usp=share_link">model_575.pth</a></td> </tr> <tr> <td> celeba_f8_adm </td> <td><a href="test_args/celeb256_adm.txt"> test_args/celeb256_adm.txt</a></td> <td>5.82</td> <td>---</td> </tr> <tr> <td> ffhq_f8_adm </td> <td><a href="test_args/ffhq_adm.txt"> test_args/ffhq_adm.txt</a></td> <td>5.82</td> <td>---</td> </tr> <tr> <td> bed_f8_adm </td> <td><a href="test_args/bed_adm.txt"> test_args/bed_adm.txt</a></td> <td>7.05</td> <td>---</td> </tr> <tr> <td> church_f8_adm </td> <td><a href="test_args/church_adm.txt"> test_args/church_adm.txt</a></td> <td>7.7</td> <td>---</td> </tr> <tr> <td> imnet_f8_adm </td> <td><a href="test_args/imnet_adm.txt"> test_args/imnet_adm.txt</a></td> <td>8.58</td> <td>---</td> </tr> </table>

请将下载的预训练模型放在 saved_info/latent_flow/<DATASET>/<EXP> 目录中，其中 <DATASET> 在 bash_scripts/run.sh 中定义。

<details> <summary>实用工具</summary>

要测量时间，请在脚本中添加 --measure_time。

要计算自适应求解器（默认：dopri5）的函数评估次数，请在脚本中添加 --compute_nfe。

要使用固定步长求解器（如 euler 和 heun），请添加 --use_karras_samplers 并更改以下两个参数：

METHOD=heun
STEPS=50

</details>

评估

要评估 FID 分数，请从此处下载预计算的统计数据，并将其放入 pytorch_fid 目录。

然后运行 bash run_test_ddp.sh 进行无条件生成，运行 bash run_test_cls_ddp.sh 进行条件生成。默认情况下，支持使用 8 个 GPU 进行多 GPU 采样以加快计算速度。

<details> <summary>为新数据集计算统计数据</summary>

为此目的提供了 pytorch_fid/compute_dataset_stat.py。

python pytorch_fid/compute_dataset_stat.py \
  --dataset <dataset> --datadir <path_to_data> \
  --image_size <image_size> --save_path <path_to_save>

</details>

致谢

我们的代码来自不同的来源：EDM、DiT、ADM、CD、Flow Matching in 100 LOC by François Rozet 和 WaveDiff。我们非常感谢这些公开可用的研究和开发资源。

联系方式

如果您遇到任何问题，请在此存储库中提出问题或发送电子邮件至 v.quandm7@vinai.io / tienhaophung@gmail.com。