使用 Hugging Face Diffusers 训练 ControlNet

Hugging Face 已发布一份详细指南和训练脚本，位于 diffusers 库中，帮助开发者训练自己的 ControlNet 模型。ControlNet 是一种神经网络结构，通过向扩散模型添加额外条件（如姿态估计、深度图或草图），实现对生成过程的细粒度控制。

三步式 ControlNet 训练工作流

为 Stable Diffusion 训练自定义 ControlNet 包含三个主要阶段：规划条件、构建数据集以及执行训练。

1. 规划条件

第一步是定义任务所需的具体条件。这包括确定期望的控制机制，并判断是否已有模型能够将普通图像转换为该特定条件。例如，Hugging Face 团队的目标是创建一个面部关键点模型，以便 Stable Diffusion 能够遵循特定的面部表情或姿势。

2. 构建数据集

一个 ControlNet 数据集需要三个特定列：

• Ground Truth Image：目标图像（例如人脸）。
• Conditioning Image：表示条件的图像（例如可视化的面部关键点掩码）。
• Prompt：描述图像的文本标题。

在 “Uncanny Faces” 示例中，团队使用了 Microsoft 的 FaceSynthetics 数据集，包含 10 万张合成面孔。由于没有现成模型能够直接将人脸转换为数据集所需的关键点格式，团队采用了最先进的 SPIGA 模型提取 iBUG 格式的 68 点面部关键点，将这些关键点转换为插画式掩码，并使用 BLIP 进行字幕生成，为每张图像生成描述。

3. 训练模型

训练由 diffusers 示例中提供的 train_controlnet.py 脚本完成。团队使用单块 A100 GPU，训练 3 个 epoch，批大小为 4。

训练观察与过拟合

团队发现训练 3 个 epoch 会导致过拟合，模型开始忽略风格并忘记与真实人脸不同的概念（例如在提示 “cat” 或 “shrek” 时无法生成）。在约 1 个 epoch（约 25K 步）左右即可收敛，此时模型能够顺利跟随姿势而不出现过拟合。由于 FaceSynthetics 数据集由合成图像构成，最终模型生成的面孔呈现 “诡异” 的 3D 效果，而非写实。

技术实现与硬件优化

训练配置

train_controlnet.py 脚本使用多个关键参数来控制输出：

• pretrained_model_name_or_path：基础 Stable Diffusion 模型（示例中使用 v2-1-base，以获得更好的面部渲染）。
• learning_rate：示例中设为 1e-5，建议范围在 1e-4 到 2e-6 之间。
• resolution：条件图像和真实图像均设为 512x512。
• validation_steps：决定模型运行验证提示和图像的频率，以跟踪训练进度。

低显存 GPU 的 VRAM 优化

虽然主要示例使用 A100，但 diffusers 脚本支持多种优化，以适配显存更低的 GPU：

GPU 显存	所需优化 / 参数
16GB	train_batch_size=1、gradient_accumulation_steps=4、gradient_checkpointing，以及通过 bitsandbytes 启用的 use_8bit_adam。
12GB	在 16GB 优化的基础上再加 set_grads_to_none。
8GB	具体配置请参见 diffusers GitHub 训练指南中的说明。

通过将批大小设为 1 并使用 4 步梯度累积，用户可以在显存受限的情况下模拟 A100 训练中使用的有效批大小 4，同时显著降低内存占用。