关于人工智能:使用-diffusers-训练你自己的-ControlNet-🧨

简介

ControlNet 这个神经网络模型使得用户能够通过施加额定条件，细粒度地管制扩散模型的生成过程。这一技术最后由 Adding Conditional Control to Text-to-Image Diffusion Models 这篇论文提出，并很快地风靡了扩散模型的开源社区。作者开源了 8 个不同的模型，使得用户能够用 8 种条件去管制 Stable Diffusion 模型（包含版本 1 到 5 ）。这 8 种条件包含姿势预计、深度图、边缘图、素描图 等等。

在这篇博客中，咱们首先介绍训练 Uncanny Faces model 的步骤。这是一个基于 3D 合成人脸的人脸姿势模型（这里的 uncanny faces 只是一个无心失去的后果，前面咱们会讲到）。

开始着手用 Stable Diffusion 训练你的 ControlNet

训练你本人的 ControlNet 须要 3 个步骤:

1. 设计你想要的生成条件: 应用 ControlNet 能够灵便地“驯服” Stable Diffusion，使它朝着你想的方向生成。预训练的模型曾经展现出了大量可用的生成条件，此外开源社区也曾经开发出了很多其它条件，比方这里 像素化的色调板。
2. 构建你本人的数据集: 当生成条件确定好后，就该构建数据集了。你既能够从头构建一个数据集，也能够应用现有数据集中的数据。为了训练模型，这个数据集须要有三个维度的信息: 图片、作为条件的图片，以及语言提醒。
3. 训练模型: 一旦数据集建好了，就能够训练模型了。如果你应用 这个基于 diffusers 的训练脚本，训练其实是最简略的。这里你须要一个至多 8G 显存的 GPU。

1. 设计你想要的生成条件

在设计你本人的生成条件前，有必要考虑一下两个问题:

1. 哪种生成条件是我想要的？
2. 是否已有现存的模型能够把失常图片转换成我的条件图片？

举个例子，如果咱们想要应用人脸特色点作为生成条件。咱们的思考过程应该是这样: 1. 个别基于特色点的 ControlNet 成果都还挺好。2. 人脸特色点检测也是一个很常见的工作，也有很多模型能够在一般图片上检测人脸特色点。3. 让 Stable Diffusion 去依据特色点生成人脸图片也挺有意思，还能让生成的人脸模拟他人的表情。

2. 构建你本人的数据集

好！那咱们当初曾经决定用人脸特色点作为生成条件了。接下来咱们须要这样构建数据集:

• 筹备 ground truth 图片 (image): 这里指的就是实在人脸图片
• 筹备 条件图片 (conditioning_image): 这里指的就是画进去的特色点
• 筹备 阐明文字 (caption): 形容图片的文字

针对这个我的项目，咱们应用微软的 FaceSynthetics 数据集: 这是一个蕴含了 10 万合成人脸的数据集。你可能会想到其它一些人脸数据集，比方 Celeb-A HQ 和 FFHQ，但这个我的项目咱们决定还是采纳合成人脸。

这里的 FaceSynthetics 数据集看起来是个不错的抉择: 它蕴含了实在的人脸图片，同时也蕴含了被标注过的人脸特色点（依照 iBUG 68 特色点的格局），同时还有人脸的宰割图。

然而，这个数据集也不是完满的。咱们后面说过，咱们应该有模型能够将实在图片转换到条件图片。但这里仿佛没有这样的模型，把人脸图片转换成咱们特色点标注模式（无奈把特色点转换为宰割图）。

所以咱们须要用另一种办法:

• 应用 FaceSynthetics 中的实在图片 (image)
• 应用一个现有的模型把人脸图片转换为 68 个特色点的模式。这里咱们应用 SPIGA 这个模型
• 应用本人的代码把人脸特色点转换为人脸宰割图，以此作为“条件图片” (conditioning_image)
• 把这些数据保留为 Hugging Face Dataset

这里 是将实在图片转换到宰割图的代码，以及将数据保留为 Hugging Face Dataset 的代码。

当初咱们筹备好了 ground truth 图片和“条件图片”，咱们还短少阐明文字。咱们强烈推荐你把阐明文字加进去，但你也能够试试应用空的阐明文字来看看成果。因为 FaceSynthetics 数据集并没有自带阐明文字，咱们应用 BLIP captioning 去给图片加上文字（代码在这里）。

至此，咱们就实现了数据集的构建。这个 Face Synthetics SPIGA with captions 数据集蕴含了 ground truth 图片、条件图片，以及对应的阐明文字，总计有 10 万条数据。所有就绪，咱们当初能够开始训练模型了。

3. 模型训练

有了 数据，下一步就是训练模型。即便这部分很难，但有了 上面的脚本，这个过程却变成了最简略的局部。咱们用了一个 A100 GPU去训练（在 LambdaLabs 每小时 1.1 美元租的）。

咱们的训练教训

咱们以 batch size 为 4 训练了 3 个 epoch。结果表明此策略有些太激进，导致后果呈现过拟合景象。模型有点遗记人脸的概念了，即便提醒语中蕴含“怪物史莱克”或“一只猫”，模型也只会生成人脸而不是“史莱克”或猫；同时模型也对各种格调变得不敏感。

如果咱们只训练 1 个 epoch (即模型仅学习了 10 万张照片)，模型倒是能遵循输出的姿势，同时也没什么过拟合。看起来还行，但因为咱们用的是合成数据，模型最终生成的都是些看起来很 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 以下显存的 GPU 训练。

• pretrained_model_name_or_path: 根底的 Stable Diffusion 模型，这里咱们应用 v2-1 版本，因为这一版生成人脸成果更好
• output_dir: 保留模型的目录文件夹
• dataset_name: 用于训练的数据集，这里咱们应用 Face Synthetics SPIGA with captions
• conditioning_image_column: 数据集中蕴含条件图片的这一栏的名称，这里咱们用 spiga_seg
• image_column: 数据集中蕴含 ground truth 图片的这一栏的名称，这里咱们用 image
• caption_column: 数据集中蕴含文字说明的这一栏的名称，这里咱们用 image_caption
• resolution: ground truth 图片和条件图片的分辨率，这里咱们用 512x512
• learning_rate: 学习率。咱们发现设成 1e-5 成果很好，但你也能够试试介于 1e-4 和 2e-6 之间的其它值
• validation_image: 这里是让你在训练过程中偷窥一下成果的。每隔 validation_steps 步训练，这些验证图片都会跑一下，让你看看以后的训练成果。请在这里插入一个指向一系列条件图片的本地门路
• validation_prompt: 这里是一句文本提醒，用于和你的验证图片一起验证以后模型。你能够依据你的须要设置
• train_batch_size: 这是训练时应用的 batch size。因为咱们用的是 V100，所以咱们还有能力把它设成 4。但如果你的 GPU 显存比拟小，咱们举荐间接设成 1。
• num_train_epochs: 训练模型应用的轮数。每一轮模型都会看一遍整个数据集。咱们试验用的是 3 轮，但仿佛最好的后果应该是呈现在一轮多一点的中央。当训练了 3 轮时，咱们的模型过拟合了。
• checkpointing_steps: 每隔这么多步，咱们都会保留一下模型的两头后果检查点。这里咱们设置成 5000，也就是每训练 5000 步就保留一下检查点。
• validation_steps: 每隔这么多步，validation_image 和 validation_prompt 就会跑一下，来验证训练过程。
• report_to: 向哪里报告训练状况。这里咱们应用 Weights and Biases 这个平台，它能够给出好看的训练报告。
• push_to_hub: 将最终后果推到 Hugging Face Hub.

然而将 train_batch_size 从 4 减小到 1 可能还不足以使模型可能在低配置 GPU 上运行，这里针对不同 GPU 的 VRAM 提供一些其它配置信息:

适配 16GB 显存的 GPU

pip install bitsandbytes

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

这里 batch size 设为 1，同时应用 4 步的梯度累计等同于你应用原始的 batch size 为 4 的状况。除此之外，咱们开启了对梯度保留检查点，以及 8 bit 的 Adam 优化器训练，以此更多地节俭显存。

适配 12GB 显存的 GPU

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

适配 8GB 显存的 GPU

请参考 咱们的教程

4. 总结

训练 ControlNet 的过程十分乏味。咱们曾经胜利地训练了一个能够模拟实在人脸姿势的模型。然而这个模型更多是生成 3D 格调的人脸图片而不是实在人脸图片，这是因为咱们应用了合成人脸的数据执行训练。当然这也让生成的模型有了独特的魅力。

试试咱们的 Hugging Face Space

下一步，为了生成实在的人脸图片，同时还不应用实在人脸数据集，咱们能够用 Stable Diffusion Image2Image 跑一遍所有的 FaceSynthetics 图片，把看起来很 3D 的人脸转换成实在人脸图片，而后再训练 ControlNet。

请持续关注咱们，接下来咱们将举办 ControlNet 训练赛事。请在 Twitter 关注 Hugging Face，或者退出咱们的 Discord 以便接管最新消息！

如果您此前曾被 ControlNet 将 Lofi Girl 的动画形象改成写实版类真人画面的能力折服，却又放心本人难以承当昂扬的训练老本，无妨报名加入咱们近期的提供收费 TPU 的 ControlNet 微调流动。试试看充斥创造力的你能有什么样的大作！

英文原文: https://hf.co/blog/train-your-controlnet

作者: Apolinário from multimodal AI art, Pedro Cuenca

译者: HoiM Y, 阿东