扩散模型学习笔记2

第5章 微调和引导
1.微调
通过微调（fine-tuning），我们将在新的数据集上重新训练已有的模型，来改变它原有的输出类型。

• 创建一个采样循环，并使用调度器（scheduler）更快地生成样本
• 在新数据集上微调一个现有的扩散模型，这包括：
  • 使用累积梯度的方法去应对训练的 batch 太小所带来的一些问题
  • 在训练过程中，将样本上传到 Weights and Biases 来记录日志，以此来监控训练过程（通过附加的实例脚本程序）
  • 将最终结果管线（pipeline）保存下来，并上传到Hub
• 通过新加的损失函数来引导采样过程，以此对现有模型施加控制，这包括：
  • 通过一个简单的基于颜色的损失来探索不同的引导方法
  • 使用 CLIP，用文本来引导生成过程

DDIM更快的采样过程
在生成的每一步中，模型都是接收一个加噪声的输入，并被要求去预测这个噪声（以此来估计完全没噪声的图片是什么样）。最初这些预测都还不算太好，因此我们把这一过程分解成很多步。但另一方面，使用1000+步来实现整个生成过程也不是必要的，因为最近的很多研究已经找到了使用尽可能少的步数来生成好的结果的方法。
在 Diffusers 库中，这些采样方法是通过调度器（scheduler）来操控的，每次更新通过step()函数完成。为了生成图片，我们从随机噪声$x$开始，每一个迭代周期（timestep）我们都送入模型一个带噪声的输入$x$并把模型预测结果再次输入step()函数。这里返回的输出都被命名为prev_sample —— 之所以是“previous”，是因为我们是在时间上“后退”，即从高噪声到低噪声（这和前向扩散过程是相反的）。

微调实验
实验中需要注意：
（1）考虑因素1：我们这里使用的 batch size 很小（只有 4），因为我们的训练是基于较大的图片尺寸的（256px），并且我们的模型也很大，如果我们的 batch size 太高，GPU 的内存可能会不够用了。你可以减小图片尺寸，换取更大的 batch size 来加速训练，但这里的模型一开始都是基于生成 256px 尺寸的图片来设计和训练的。
现在我们看看训练循环。我们把要优化的目标参数设定为image_pipe.unet.parameters()，以此来更新预训练过的模型的权重。其它部分的代码基本上和第一单元例子中的对应部分一样。在 Colab 上跑的话，大约需要10分钟，你可以趁这个时间喝杯茶休息一下。
（2）考虑因素2：*我们的损失值曲线简直像噪声一样混乱！这是因为每一次迭代我们都只用了四个训练样本，而且加到它们上面的噪声水平还都是随机挑选的。这对于训练来讲并不理想。一种弥补的措施是，我们使用一个非常小的学习率，限制每次更新的幅度。但我们还有一个更好的方法，既能得到和使用更大的 batch size 一样的收益，又不需要让我们的内存爆掉。
如果我们多运行几次loss.backward()后再调用optimizer.step()和optimizer.zero_grad()，PyTorch 就会把梯度累积（加和）起来，这样多个批次的数据产生的更新信号就会被高效地融合在一起，产出一个单独的（更好的）梯度估计用于参数更新。这样做会减少参数更新的总次数，就正如我们使用更大的 batch size 时希望看到的一样。梯度累积是一个很多框架都会替你做的事情但这里我们从头实现一遍也挺好的，因为这对你在 GPU 内存受限时训练模型非常有帮助。正如你在上面代码中看到的那样，其实也不需要你写很多代码。
（3）考虑因素3：*即使这样，我们的训练还是挺慢的，而且每遍历完一轮数据集才打印出一行更新，这也不足以让我们知道我们的训练到底怎样了。我们也许还应该：

• 训练过程中时不时地生成点图像样本，供我们检查模型性能
• 在训练过程中，把诸如损失值和生成的图片样本在内的一些东西记录到日志里。你可以使用诸如 Weights and Biases 或 tensorboard 之类的工具
  （4）考虑因素4：*微调这个过程可能是难以预知的。如果我们训练很长时间，我们也许能看见一些生成得很完美的蝴蝶，但中间过程从模型自身讲也极其有趣，尤其是你对艺术风格感兴趣时！你可以试试短时间或长时间地观察一下训练过程，并试着该百年学习率，看看这会怎么影响模型的最终输出。
  （5）考虑因素5：*一般而言，判断一次微调到底有多管用并不容易，而且“足够好的性能”在不同应用场景下代表什么水平也会有所不同。比如，如果你在一个很小的数据集上微调一个文本条件模型，比如stable diffusion模型，你可能希望模型极可能保留它原始训练所学习的东西以便于它能理解你的数据集没有涵盖的各种文本提示；同时你又希望它适配你的数据，以便它生成的东西和你的数据风格一致。这可能意味着，你需要使用一个很低的学习率，并配合对模型进行指数平均，就像这个关于创建一个宝可梦版 stable diffusion 模型的博客中的做法一样。别的情况下，你可能还想在一个新数据集上完全重训一个模型（就像我们前面从卧室图片到 wikiart 图片微调一样），那你就需要大的学习率和长时间训练了。即使从这里的损失值曲线看不出模型在进步，但生成样本已经很清楚地显示出了一个从原始数据到更有艺术范的风格迁移的过程了，虽然看着还是不太协调。

2、引导
通过引导（guidance），我们将在推理阶段引导现有模型的生成过程，以此来获取额外的控制。
第一步，我们先创建一个函数，定义我们希望优化的一个指标（损失值）。这里是一个让生成的图片趋向于某种颜色的例子，它将图片像素值和目标颜色（这里用的是一种浅蓝绿色）对比，返回平均的误差。
接下来，我们要修改采样循环，在每一步，我们要做这些事情：

• 创建一个新版的 x，并且 requires_grad = True
• 算出去噪后的版本（x0）
• 将预测出的x0送入我们的损失函数中
• 找到这个损失函数对于 x 的梯度
• 在我们使用调度器前，用这个梯度去修改 x ，希望 x 朝着能减低损失值的方向改进

这里有两种实现方法，你可以探索一下哪一种更好。第一，我们是在从 UNet 得到噪声预测后才给 x 设置 requires_grad 的，这样对内存来讲更高效一点（因为我们不用穿过扩散模型去追踪梯度），但这样做梯度的精度会低一点。第二种方法是，我们先给 x 设置 requires_grad，然后再送入 UNet 并计算预测出的 x0。

CLIP 引导
引导生成的图片向某种颜色倾斜确实让我们多少对生成有所控制，但如果我们能仅仅打几行字描述一下就得到我们想要的图片呢？

CLIP 是一个由 OpenAI 开发的模型，它可以让我们拿图片和文字说明去作比较。这是个非常强大的功能，因为它让我们能量化一张图和一句提示语有多匹配。另外，由于这个过程是可微分的，我们可以使用它作为损失函数去引导我们的扩散模型。

这里我们不深究细节。基本的方法是：

• 给文字提示语做嵌入（embedding），为 CLIP 获取一个 512 维的 embedding
• 对于扩散模型的生成过程的每一步：
  • 做出多个不同版本的预测出来的去噪图片（不同的变种可以提供一个更干净的损失信号）
  • 对每一个预测出的去噪图片，用 CLIP 给图片做嵌入（embedding），并将这个嵌入和文字的嵌入做对比（用一种叫 Great Circle Distance Squared 的度量方法）
• 计算这个损失对于当前带噪的 x 的梯度，并在用调度器（scheduler）更新它之前用这个梯度去修改 x