用扩散器训练你的控制网

简介

ControlNet 是一种神经网络结构，允许通过添加额外条件对扩散模型进行细粒度控制。迅速接管了开源扩散社区作者发布的 8 种不同条件来控制稳定扩散 v1-5，包括姿势估计、深度图、精明边缘、草图等。

在这篇博客文章中，我们将详细介绍我们如何训练Uncanny Faces模型的每一步——一个基于3D合成面孔的面部姿势模型。

开始训练控制网以实现稳定的扩散

训练你自己的控制网需要3个步骤

1. 规划您的状况：ControlNet足够灵活，可以驯服对许多任务的稳定扩散。预先训练的模型展示了广泛的条件，社区已经构建了其他条件，例如像素化调色板的条件。
2. 构建数据集：确定条件后，就该构建数据集了。为此，您可以从头开始构建数据集，也可以使用现有数据集的子集。数据集上需要三列来训练模型：真实图像、conditioning_image和提示。
3. 训练模型：数据集准备就绪后，就可以训练模型了。由于扩散器训练脚本，这是最简单的部分。您需要一个至少具有 8GB VRAM 的 GPU。

规划条件

为了计划你的状况，考虑两个问题

1. 我想使用哪种条件反射？
2. 是否有已经存在的模型可以将“常规”图像转换为我的状况？

在我们的示例中，我们考虑使用面部特征点调理。我们的理由是：
1.一般地标条件控制网运行良好。
2.面部特征点是一种足够普遍的技术，并且有多种模型可以计算常规图片上的面部特征点
3.驯服稳定扩散以遵循某个面部特征或模仿自己的面部表情可能会很有趣。

2. 构建数据集

所以我们决定做一个面部标志稳定扩散调理。因此，要准备数据集，我们需要：

地面实况图像：在这种情况下，是人脸的图像
条件反射图像：在本例中，可视化面部特征的图像
标题：描述正在使用的图像的标题

对于这个项目，我们决定使用微软的FaceSynthetics数据集：它是一个包含100K合成人脸的数据集。其他具有真实面孔的人脸研究数据集，例如Celeb-A HQ，FFHQ – 但我们决定在这个项目中使用合成人脸。

FaceSynthetics数据集听起来是一个很好的开始：它包含面部的真实图像，以iBUG 68面部特征点格式注释的面部特征点，以及面部的分割图像。

好极了。对吧？不幸的是，并非如此。还记得“规划你的条件”步骤中的第二个问题吗? 那就是我们应该建立模型，将常规图像转换为条件反射？事实证明，目前还没有已知的模型可以将人脸转换为这个数据集的带注释的地标格式。

所以我们决定走另一条路：

• 使用面部合成数据的基本真实人脸图像
• 使用已知模型，可以将任何面部图像转换为iBUG的68面部特征点格式（在我们的例子中，我们使用SOTA模型SPIGA）
• 使用自定义代码将面部特征点转换为漂亮的插图蒙版以用作conditioning_image
• 将其另存为拥抱面部数据集

在这里，您可以找到用于将 FaceSynthetics 数据集中的实况图像转换为插图蒙版并将其另存为拥抱面部数据集的代码。

现在，对于真实图像和数据集上的conditioning_image，我们缺少一个步骤：每个图像的标题。强烈建议执行此步骤，但您可以尝试空提示并报告结果。由于我们没有FaceSynthetics数据集的字幕，我们通过BLIP字幕运行它。您可以在此处检查用于为所有图像添加字幕的代码

就这样，我们到达了最终数据集！带有标题的面部合成 SPIGA 包含真实图像、分割和 FaceSynthetics 数据集的 100K 图像的标题。我们已准备好训练模型！

3.训练模型

准备好我们的数据集后，是时候训练模型了！尽管这应该是过程中最困难的部分，但使用扩散器训练脚本，结果证明这是最简单的。我们在LambdaLabs上以1.10美元/小时的价格租用了一架A100。

我们对模型进行了 3 个 epoch 的训练（这意味着向模型显示 3 次 100K 图像的批次）和 4 个批量大小（每个步骤向模型显示 4 个图像）。事实证明这是过度和过度拟合的（所以它忘记了有点真实面孔的概念，所以例如提示中的“史莱克”或“猫”不会变成怪物或猫，而是一个人，也开始忽略样式）。

只有 1 个纪元（所以在模型“看到”100K 图像之后），它已经收敛到跟随姿势而不是过度拟合。所以它奏效了，但是…当我们使用面部合成数据集时，该模型最终学习了不可思议的3D面部，而不是逼真的面部。考虑到我们使用的是合成人脸数据集而不是真实人脸数据集，这是有道理的，并且可以用于娱乐/模因目的。这是uncannyfaces_25K模型。

在这个交互式表中，您可以使用下面的表盘来查看模型经历了多少训练步骤，以及它如何影响训练过程。在大约15K步的时候，它已经开始学习姿势了。它大约成熟了25K步。在这里

我们是如何进行培训的

我们所要做的就是安装依赖项

pip install git+https://github.com/huggingface/diffusers.git transformers accelerate xformers==0.0.16 wandb
huggingface-cli login
wandb login 

然后运行train_controlnet.py代码

!accelerate launch train_controlnet.py 
 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base" 
 --output_dir="model_out" 
 --dataset_name=multimodalart/facesyntheticsspigacaptioned 
 --conditioning_image_column=spiga_seg 
 --image_column=image 
 --caption_column=image_caption 
 --resolution=512 
 --learning_rate=1e-5 
 --validation_image "./face_landmarks1.jpeg" "./face_landmarks2.jpeg" "./face_landmarks3.jpeg" 
 --validation_prompt "High-quality close-up dslr photo of man wearing a hat with trees in the background" "Girl smiling, professional dslr photograph, dark background, studio lights, high quality" "Portrait of a clown face, oil on canvas, bittersweet expression" 
 --train_batch_size=4 
 --num_train_epochs=3 
 --tracker_project_name="controlnet" 
 --enable_xformers_memory_efficient_attention 
 --checkpointing_steps=5000 
 --validation_steps=5000 
 --report_to wandb 
 --push_to_hub

让我们来分析一些设置，并讨论一些优化技巧，以低至8GB的VRAM进行训练：

• pretrained_model_name_or_path：您想要使用的稳定扩散基础模型（我们在这里选择了 v2-1，因为它可以更好地渲染面部）
• output_dir：要保存模型的目录
• dataset_name：将用于训练的数据集
• conditioning_image_column：数据集中包含条件图像的列的名称（在本例中为 spiga_seg）
• image_column：数据集中包含真实地表图像的列的名称（在我们的案例图像中）
• caption_column：数据集中包含图像标题的列的名称（在本例中为 image_caption）
• resolution：条件反射图像和真实地面图像的分辨率（在本例中为 512×512）
• learning_rate：出借率。我们发现 1e-5 适用于这些示例，但您可以尝试 1e-4 和 2e-6 之间的不同值。
• validation_image：这是让你在训练期间先睹为快！将针对每个validation_steps量运行验证图像，以便您可以查看训练的进展情况。在此处插入任意数量的调节图像的本地路径
• validation_prompt：与验证映像一起运行的提示。可以是任何可以测试模型是否训练良好的内容
• train_batch_size：这是适合 GPU 的训练批次的大小。由于拥有 A4，我们可以负担得起 100，但如果您的 GPU 具有较低的 VRAM，我们建议您将此值降低到 1。
• num_train_epochs：每个纪元对应于训练集中的图像将被模型“看到”的次数。我们试验了 3 个 epoch，但结果发现最好的结果只需要 1 个 epoch 多一点，我们的模型过度拟合了 3 个 epoch。
• checkpointing_steps：每隔 x 步保存一个中间检查点（在本例中为 5000）。每 5000 步，就会保存一个中间检查点。
• validation_steps：每 x 步运行validaton_prompt和validation_image。
• report_to：向何处报告您的培训。在这里，我们使用了权重和偏差，它给了我们这个很好的报告。但是将train_batch_size从 4 减少到 1 可能不足以让训练适合小型 GPU，以下是为每个 GPU VRAM 大小添加的一些附加参数：
• push_to_hub：将最终训练的模型推送到拥抱面部中心的参数。

安装在16GB VRAM GPU上

pip install bitsandbytes

--train_batch_size=1 
--gradient_accumulation_steps=4 
--gradient_checkpointing 
--use_8bit_adam

此外，我们还启用了梯度检查点和 8 位 Adam 以节省更多内存。

安装在12GB VRAM GPU上

--gradient_accumulation_steps=4 
--gradient_checkpointing 
--use_8bit_adam
--set_grads_to_none

4.结论！

这种训练 ControlNet 的经历非常有趣。我们成功地训练了一个可以遵循真实面部姿势的模型 – 但是它学会了制作不可思议的3D面部而不是真正的3D面部，因为这是它训练的数据集，它有自己的魅力和耀斑。

*本文正在参加 人工智能创作者扶持计划 ”