摘要

本仓库包含在ImageNet 64x64上训练的DDPM、DDIM和无分类器引导模型。更多信息请见下文。

我还写了一篇文章来解释其背后的算法。

当前功能

本仓库具有以下扩散模型功能:

普通DDPM
改进的DDPM,使用余弦调度器和方差预测
用于更快推理的DDIM
无分类器引导以提高图像质量

我写了一篇文章详细解释了每个部分,而不是在这里逐一介绍。

环境配置

首先,使用以下命令行从此仓库下载数据:

git clone https://github.com/gmongaras/Diffusion_models_from_scratch.git
cd Diffusion_models_from_scratch/

(可选)如果你不想更改当前环境,可以先创建一个虚拟环境:

pip install virtualenv
python -m venv MyEnv/

激活虚拟环境:https://docs.python.org/3/library/venv.html#how-venvs-work

Windows: MyEnv\Scripts\activate.bat

Linux: source MyEnv/bin/activate

在运行任何脚本之前,请确保下载正确的包和包版本。你可以运行以下命令来升级pip并安装必要的包版本:

pip install pip -U
pip install -U requirements.txt

注意:如果要进行训练,PyTorch应该安装启用cuda的版本,如果要生成图像,可能也需要cuda,但不是必需的。下载的cuda版本可能与所需的版本不同。cuda版本及其下载方法可在以下链接找到:

https://pytorch.org/get-started/locally/

现在环境应该正确配置好了。

下载预训练模型

预训练模型说明

我提供了几个可下载的预训练模型,具有不同的模型架构类型。根据U-Net块的构造,有5种模型类型:

res ➜ conv ➜ clsAtn ➜ chnAtn (Res-Conv)
res ➜ clsAtn ➜ chnAtn (Res)
res ➜ res ➜ clsAtn ➜ chnAtn (Res-Res)
res ➜ res ➜ clsAtn ➜ atn ➜ chnAtn (Res-Res-Atn)
res ➜ clsAtn ➜ chnAtn 通道数为192 (Res Large)

上述符号来自训练模型部分的blk_types参数。

除非另有说明,每个模型都使用以下参数进行训练:

图像分辨率:64x64
初始嵌入通道:128
通道乘数 - 1
U-net块数 - 3
时间步数 - 1000
VLB权重Lambda - 0.001
Beta调度器 - 余弦
批量大小 - 128 (跨8个GPU,因此为1024)
梯度累积步数 - 1
步数(注意:这不是epoch,一步是模型的单次梯度更新) - 600,000
学习率 - 3*10^-4 = 0.0003
时间嵌入维度大小 - 512
类别嵌入维度大小 - 512
无分类器引导的空类概率 - 0.2
注意力分辨率 - 16

以下是一些训练注意事项:

我使用了8个并行GPU,每个GPU的批量大小为128。因此,总批量大小为8*128 = 1024。
我训练每个模型总共600,000步。请注意,这不是epoch数,而是模型更新次数。如果你希望从预训练检查点继续训练,这些模型似乎可以训练更长时间,因为即使在600,000步时,FID值看起来仍在下降。
训练较小的模型(res-conv、res、res-res)花了6-7天,而较大的模型在8个A100上花了约8天。

选择模型

要选择模型,我建议查看结果。FID分数越低,模型输出质量越好。根据结果,最佳模型是:

res-res-atn:
- 358个epoch (358e)
- 450000步 (450000s)
- 该模型的文件包括:
  - 模型文件: model_358e_450000s.pkl
  - 模型文件: model_params_358e_450000s.json
  - 优化器文件: optim_358e_450000s.pkl
res-res:
- 438个epoch (438e)
- 550000步 (550000s)
- 该模型的文件包括:
  - 模型文件: model_438e_550000s.pkl
  - 模型元数据: model_params_438e_550000s.json
  - 优化器: optim_438e_550000s.pkl

下载模型

选择模型后,你可以从以下链接下载模型:

Google Drive链接

从检查点继续训练需要下载模型的三个文件:

模型.pkl文件(例如:model_438e_550000s.pkl)
模型元数据.json文件(例如:model_params_438e_550000s.json)
优化器.pkl文件(例如:optim_438e_550000s.pkl)

推理/生成只需下载模型的两个文件:

模型.pkl文件(例如:model_438e_550000s.pkl)
模型元数据.json文件(例如:model_params_438e_550000s.json)

将这些文件放入models/目录,以便在训练/生成时轻松加载。

下载训练数据

可以从以下链接下载ImageNet数据: https://image-net.org/download-images.php

要获取数据,你必须先申请访问权限并获得批准,然后才能下载ImageNet数据。我在ImageNet 64x64上训练了我的模型

下载后,你应该将Imagenet64_train_part1.zip和Imagenet64_train_part2.zip放在data/目录中。

zip文件位于正确的目录后,运行以下脚本将数据加载到必要的格式:

python data/loadImagenet64.py

如果你希望在训练前将数据加载到内存中,请运行以下脚本。否则,数据将根据需要从磁盘中提取。

python data/make_massive_tensor.py

下载所有数据后,目录结构应如下所示:目录结构

目录结构

如果您下载了预训练模型和训练数据，您的目录结构应如下所示：

.
├── data
│   ├── Imagenet64
|   |   ├── 0.pkl
|   |   ├── ...
|   |   ├── metadata.pkl
│   ├── Imagenet64_train_part1.zip
│   ├── Imagenet64_train_part2.zip
│   ├── README.md
│   ├── archive.zip
│   ├── loadImagenet64.py
│   ├── make_massive_tensor.py
├── eval
|   ├── __init__.py
|   ├── compute_FID.py
|   ├── compute_imagenet_stats.py
|   ├── compute_model_stats.py
|   ├── compute_model_stats_multiple.py
├── models
|   ├── README.md
|   ├── [模型参数名].json
|   ├── [模型名].pkl
├── src
|   ├── blocks
|   |   ├── BigGAN_Res.py
|   |   ├── BigGAN_ResDown.py
|   |   ├── BigGAN_ResUp.py
|   |   ├── ConditionalBatchNorm2D.py
|   |   ├── Efficient_Channel_Attention.py
|   |   ├── Multihead_Attn.py
|   |   ├── Non_local.py
|   |   ├── Non_local_MH.py
|   |   ├── PositionalEncoding.py
|   |   ├── Spatial_Channel_Attention.py
|   |   ├── __init__.py
|   |   ├── clsAttn.py
|   |   ├── convNext.py
|   |   ├── resBlock.py
|   |   ├── wideResNet.py
|   ├── helpers
|   |   ├── PixelCNN_PP_helper_functions.py
|   |   ├── PixelCNN_PP_loss.py
|   |   ├── image_rescale.py
|   |   ├── multi_gpu_helpers.py
|   ├── models
|   |   ├── PixelCNN.py
|   |   ├── PixelCNN_PP.py
|   |   ├── U_Net.py
|   |   ├── Variance_Scheduler.py
|   |   ├── diff_model.py
|   ├── CustomDataset.py
|   ├── __init__.py
|   ├── infer.py
|   ├── model_trainer.py
|   ├── train.py
├── tests
|   ├── BigGAN_Res_test.py
|   ├── U_Net_test.py
|   ├── __init__.py
|   ├── diff_model_noise_test.py
├── .gitattributes
├── .gitignore
├── README.md

训练模型

在训练模型之前，请确保您已经设置好环境并下载了数据

完成上述步骤后，您可以从此仓库的根目录运行训练脚本，如下所示：

torchrun --nproc_per_node=[GPU数量] src/train.py --[参数]

[GPU数量]替换为您希望用于并行训练的GPU数量
[参数]替换为下列任何参数

例如：

torchrun --nproc_per_node=8 src/train.py --blk_types res,res,clsAtn,chnAtn --batchSize 32

上述示例使用以下参数运行代码：

在8个并行GPU上运行。
每个U-net块由res->res->clsAtn->chnAtn顺序块组成。
每个GPU的batchSize为32，因此总批量大小为8*32 = 256。

torchrun --nproc_per_node=1 src/train.py --loadModel True --loadDir models/models_res --loadFile model_479e_600000s.pkl --optimFile optim_479e_600000s.pkl --loadDefFile model_params_479e_600000s.json --gradAccSteps 2

上述示例加载预训练模型的检查点：

使用1个GPU
加载模型文件model_479e_600000s.pkl
加载优化器文件optim_479e_600000s.pkl
加载模型元数据文件model_params_479e_600000s.json
使用2个梯度累积步骤

脚本的参数如下：

数据参数

inCh [3] - 输入数据的通道数。
data_path [data/Imagenet64] - ImageNet 64x64数据集的路径。
load_into_mem [True] - True表示将所有ImageNet数据加载到内存中，False表示根据需要从磁盘加载数据。

模型参数

embCh [128] - U-net顶层的通道数。注意，这在每个U-net层都会按2^(chMult*layer)进行缩放。
chMult [1] - 在每个U-net层，通道应该按什么比例乘以？每一层有embCh2^(chMultlayer)个通道。
num_layers [3] - U-net层数。值为3表示深度为3，意味着U-net中有3个下采样层和3个上采样层。
blk_types [res,clsAtn,chnAtn] - 残差块应如何构建？（res、conv、clsAtn、atn和/或chnAtn的列表）
- res：普通残差块。
- conv：ConvNext块。
- clsAtn：包含额外类别信息的注意力块（Q、K是类别特征）。
- atn：隐藏特征的注意力ViT块。
- chnAtn：特征通道上的高效轻量级注意力块。
- 例如：res,res,conv,clsAtn,chnAtn
T [1000] - 扩散过程中的时间步数。
beta_sched [cosine] - 扩散过程中使用的噪声调度器。可以是linear或cosine。
t_dim [512] - 时间信息编码向量的维度。
c_dim [512] - 类别信息编码向量的维度。注意：使用-1表示无类别信息。
atn_resolution [16] - 注意力块（atn）的分辨率。该分辨率将图像分割成该分辨率的块作为向量。例如：分辨率16会创建16x16的块并将它们展平为特征向量。

训练参数

Lambda [0.001] - 模型中方差和均值损失之间的权重项。
batchSize [128] - 单个GPU上的批量大小。如果使用多个GPU，此批量大小将乘以GPU数量。
gradAccSteps [1] - 将batchSize分解成的步骤数。不是一次加载整个批次到内存中进行一次大的更新，而是将batchSize分解成batchSize//gradAccSteps大小的小批次，以便能够装入内存。从数学上讲，更新效果与单个批次更新相同，但更新被分散到较小的更新中以适应内存。较高的值需要更多时间，但使用更少的内存。
device [gpu] - 将模型放置在哪个设备上。使用"gpu"将其放置在一个或多个Cuda设备上，或使用"cpu"将其放置在CPU上。CPU应仅用于测试。
epochs [1000000] - 训练的轮数。
lr [0.0003] - 模型学习率。
p_uncond [0.2] - 用于无分类器引导的空类训练概率。注意，好的值是0.1或0.2。（仅当c_dim不为None时使用）
use_importance [False] - True表示对t值使用重要性采样，False表示使用均匀采样。

保存参数

saveDir [models/] - 保存模型检查点的目录。注意将保存三个文件：模型.pkl文件、模型元数据.json文件和用于训练重新加载的优化器.pkl文件。
numSaveSteps [10000] - 保存新模型检查点的步数。这不是轮数，而是模型更新的次数。注意将保存三个文件：模型.pkl文件、模型元数据.json文件和用于训练重新加载的优化器.pkl文件。

模型加载参数

loadModel [False] - 设为 True 以从检查点加载预训练模型。设为 False 则使用随机初始化的模型。请注意,成功重启需要三个模型文件:模型 .pkl 文件、模型元数据 .json 文件和优化器 .pkl 文件。
loadDir [models/] - 要加载的模型文件所在目录。
loadFile [""] - 要加载的模型 .pkl 文件名。格式类似于:model_10e_100s.pkl
optimFile [""] - 要加载的优化器 .pkl 文件名。格式类似于:optim_10e_100s.pkl
loadDefFile [""] - 要加载的模型元数据 .json 文件名。格式类似于:model_params_10e_100s.json

数据加载参数

reshapeType [""] - 如果数据大小不均匀,使用此参数将图像向上或向下重塑为 2 的幂次方,或者保持不变("up"、"down"、"")

使用预训练模型生成图像

在训练模型之前,请确保你已经设置好环境并下载了预训练模型

完成上述步骤后,你可以从该仓库的根目录运行以下脚本:

python -m src.infer --loadDir [模型目录位置] --loadFile [.pkl模型文件名] --loadDefFile [.json模型参数文件名] --[其他参数]

例如,如果我下载了 models_res_res_atn 模型的 model_358e_450000s 文件,并想在 CPU 上使用步长为 20,我可以在命令行中使用以下命令:

python -m src.infer --loadDir models/models_res_res_atn --loadFile model_358e_450000s.pkl --loadDefFile model_params_358e_450000s.json --device cpu --step_size 20

推理脚本的参数如下:

必需参数:

loadDir - 要加载的模型所在位置。
loadFile - 要加载的 .pkl 模型文件名。例如:model_358e_450000s.pkl
loadDefFile - 要加载的 .json 模型文件名。例如:model_params_358e_450000s.pkl

生成参数

step_size [10] - 生成时的步长。对于在 1000 步上训练的模型,步长为 10 时需要 100 步来生成。较小的值生成速度更快,但图像质量较低。
DDIM_scale [0] - 必须 >= 0。当该值为 0 时,使用 DDIM。当该值为 1 时,使用 DDPM。低标量在高步长时表现更好,高标量在低步长时表现更好。
device ["gpu"] - 放置模型的设备。使用 "gpu" 或 "cpu"。
guidance [4] - 分类器引导尺度,必须 >= 0。值越高,图像质量越好,但图像多样性越低。
class_label [0] - 从 0 开始索引的类别值。使用 -1 表示随机类别,使用 >= 0 的其他类别值表示其他类别。对于 ImageNet,类别值范围从 0 到 999,可以在 data/class_information.txt 中找到。
corrected [False] - 设为 True 对生成进行限制,设为 False 则不限制生成。如果模型生成单一颜色的图像,可能需要将此标志设为 True。注意:这种限制通常在生成长序列(低步长)时需要。注意:使用更高的引导 w 时,校正通常会干扰生成。

输出参数

out_imgname ["fig.png"] - 保存输出图像的文件名。
out_gifname ["diffusion.gif"] - 保存输出动图的文件名。
gif_fps [10] - 输出动图的帧率。

注意:类别值和标签从零开始索引,可以在这个文档中找到。

计算预训练模型的 FID

一旦你训练好了模型,你可以使用这些脚本来评估它们。

注意:本节的所有脚本都位于 eval/ 目录中。

计算 FID 需要三个步骤:

1: 计算 ImageNet 数据的统计信息

在这一步,运行 compute_imagenet_stats.py 来计算 ImageNet 数据集的 FID。

python -m eval.compute_imagenet_stats

此脚本有以下参数:

archive_file1 - 第一个 ImageNet 64x64 zip 文件的路径。
archive_file2 - 第二个 ImageNet 64x64 zip 文件的路径。
batchSize - 用于并行生成统计信息的批量大小。
num_imgs - 从数据集中采样计算统计信息的图像数量。
device - 用于计算统计信息的设备。

2: 计算预训练模型的统计信息

这一步有两种选择。如果你想为单个预训练模型生成 FID,使用 compute_model_stats.py:

python -m eval.compute_model_stats

此脚本有以下参数(可以通过编辑文件进行访问):

model_dirname - 要计算统计信息的预训练模型目录。
model_filename - 预训练模型的文件名。
model_params_filename - 预训练模型元数据的文件名。
device - 运行模型推理的设备
gpu_num - 运行模型推理的 GPU 编号(如果只有 1 个 GPU,则使用 0)
num_fake_imgs - 在计算模型统计信息之前生成的图像数量。注意,不建议使用小于 10,000 的值,因为统计信息将不准确。
batchSize - 同时生成的图像批量大小。较高的值可以加快进程,但需要更多的 GPU 内存。
step_size - 扩散模型的步长(>= 1)。这个步长将生成过程减少 step_size 倍。如果模型需要 1000 步来生成单个图像,但步长为 4,那么生成一张图像将需要 1000/4 = 250 步。注意,更高的步长意味着更快的生成,但图像质量也更低。
DDIM_scale - 使用 0 表示 DDIM,使用 1 表示 DDPM(>= 0)。有关此的更多信息位于训练部分。
corrected - True 表示对生成进行限制,False 表示不限制生成。如果模型生成全黑或全白图像,则可能需要这个限制。低步长通常需要限制。
file_path - 保存统计文件的路径。
mean_filename - 保存均值统计信息的文件名。
var_filename - 保存方差统计信息的文件名。

如果你想为多个模型生成 FID 并有多个 GPU 可用,你可以并行化这个过程。compute_model_stats_multiple.py 允许这种并行化,可以使用以下命令运行:

python -m eval.compute_model_stats_multiple

注意:每个列表中的项目数量最多应等于你想使用的 GPU 数量。

此脚本有以下可以在脚本文件内更改的参数:

dir_name - 加载所有模型文件的目录。
model_filenames - 要计算 FID 的模型文件名列表。
model_params_filenames - 模型元数据文件名列表。
gpu_nums - 放置每个模型的 GPU 编号。GPU 编号的索引对应文件名的索引。
step_size - 扩散模型的步长(>= 1)。这个步长将生成过程减少 step_size 倍。如果模型需要 1000 步来生成单个图像,但步长为 4,那么生成一张图像将需要 1000/4 = 250 步。注意,更高的步长意味着更快的生成,但图像质量也更低。
DDIM_scale - 使用 0 表示 DDIM,使用 1 表示 DDPM(>= 0)。有关此的更多信息位于训练部分。
corrected - True 表示对生成进行限制,False 表示不限制生成。如果模型生成全黑或全白图像,则可能需要这个限制。低步长通常需要限制。
num_fake_imgs - 在计算模型统计数据之前生成的图像数量。建议不要低于10,000张,否则统计数据可能不准确。
batchSize - 一次生成的图像批量大小。较大的值可以加快处理速度,但需要更多GPU内存。
file_path - 保存所有模型统计数据的目录。
mean_filenames - 保存每个模型平均统计数据的文件名。
var_filenames - 保存每个模型方差统计数据的文件名。