k-diffusion

这是一个基于PyTorch的《阐明扩散型生成模型的设计空间》（Karras等人，2022）的实现，并增加了一些改进和额外功能，如更好的采样算法和基于transformer的扩散模型。

沙漏扩散transformer

k-diffusion包含一种新的模型类型image_transformer_v2，它结合了沙漏Transformer和DiT的理念。

要求

要使用这种新的模型类型，你需要安装自定义CUDA核心：

NATTEN用于层级结构低层的稀疏（邻域）注意力。模型类型也有一个移位窗口注意力版本，不需要自定义CUDA核心，但性能不如前者，且训练和推理速度较慢。
FlashAttention-2用于全局注意力。如果未安装，将回退到普通的PyTorch。

此外，请确保你的PyTorch安装能够使用torch.compile()。如果torch.compile()不可用，将回退到急切模式，但训练速度会更慢且内存使用更多。

使用方法

演示

要在Oxford Flowers数据集上训练256x256 RGB模型，无需安装自定义CUDA核心，请先安装Hugging Face Datasets：

pip install datasets

然后运行：

python train.py --config configs/config_oxford_flowers_shifted_window.json --name flowers_demo_001 --evaluate-n 0 --batch-size 32 --sample-n 36 --mixed-precision bf16

如果内存不足，尝试添加--checkpointing或减小批量大小。如果使用较旧的GPU（pre-Ampere），请去掉--mixed-precision bf16以使用FP32进行训练。不建议使用FP16训练。

如果你已安装并正常运行NATTEN（推荐），可以通过指定--config configs/config_oxford_flowers.json来使用邻域注意力而非移位窗口注意力进行训练。

配置文件

在配置文件的"model"键中：

将"type"键设置为"image_transformer_v2"。
基础patch大小由"patch_size"键设置，如"patch_size": [4, 4]。
层级结构每个级别的模型深度由"depths"配置键指定，如"depths": [2, 2, 4]。这将构建一个模型，第一级（4x4 patches）有两个transformer层，第二级（8x8 patches）有两个，最高级（16x16 patches）有四个，然后第二级再有两个，第一级再有两个。
层级结构每个级别的模型宽度由"widths"配置键指定，如"widths": [192, 384, 768]。宽度必须是注意力头维度的倍数。
层级结构每个级别的自注意力机制由"self_attns"配置键指定，如：
```
"self_attns": [
    {"type": "neighborhood", "d_head": 64, "kernel_size": 7},
    {"type": "neighborhood", "d_head": 64, "kernel_size": 7},
    {"type": "global", "d_head": 64},
]
```
如果未指定，除最高级外的所有级别都使用64维头和7x7核的邻域注意力。最高级使用64维头的全局注意力。因此，除最高级外，每个级别的token数量可以非常大。
作为无法使用NATTEN的备选方案，你也可以在层级结构的低级别使用移位窗口注意力来训练模型。移位窗口注意力的性能不如邻域注意力，训练和推理速度也较慢，但不需要自定义CUDA核心。指定方式如下：
```
"self_attns": [
    {"type": "shifted-window", "d_head": 64, "window_size": 8},
    {"type": "shifted-window", "d_head": 64, "window_size": 8},
    {"type": "global", "d_head": 64},
]
```
每个级别的窗口大小必须能够均匀地划分该级别的图像大小。使用一种注意力类型训练的模型必须经过微调才能使用不同类型。

推理

待完成：编写此部分

安装

k-diffusion可以通过PyPI安装（pip install k-diffusion），但这不包括训练和推理脚本，只包含其他人可以依赖的库代码。要运行训练和推理脚本，请克隆此仓库并运行pip install -e <仓库路径>。

训练

要训练模型：

$ ./train.py --config 配置文件 --name 运行名称

例如，要在MNIST上训练模型：

$ ./train.py --config configs/config_mnist_transformer.json --name 运行名称

配置文件允许你指定数据集类型。目前支持的类型有"imagefolder"（递归查找该文件夹及其子文件夹中的所有图像），"cifar10"（CIFAR-10）和"mnist"（MNIST）。也支持"huggingface"Hugging Face Datasets。

支持使用Hugging Face Accelerate进行多GPU和多节点训练。你可以通过运行以下命令配置Accelerate：

$ accelerate config

然后运行：

$ accelerate launch train.py --config 配置文件 --name 运行名称

改进/额外功能

k-diffusion支持高效的层级transformer模型类型。
k-diffusion支持Min-SNR损失加权的软版本，可在高分辨率下改善训练效果，所需超参数比Karras等人（2022）使用的损失加权更少。
k-diffusion为v-diffusion-pytorch、OpenAI diffusion和CompVis diffusion模型提供了封装器，允许它们与其采样器和ODE/SDE一起使用。
k-diffusion实现了DPM-Solver，在相同的函数评估次数下产生比Karras算法2更高质量的样本，并支持自适应步长控制。现在还实现了DPM-Solver++(2S)和(2M)，可在较少步数下提高质量。
k-diffusion支持从无条件扩散模型进行CLIP引导采样（参见sample_clip_guided.py）。
k-diffusion支持原生模型和所有封装模型的对数似然计算（非变分下界）。
k-diffusion可以在训练过程中计算相对于训练集的FID和KID。
k-diffusion可以在训练过程中计算梯度噪声尺度（1 / SNR），来自《大批量训练的经验模型》（https://arxiv.org/abs/1812.06162）。