Alan Hou的个人博客

生成式AI第六章 Stable Diffusion微调

在上一章中,我们介绍了如何通过微调使语言模型学会以特定风格写作或学习特定领域的概念。我们可以将相同的原理应用于文生图模型,从而在仅使用单个GPU的情况下(对应预训练Stable Diffusion等模型所需的多GPU节点)定制这些模型。

本章中,我们将使用第四章中学习的预训练Stable Diffusion基础模型,并扩展让其学习它可能不了解的风格和概念,例如your pet或某种特定的绘画风格。我们还将学习如何赋予它新的能力,如超分辨率、图像修补或新的输入条件。

这里我们不会从头编写代码,而是理解并运行现有的用于微调模型的脚本。建议您克隆diffusers库,因为大多数示例都位于该库的examples文件夹中:

Stable Diffusion全微调

全模型(Full model)是指在出现一些模型定制技术(如低秩适配(LoRA)、文本反转和Dreambooth)后,用于微调的描述词。我们将在本章进一步讨论这些技术(我们在第五章中探讨了LLM的LoRA),但读者已经了解到这些技术可能不会微调整个模型!

在这些技术出现之前,全模型微调就被简单地称为微调。彼时,微调意味着进一步训练扩散模型——就像我们在第三和第四章中学到的那样,但目标是将其引导至想添加的特定知识。可以使Stable Diffusion学会通过提示词无法获取的风格或主题,或在模型发布后才出现的内容。正如全模型这个限定词所隐含的那样,一旦模型被微调,虽然在教它的风格或主题上会变得很好,但它可能会专门仅生成这种内容。本节将使用一个预制脚本来执行全微调。我们将使用diffusers库中的脚本 diffusers/examples/text_to_image/train_text_to_image.py

图6-1: Stable Diffusion微调

准备数据集

全模型微调需要一个相对较大的数据集,包含500多张图像。虽然听起来很多,但与训练整个Stable Diffusion模型所需的数十亿张图像相比,几百张图像的数据集显得微不足道。我们在本章进一步讨论的特定模型定制技术中,将学习如何使用少至四张图像的定制模型。

回到全模型微调!在我们调校文生图模型时,必须向其展示包含图像及描述这些图像的相应标题的数据集,就像模型在其预训练期间所展示的一样。如需一些想法,可以参考以下示例:

:虽然我们鼓励探索,但对于任何超出学习和教育内容的微调部署,如使用的图像是某位艺术家的风格、某个人的脸或任何持有知识产权的材料,询问作者或知识产权持有人是否可以这样做不仅是讲礼貌,还有可能是法律所要求的。

在我们的示例中,将使用哈勃望远镜图像,这些图像在拍摄后6个月内由NASA公之于众。创建哈勃望远镜数据集并微调Stable Diffusion以适应哈勃图像是德克萨斯大学达拉斯分校研究员Maxwell Weinzierl开创的,他制作了esa-hubble数据集和Hubble Diffusion 1和2模型。对于这个示例,我们的目标是通过微调Stable Diffusion v1.5与esa-hubble数据集来重新创建Hubble Diffusion 1模型。

esa-hubble数据集是通过从欧洲航天局网站抓取哈勃望远镜拍摄的图像创建的。所幸这些图像所描述的天文现象的标题也有用,因此可以将两者保存并置于数据集库兼容的格式中,以便用于微调。如何收集数据(通过网络抓取或其他方式)超出了本文的范围,但读者可以研究像Scrapy或BeautifulSoup这样的Python爬虫工具。留意每个网站关于抓取或爬取的政策。

一旦有了带有图像-文本对的数据集,可以将它们加载到数据集库中。以下部分将展示如何操作。如果没有可用的图像数据集,可以使用提供的esa-hubble数据集继续操作。

imagefolderload_dataset的特殊模式,可加载目录中的图片以及包含每张图标题的元数据文件。这一模式要求将所有图片放在/path/to/folder中,同时图片标题放在metadata.csv文件中。所有文件夹的结构是这样的:

metadata.csv文件的内容如下:

加载好数据集,可以通过一个简单的命令将其推送到Hugging Face Hub,与社区分享。

这时数据集已经保存并且可以用来微调模型了!有时你拥有一个完美的图像数据集,但没有人为其制作的标题。这时可以使用图像到文本模型来创建标题,然后可以用这些标题来微调模型。像Salesforce的BLIP和BLIP-2这样的模型广泛用于此任务。可以查看图像到文本任务页面了解更多信息。如果没有数据集,也不用担心,可以使用提供的esa-hubble数据集!

微调模型

前面讲到微调模型需要一个数据集,还需要一个训练脚本以及待微调模型的权重。幸运的是,我们可以使用diffusers库(提供示例训练脚本)和accelerate库(用于高效训练,使PyTorch的训练循环可以在多GPU、TPU或不同精度(如bf16)上运行)来轻松配置这些内容。这个脚本非常有用,因为它包含了高效训练Stable Diffusion的UNet所需的所有代码,同时通过公开的超参数赋予用户控制权。

需要一个至少具有16GB VRAM的GPU来执行此微调,或者使用诸如Google Colab Pro之类的服务。我们会进一步学习的自定义技术可在更小的GPU或Google Colab的免费版本上进行训练。

读者可以使用自己的数据集,或者使用Maxwell Weinzierl的esa-hubble数据集来复刻Hubble Diffusion模型。

要开始微调模型,首先需要从diffusers库中克隆必要的训练脚本:

然后使用train_text_to_image.py脚本。有关VRAM约束的设置,我们将使用use_8bit_adam,这需要bitsandbytes库,允许使用8位Adam优化器。

在我们深入讨论参数之前,建议先运行脚本并继续阅读本章,因为训练将花费一些时间(取决于你的GPU,大约1-3小时),然后在阅读完本章后再进一步调整超参数。

这里不详细介绍train_text_to_image.py脚本的工作原理,但基本概念与我们在第4章的训练模型小节中学到的相同。但理解关键的超参数(即在开始微调模型之前设置的参数)仍然非常重要,我们将逐一介绍上述训练脚本中的设置。

最重要的概念是学习率(learning_rate)和训练轮数(num_train_epochs)。

我们还将简要介绍其他超参数:

train_text_to_image.py训练脚本有更多参数和配置。Stable Diffusion XL也可以进行类似的微调。模型训练完成并推送到Hub后,就可以对其运行推理了!

推理

微调后,模型可以像我们在第4章中学习的常规Stable Diffusion模型那样用于推理,只不过这次加载的是一个新训练的模型。如果你没有计算资源用于训练模型,这里有一个在相同Hubble数据集上训练的模型可用于推理:Supermaxman/hubble-diffusion-1。你也可以通过Hugging Face平台分享自己的模型供他人使用。

在你尝试时,你可能会发现模型非常擅长生成类似Hubble望远镜拍摄的图像(或我们对其微调的内容)。然而,它变成了一个专门的模型。如果用提示词生成其他东西,它可能会生成一些银河系的输出或只是一些无意义的东西。这是因为微调整个模型会带来灾难性遗忘,即整个模型会朝你引导的方向调整。同时还需要用相当多的图像来训练它。Dreambooth和低秩适配(LoRA)等技术可以克服这些限制。

Dreambooth

Dreambooth是一种微调Stable Diffusion的自定义技术,首次出现在Google Research的Nataniel Ruiz等人的论文《DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation》中。Dreambooth技术通过全微调Stable Diffusion的UNet实现,同时为模型提供底座,例如用于先验保持损失的类图像 – 与训练的重构损失结合使用;并提供一个唯一的触发词来激活模型中的概念知识。

与全模型微调文生图扩散模型相比,Dreambooth带来了三个令人兴奋的进展:

原论文中的Dreambooth技术用于Google的专有扩散模型Imagen。然而,开源社区成员XavierXiao将该技术适配到Stable Diffusion,自那以后,许多社区实现,如TheLastBen和kohya-ss的相继出现。diffusers库也有Dreambooth训练脚本。

总体而言,社区的发现表明:

我们将使用diffusers库中的脚本diffusers/examples/dreambooth/train_dreambooth.py

:Dreambooth并不是第一个具有这一目标的自定义技术。基于开创性论文An Image is Worth One Word(Rinon Gal等人)的文本反转技术展示了如何为Stable Diffusion模型的文本编码器训练新的嵌入以包含新主题。该技术仍有效,因为训练的嵌入可以很小(只有几千字节)。然而,存在一个大小与质量的权衡,而Dreambooth的质量使其在文生图社区中占据主导地位。结合文本反演和Dreambooth的技术(称为Pivotal Tuning)的实验也显示出良好的效果。使用Dreambooth,你需要找到一个独特的触发词。然而,通过使用文本反转,我们可以创建这些触发词作为新token,这有助于更好地注入新概念。

准备数据集

5到20个样本通常足够训练一个新对象或人脸。如果模型在这个范围内学习困难,增加样本数量有助改进效果。由于Dreambooth不需要对图像进行描述,只需将训练图像放在所选文件夹中,并使用训练代码指向它们。例如,可以下载您的宠物照片!

在本例中,我们将训练一个模型识别一位作者的脸。如想跟着一起操作,可以训练一个识别您自己脸部的模型(或者是您的宠物)。

先验保存

使用Dreambooth,可以选择利用先验保存类。它通过在训练时提供先验保存损失来实现,使模型理解它生成的是同一类的元素。例如,在教模型识别您的脸时,拥有一组人脸图像将帮助模型“理解”要训练的类别是人脸,因此即使提供的样本较少,也能基于人脸图像进行训练。如果模型已经对要创建的类别有了知识,甚至可以自己生成先验图像(训练代码中有一个标志允许我们这样做)或将它们上传到特定文件夹。

可以在训练脚本中配置一些参数启用先验保存功能:

Dreambooth模型训练

与对整个模型进行微调类似,最重要的变量是learning_ratenum_train_epochs。低学习率加上逐渐增加的训练轮次或步骤可以作为获得高质量的结果很好的起点。另一种探索方式是固定训练轮次或步骤,并增加学习率。两者可以结合起来达到最优的超参数。

我们来看一些Dreambooth特有的参数:

我们使用train_dreambooth.py脚本来训练一个识别作者之一脸部的模型。如果想跟着一起操作,可以训练一个识别您自己脸部的模型。

图6-2: Polis 的脸

推理

尽管模型的整体结构得以保留并且只改变了基于instance_prompt的新token,新的Dreambooth模型仍然是完整的Stable Diffusion权重。因此,它可以作为推理加载。

这里是用作者之一的脸训练的模型的一些结果:

图6-3: Polis脸的Dreambooth

训练LoRA

全模型微调和Dreambooth有一个问题:一旦我们完成模型的微调,最终会得到与原始Stable Diffusion模型一样大的新权重。这种情况对于共享、本地托管、堆叠模型、云端服务以及其他下游应用来说都不理想。为此,我们可以使用低秩适配(LoRA),就像我们在上一章中对LLM所训练的那样!

图6-4:LoRA图

如上一章所见,LoRA可以冻结预训练模型权重并注入秩分解矩阵,显著减少需要训练的参数数量。LoRA训练的秩也可以作为工件(artifact)共享,注入到模型中而不会增加推理延迟。

听起来很棒!但我们在微调扩散模型,而原始的LoRA专注于transformer。这就是Simo Ryu的Stable Diffusion LoRA GitHub库出现的契机。意识到LoRA秩可以像添加到transformer LLM上一样附加到Stable Diffusion Unet和文本编码器上,现在已经解锁了LoRA对扩散模型的威力。

diffusers库再次发挥了作用,包含了一个用于LoRA训练的脚本,适用于全模型微调和dreambooth微调。使用diffusers训练LoRA权重几乎与全模型微调或使用diffusers进行dreambooth微调相同,但有一些主要区别:

赋予Stable Diffusion新功能

通过微调教授模型新风格或对象是非常棒的。但是,假如我们可以使用微调赋予Stable Diffusion更多的功能呢?通过使用一些特殊的技术进行微调,我们可以使模型具有修补(inpainting)、超分辨率或添加额外条件的能力。

修补

通过在UNet中包含额外的输入通道,可以赋予预训练的文生图扩散模型修补能力。在RunwayML的修补专家模型Stable Diffusion v1中,他们通过在UNet中添加五个清零的输入通道(四个用于编码的掩码图像,一个用于掩码本身)进行了约40万步的训练。在训练期间,生成了合成掩码,掩盖25%的内容。由于有掩码背后的图像真值,模型学习如何根据提示填充掩码区域,成为一个强大的图像编辑工具。

对于更高级的模型,如Stable Diffusion XL,一些修补功能在没有特定进一步微调就可以使用,这让一些人质疑一个专业微调模型是否可以提升这一能力。然而,SDXL专家修补模型发布了一些额外的功能,显示了这种技术在更大和更先进的模型中的潜力。虽然这种技术在家用硬件上无法访问,要求进行数十万步的全微调,但微调模型是可以的,任何人都可以使用。

额外输入用于特殊条件

就像可以为模型添加新的输入通道学习如何进行修补一样,也可以添加其他条件。一个应用的例子是Stable Diffusion 2 Depth模型,它从stable-diffusion-2-base恢复并进行了20万步的微调,添加了一个额外的输入通道,该通道处理用户提示词和由MiDaS生成的单目深度(相对于摄像机的距离)预测图像。

虽然这种技术效果很好,但为了获得新的条件,微调基模型数十万步将这一处理限定在少数公司和实验室。然而,附加在模型上层的适配器技术,如ControlNets、ControLoras和T2I适配器,出现了,使这一过程在训练和推理方面更有效率;我们将在下一章中探讨更多这些创造性的文生图应用。

项目实操:自行训练一个SDXL Dreambooth LoRA

微调是为文生图扩散模型带来更多知识的好方法,我们在本章已经学到,Dreambooth允许仅用少量样本图像进行微调,而LoRA训练相对于全模型微调在GPU使用量上更低、模型更小。对于这个项目,读者将微调一个扩散模型的Dreambooth LoRA。在学习了上面的基础知识后,读者可以使用我们的高级脚本。如果您没有16GB或以上显存的GPU,建议使用Google ColabHugging Face空间完成本项目。

目标是能够在Stable Diffusion中生成一个新的、尚不存在的对象或风格。

第一步:数据集创建:

第二步:模型训练:

结论

由于从头训练一个大型文生图模型需要大量计算资源,微调使得单GPU操作能够自定义预训练模型生成所需的内容。在本章中,我们学到了微调扩散模型如何扩展其知识并根据特定需求进行定制,同时保留其整体知识。我们学会了如何进行全模型微调、Dreambooth以特定角色或风格进行微调,以及使用LoRA提高效率。我们还了解到微调扩散模型可以赋予它们新的功能。总的来说,微调是一个强大的工具。

本章探讨了微调Stable Diffusion模型以教授其风格、对象或功能的技术。从全模型微调开始,我们学会了如何改变模型的行为来生成所需风格或主题的图像。然后,我们介绍了Dreambooth和低秩适应(LoRA)等技术,这些技术可使用更少的样本和更少的灾难性遗忘风险进行定制。

我们还讨论了通过微调添加新功能的潜力,例如修补和特殊条件,扩展了Stable Diffusion原始配置外的实用性。

练习

  1. 说明全模型微调和Dreambooth的主要区别?
  2. 使用LoRA相对于全模型微调在计算资源和模型适应性方面的优势是什么?
  3. 在进行Dreambooth训练时使用唯一token的重要性是什么?
  4. 除了教授新概念,微调还能为模型添加哪些新功能?列举应用微调技术后模型学到的两个功能。
  5. 讨论超参数的选取如何影响微调扩散模型的结果。
  6. 描述在偏见数据集上微调文生图模型的潜在风险。
    挑战
  7. LoRA与全模型微调比较:训练一个使用LoRA和全模型微调的Dreambooth模型并比较结果。尝试修改LoRA的rank超参数,看看它对结果的影响。
退出移动版