k-diffusion

这是 Elucidating the Design Space of Diffusion-Based Generative Models（Karras 等，2022）在 PyTorch 上的实现，并进行了增强和添加了额外功能，例如改进的采样算法以及基于 Transformer 的扩散模型。

Hourglass 扩散 Transformer

k-diffusion 包含一种新的模型类型 image_transformer_v2，它借鉴了 Hourglass Transformer 和 DiT 中的思想。

要求

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

• NATTEN 用于在层次结构的低层使用的稀疏（邻域）注意力机制。该模型类型还有一个 移位窗口注意力 版本，不需要自定义 CUDA 内核，但性能较差，训练和推理速度也较慢。

• FlashAttention-2 用于全局注意力机制。如果未安装，将回退到普通的 PyTorch 实现。

此外，您还需要确保您的 PyTorch 安装支持 torch.compile()。如果无法使用 torch.compile()，则会回退到急切模式，但这会导致训练速度变慢并占用更多内存。

使用方法

演示

要在不安装自定义 CUDA 内核的情况下，在 Oxford Flowers 数据集上训练一个 256x256 RGB 模型，请先安装 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（Ampere 之前），请省略 --mixed-precision bf16 以在 FP32 下进行训练。不建议使用 FP16 进行训练。

如果您已经安装并正确配置了 NATTEN（推荐），可以通过指定 --config configs/config_oxford_flowers.json 来使用邻域注意力代替移位窗口注意力进行训练。

配置文件

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

1. 将 "type" 键设置为 "image_transformer_v2"。

2. 基础补丁大小由 "patch_size" 键指定，例如 "patch_size": [4, 4]。

3. 层次结构中每一层的模型深度由 "depths" 配置键指定，例如 "depths": [2, 2, 4]。这将构建一个模型：第一层（4x4 补丁）有两个 Transformer 层，第二层（8x8 补丁）也有两个，最高层（16x16 补丁）有四个，随后在第二层再加两个，最后在第一层再加两个。

4. 层次结构中每一层的模型宽度由 "widths" 配置键指定，例如 "widths": [192, 384, 768]。宽度必须是注意力头维度的倍数。

5. 层次结构中每一层的自注意力机制由 "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 数量可能会非常大。

6. 作为备选方案，如果您或您的用户无法使用 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 CONFIG_FILE --name RUN_NAME

例如，要在 MNIST 数据集上训练一个模型：

$ ./train.py --config configs/config_mnist_transformer.json --name RUN_NAME

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

多 GPU 和多节点训练可通过 Hugging Face Accelerate 实现。您可以先运行以下命令来配置 Accelerate：

$ accelerate config

然后运行：

$ accelerate launch train.py --config CONFIG_FILE --name RUN_NAME

增强/附加功能

• k-diffusion 支持一种高效的分层 Transformer 模型类型。

• k-diffusion 支持 Min-SNR 损失加权 的软版本，用于在高分辨率下进行改进的训练，且所需的超参数比 Karras 等人（2022）使用的损失加权更少。

• k-diffusion 为 v-diffusion-pytorch、OpenAI diffusion 和 CompVis diffusion 模型提供了封装器，使其能够与 k-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 在训练过程中还可以根据 大型批量训练的经验模型（https://arxiv.org/abs/1812.06162），计算梯度噪声尺度（1 / SNR）。

待办事项

• 隐式扩散

k-diffusion 快速上手指南

k-diffusion 是一个基于 PyTorch 的扩散模型实现库，源自论文《Elucidating the Design Space of Diffusion-Based Generative Models》。它提供了高效的采样算法（如 DPM-Solver）、改进的损失函数以及新型的分层 Transformer 架构（Hourglass Diffusion Transformer）。

环境准备

系统要求

• 操作系统: Linux (推荐) 或 macOS
• GPU: NVIDIA GPU (支持 CUDA)，建议使用 Ampere 架构或更新版本以启用 bf16 混合精度训练。
• Python: 3.8+
• PyTorch: 需安装支持 torch.compile() 的版本以获得最佳性能。

前置依赖

若使用新型的 image_transformer_v2 模型（分层 Transformer），建议安装以下自定义 CUDA 内核以提升训练和推理速度：

1. NATTEN: 用于低层级的稀疏邻域注意力机制。
2. FlashAttention-2: 用于全局注意力机制（未安装时会自动回退到原生 PyTorch 实现）。

注意：若不安装 NATTEN，可使用“移位窗口注意力”（shifted window attention）模式，该模式无需自定义内核，但性能和速度略低。

安装步骤

1. 克隆仓库并安装

为了运行训练和推理脚本，建议从源码安装（可编辑模式）：

git clone https://github.com/crowsonkb/k-diffusion.git
cd k-diffusion
pip install -e .

国内加速提示：如果下载速度慢，可配置 pip 使用国内镜像源（如清华源）：

pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装可选依赖（针对高性能 Transformer 模型）

如果需要训练 image_transformer_v2 模型，请安装以下组件：

# 安装 Hugging Face Datasets (用于数据加载)
pip install datasets

# 安装 FlashAttention-2
pip install flash-attn --no-build-isolation

# 安装 NATTEN (需根据官方文档编译，此处为通用安装命令示例)
pip install natten

3. 配置多 GPU 环境（可选）

若需进行多卡或多节点训练，需配置 Hugging Face Accelerate：

accelerate config

基本使用

快速演示：训练一个简易模型

以下示例展示如何在 Oxford Flowers 数据集上训练一个 256x256 的 RGB 模型。此命令无需安装自定义 CUDA 内核（使用移位窗口注意力），适合快速验证环境。

1. 确保已安装 datasets 库：

   pip install datasets
   

2. 运行训练脚本：

   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
   

   参数说明与调整建议：

   • --mixed-precision bf16: 启用 bf16 混合精度加速。注意：如果您使用的是 Ampere 之前的旧款 GPU，请移除此参数以使用 FP32 训练（不建议使用 FP16）。
   • 显存不足时：添加 --checkpointing 参数或减小 --batch-size。

3. 使用高性能模式（可选）： 如果您已成功安装 NATTEN，可以使用邻域注意力（性能更好）进行训练：

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

配置文件关键点

若要自定义模型架构（在配置文件的 "model" 字段中）：

• 类型: 设置 "type": "image_transformer_v2"。
• 层级深度: 通过 "depths" 定义每层的 Transformer 层数（例如 [2, 2, 4] 表示沙漏型结构）。
• 注意力机制: 通过 "self_attns" 定义。
  • 推荐配置（需 NATTEN）：低层使用 "neighborhood"，顶层使用 "global"。
  • 兼容配置（无 NATTEN）：低层使用 "shifted-window"。

其他训练场景

支持的数据集类型包括本地文件夹 (imagefolder)、CIFAR-10、MNIST 以及 Hugging Face 数据集。

单卡训练示例 (MNIST):

./train.py --config configs/config_mnist_transformer.json --name mnist_run

多卡训练示例:

accelerate launch train.py --config configs/config_mnist_transformer.json --name mnist_run_multi