ml_mdm - 套娃扩散模型

ml_mdm是一个Python包,用于高效训练高质量的文本到图像扩散模型 — 由Luke Carlson、Jiatao Gu、Shuangfei Zhai和Navdeep Jaitly推出。

---

<div align="center">

本软件项目伴随研究论文Matryoshka Diffusion Models。

Jiatao Gu, Shuangfei Zhai, Yizhe Zhang, Josh Susskind, Navdeep Jaitly

[论文] [引用]

</div>

目录

章节	描述
简介	套娃扩散模型的简要概述
安装	开始使用ml_mdm训练模型和生成样本
预训练模型	下载我们预训练模型(64, 356, 1024)的链接
网页演示	使用我们的网页界面生成图像
代码库结构	Python模块概述
概念	核心概念和设计原则
教程	在CC12m上逐步训练MDM模型

简介

扩散模型是生成高质量图像和视频的事实标准方法,但由于计算和优化挑战,学习高维模型仍然是一项艰巨的任务。

ml_mdm是一个端到端的高分辨率图像和视频合成框架 — 它以我们的技术命名:套娃扩散模型。

值得注意的是,我们可以在高达1024x1024像素的分辨率下训练单个像素空间模型,使用仅包含1200万张图像的CC12M数据集展示了强大的零样本泛化能力。

安装

默认安装依赖项在pyproject.toml中定义,选择这些依赖项是为了让您即使在只有CPU的机器上也能安装此库。

用户已经在Python 3.9、3.10和cuda_12、cuda-11.8上运行过这个代码库

> pip install -e .

开发人员还应该使用pre-commit install设置pre-commit。

运行测试用例

> pytest   # 将运行所有测试用例 - 包括需要GPU的测试
> pytest  -m "not gpu"  # 运行只需要CPU的测试用例

预训练模型

我们已经上传了模型检查点到:

• https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr64/vis_model.pth
• https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr256/vis_model.pth
• https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr1024/vis_model.pth

注意:我们发布的模型是在从Flickr收集的5000万对文本-图像对上训练的。在这个仓库中,我们提供了下载CC12M的脚本,以及在CC12M数据上训练等效模型的配置。

您可以随意下载模型或继续往下阅读以训练自己的模型。一旦预训练模型在本地下载完成,您可以在我们的网页演示中使用它,将其作为参数传递给训练、采样等。

export ASSET_PATH=https://docs-assets.developer.apple.com/ml-research/models/mdm

curl $ASSET_PATH/flickr64/vis_model.pth --output vis_model_64x64.pth
curl $ASSET_PATH/flickr256/vis_model.pth --output vis_model_256x256.pth
curl $ASSET_PATH/flickr1024/vis_model.pth --output vis_model_1024x1024.pth

网页演示

下载检查点后,您可以使用以下命令运行自己的网页演示实例:

torchrun --standalone --nproc_per_node=1  ml_mdm/clis/generate_sample.py --port $YOUR_PORT

代码库

模块	描述
ml_mdm.models	核心模型实现
ml_mdm.diffusion	模型管道,例如DDPM
ml_mdm.config	使用simple parsing将配置数据类与相关模型、管道和命令行接口连接
ml_mdm.clis	项目中的所有命令行工具,最相关的是train_parallel.py
tests/	单元测试和样本训练文件

概念

ml_mdm.models

在ml_mdm.models子模块中,我们开源了以下实现:

• U-Nets
• 嵌套U-Nets

ml_mdm.config

ml_mdm.config包含核心配置和命令行接口逻辑。这个代码库中的许多模型、命令行接口和函数都是通过传入一个dataclass对象来配置的。我们使用SimpleParsing动态创建命令行接口,并允许使用--config_path传入yaml config表示。

本质上,simple_parsing会将所有传入的命令行参数和yaml文件转换为干净的配置类,如ml_mdm.reader.ReaderConfig、ml_mdm.diffusion.DiffusionConfig。

教程

使用预训练检查点生成自己的图像

安装ml_mdm后,将这些检查点下载到仓库目录中。

curl https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr64/vis_model.pth --output vis_model_64x64.pth
curl https://docs-assets.developer.apple.com/ml-research/models/mdm/flickr256/vis_model.pth --output vis_model_256x256.pth

网页演示将使用相应的配置加载每个模型:

• vis_model_64x64.pth将使用configs/models/cc12m_64x64.yaml中的设置加载
• vis_model_256x256.pth将使用configs/models/cc12m_256x256.yaml中的设置加载
• vis_model_1024x1024.pth将使用configs/models/cc12m_1024x1024.yaml中的设置加载

在演示中,您可以更改各种设置并查看模型的内部结构。通过替换$YOUR_PORT设置您想使用的端口,然后运行:

torchrun --standalone --nproc_per_node=1  ml_mdm/clis/generate_sample.py --port $YOUR_PORT

在虚拟数据上训练

如果您只想在不下载大型数据集的情况下逐步完成训练模型和运行管道的过程,我们为您准备了一个最小示例。它使用tests/test_files/中的虚拟数据

随意尝试直接在命令行或通过编辑配置yaml文件来更改各种--args

torchrun --standalone --nproc_per_node=1 ml_mdm/clis/train_parallel.py \
 --file-list=tests/test_files/sample_training_0.tsv \
 --multinode=0 \
  --output-dir=outputs    --config_path configs/models/cc12m_64x64.yaml \
  -num_diffusion_steps=10 \
	--num-training-steps=10

您应该看到一个outputs/vis_model_000100.pth文件。现在让我们做一些更有意义的事:

让我们在CC12m上训练一个MDM模型

1. 数据准备:

(可选)使用此示例参数下载CC12m的前1000个文件

该脚本基于img2dataset的CC12M脚本。

curl https://storage.googleapis.com/conceptual_12m/cc12m.tsv | head -n 1000 > cc12m_index.tsv

# 向文件添加标题
sed -i '1s/^/url\tcaption\n/'  cc12m_index.tsv

注意:如果您想要所有的cc12m,请从调用中删除| head -n 1000

然后准备并分割为训练/验证集

此脚本需要img2dataset,运行pip install '.[data_prep]'或只运行pip install img2dataset

python3 -m ml_mdm.clis.scrape_cc12m \
  --cc12m_index cc12m_index.tsv \
  --cc12m_local_dir cc12m_download

运行此命令后,您将看到以下文件:

training.0.tsv # 训练索引文件
validation.tsv # 验证索引文件
cc12m_download/
   00000.parquet  00000.tar  00000.tsv  00000_stats.json  validation.tsv
   00001.parquet ....

2. 训练

现在我们有了训练文件,我们可以选择一个模型配置并传递任何额外的训练参数:

# 修改torchrun参数以适应您的GPU设置
torchrun --standalone --nproc_per_node=8 ml_mdm/clis/train_parallel.py \
  --file-list=training_0.tsv \
  --multinode=0 --output-dir=/mnt/data/outputs \
  --config_path configs/models/cc12m_64x64.yaml \
  --num-training-steps=100   --warmup-steps 10

注意:configs/models/cc12m_64x64.yaml包含更多参数,查看它以了解更多详情。 如果你已下载预训练模型，可以将 --pretrained-vision-file 参数设置为指向磁盘上的模型位置

训练完成后，你会在 --output-dir 参数定义的文件夹中找到模型：

2024-07-22:17:58:46,649 INFO     [model_ema.py:33] 保存 EMA 模型文件：/mnt/data/outputs/vis_model_000100.pth
2024-07-22:17:58:47,448 INFO     [unet.py:794] 保存模型文件：/mnt/data/outputs/vis_model_noema_000100.pth

3. 从模型中采样

现在我们有了训练好的模型，我们可以从扩散模型中生成样本：

torchrun --standalone --nproc_per_node=1 ml_mdm/clis/generate_batch.py \
  --config_path configs/models/cc12m_64x64.yaml \
  --min-examples 3 --test-file-list validation.tsv \
  --sample-image-size 64 --model-file /mnt/data/outputs/vis_model_000100.pth

如果你想跳过训练步骤，可以将 --model-file 更新为指向我们的预训练模型之一

数据集存储

对于长期存储，你可以选择将数据上传到 s3://{your_bucket}/datasets/{datasetname}/*.[tar,tsv]。

然后更新 configs/datasets/cc12m.yaml 以指向你的 s3 路径。

# configs/datasets/cc12m.yaml
train:
  files:
    - s3://mlx/datasets/cc12m-64x64/images_00.*.tsv
eval:
  files:
    - s3://mlx/datasets/cc12m-64x64/validation.tsv

# configs/datasets/reader_config.yaml
reader_config:
  append_eos: true
  bucket: ${your_bucket} # 添加你的 s3 存储桶
  endpoint_url: None # boto 将自动推断端点

然后你可以使用我们的数据集下载助手：

python -m ml_mdm.clis.download_tar_from_index \
  --dataset-config-file configs/datasets/cc12m.yaml \
  --subset train --download_tar

python -m ml_mdm.clis.download_tar_from_index \
  --dataset-config-file configs/datasets/cc12m.yaml \
  --subset eval --download_tar

S3 数据集选择

查看 configs/datasets/cc12m.yaml。

代码允许提供多个正则表达式。请记住，这些正则表达式不是通配符 -- 它们是来自 Python re 库的正则表达式。所以如果你想在训练中只使用 WebImage 中 1000 个 tar 文件中的 100 个，你可以这样做：

train:
  files:
    - s3://mlx/datasets/example-dataset-100M_64px/example-dataset-100M-00[0-1]..-[0-9]*-of-01000.tsv
eval:
  files:
    - s3://mlx/datasets/example-dataset-100M_64px/validation.tsv

你也可以混合搭配文件。所以如果你想合并 CC12m 和 ImageNet，你可以创建一个新的 yaml 文件，内容如下：

train:
  files:
    - s3://mlx/datasets/imagenet-64px/imagenet-train-000??-of-00100.tsv
    - s3://mlx/datasets/cc12m-64x64/images_00.*.tsv
eval:
  files:
    - s3://mlx/datasets/cc12m-64x64/validation.tsv

数据集结构

S3 存储桶包含一系列以这种格式的文件，查看 ml_mdm/clis/scrape_cc12m.py 以生成你自己的。

2023-04-01 01:31:30   36147200 images_00000.tar
2023-05-10 11:34:49    1108424 images_00000.tsv
2023-04-01 01:31:26   36454400 images_00001.tar
2023-05-10 11:34:49    1109588 images_00001.tsv
2023-04-01 01:31:53   36116480 images_00002.tar
...

这些文件的最小表示可以在 tests/test_files/ 中找到。

引用

如果你发现我们的工作有用，请考虑引用我们：

@misc{gu2023matryoshkadiffusionmodels,
      title={Matryoshka Diffusion Models},
      author={Jiatao Gu and Shuangfei Zhai and Yizhe Zhang and Josh Susskind and Navdeep Jaitly},
      year={2023},
      eprint={2310.15111},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2310.15111},
}