随着 Meta 上周发布的 Segment Anything Model (SAM)，计算机视觉迎来了它的 ChatGPT 时刻。SAM 训练了超过 110 亿个分割掩码，是预测 AI 用例而非生成 AI 的基础模型。虽然它在对广泛的图像模式和问题空间进行分割的能力方面表现出了令人难以置信的灵活性，但它在发布时没有“微调”功能。
本教程将概述使用掩码解码器微调 SAM 的一些关键步骤，特别是描述 SAM 中使用哪些函数对数据进行预处理/后处理，使其处于良好状态以进行微调。

什么是SAM？

Segment Anything Model (SAM) 是由 Meta AI 开发的一种分割模型。它被认为是计算机视觉的第一个基础模型。SAM 在包含数百万图像和数十亿掩码的庞大数据集上进行了训练，使其非常强大。顾名思义，SAM 能够为各种图像生成准确的分割掩码。Sam 的设计允许它考虑人类提示，这使得它对于 Human In The Loop 注释特别强大。这些提示可以是多模式的：它们可以是要分割区域上的点、要分割对象周围的边界框或关于应该分割什么的文本提示。

该模型分为 3 个组件：图像编码器、提示编码器和掩码解码器。

图像编码器为被分割的图像生成嵌入，而提示编码器为提示生成嵌入。图像编码器是模型中特别大的组件。这与轻量级掩码解码器形成对比，轻量级掩码解码器根据嵌入预测分割掩码。Meta AI 已将在 Segment Anything 10 亿掩码 (SA-1B) 数据集上训练的模型的权重和偏差用作模型检查点。在此处的解释器博客文章中了解有关 Segment Anything 工作原理的更多信息。

什么是模型微调？

公开可用的最先进模型具有自定义架构，通常提供预训练模型权重。如果这些架构在没有权重的情况下提供，那么用户将需要从头开始训练模型，用户将需要使用大量数据集来获得最先进的性能。

模型微调是采用预训练模型（架构+权重）并显示特定用例数据的过程。这通常是模型以前没有见过的数据，或者在其原始训练数据集中代表性不足的数据。

微调模型和从头开始的区别在于权重和偏差的起始值。如果我们从头开始训练，这些将根据某种策略随机初始化。在这样的初始配置中，模型对手头的任务“一无所知”并且表现不佳。通过使用预先存在的权重和偏差作为起点，我们可以“微调”权重和偏差，以便我们的模型在我们的自定义数据集上更好地工作。例如：学习识别猫的信息（边缘检测、计数爪子）将对识别狗有用。

为什么要微调模型？

微调模型的目的是在预训练模型以前没有见过的数据上获得更高的性能。例如，在从手机摄像头收集的广泛数据集上训练的图像分割模型将主要从水平视角看到图像。

如果我们尝试将此模型用于从垂直角度拍摄的卫星图像，它的性能可能不会那么好。如果我们试图分割屋顶，该模型可能不会产生最佳结果。预训练很有用，因为模型已经学会了一般如何分割对象，所以我们想利用这个起点来构建一个可以准确分割屋顶的模型。此外，我们的自定义数据集可能不会有数百万个示例，因此我们希望微调而不是从头开始训练模型。

微调是可取的，这样我们就可以在特定用例上获得更好的性能，而不必承担从头开始训练模型的计算成本。

如何微调 Segment Anything 模型

背景与架构

我们在介绍部分概述了 SAM 体系结构。图像编码器具有包含许多参数的复杂架构。为了微调模型，我们将重点放在轻量级的掩码解码器上是有意义的，因此微调更容易、更快、内存效率更高。

为了微调 SAM，我们需要提取其架构的底层部分（图像和提示编码器、掩码解码器）。我们不能使用SamPredictor.predict (链接) 有两个原因：

• 我们只想微调掩码解码器
• 这个函数调用SamPredictor.predict_torch，它有 @torch.no_grad()装饰器（链接），它阻止我们计算梯度。

因此，我们需要检查SamPredictor.predict函数并调用适当的函数，并在我们想要微调的部分（掩码解码器）上启用梯度计算。这样做也是了解更多有关 SAM 工作原理的好方法。

创建自定义数据集

我们需要三件事来微调我们的模型：

• 在其上绘制分割的图像
• 分割真值掩码
• 提示输入模型，我正在使用边界框

我选择了印章验证数据集（链接），因为它包含 SAM 在其训练中可能没有看到的数据（即文档上的印章）。我可以通过使用预先训练的权重运行推理来验证它在这个数据集上表现良好，但并不完美。ground truth masks 也非常精确，这将使我们能够计算准确的损失。最后，该数据集包含分割掩码周围的边界框，我们可以将其用作 SAM 的提示。示例图像如下所示。这些边界框与人工注释者在寻找生成分段时所经历的工作流程非常吻合。

输入数据预处理

我们需要预处理从 numpy 数组到 pytorch 张量的扫描。为此，我们可以关注SamPredictor.set_image (链接) 和SamPredictor.set_torch_image (链接) 内部发生的事情，它们对图像进行预处理。首先，我们可以使用utils.transform.ResizeLongestSide来调整图像大小，因为这是预测器内部使用的转换器 ( link )。然后我们可以将图像转换为pytorch张量，并使用 SAM 预处理方法（链接）完成预处理。

训练设置

我们下载vit_b模型的模型检查点并将它们加载到：

sam_model = sam_model_registry['vit_b'](checkpoint='sam_vit_b_01ec64.pth')
我们可以使用默认设置 Adam 优化器，并指定要调整的参数是掩码解码器的参数：

optimizer = torch.optim.Adam(sam_model.mask_decoder.parameters())
同时，我们可以设置我们的损失函数，例如Mean Squared Error

loss_fn = torch.nn.MSELoss()

循环训练

在主训练循环中，我们将迭代我们的数据项，生成掩码并将它们与我们的真实掩码进行比较，以便我们可以根据损失函数优化模型参数。

在此示例中，我们使用 GPU 进行训练，因为它比使用 CPU 快得多。在适当的张量上使用.to(device)非常重要，以确保我们不会在 CPU 上使用某些张量而在 GPU 上使用其他张量。

我们希望通过将编码器包装在torch.no_grad()上下文管理器中来嵌入图像，否则我们将遇到内存问题，以及我们不希望微调图像编码器的事实。

with torch.no_grad():
	image_embedding = sam_model.image_encoder(input_image)

我们还可以在 no_grad 上下文管理器中生成提示嵌入。我们使用我们的边界框坐标，转换为 pytorch 张量。

with torch.no_grad():
      sparse_embeddings, dense_embeddings = sam_model.prompt_encoder(
          points=None,
          boxes=box_torch,
          masks=None,
      )

最后，我们可以生成掩码。请注意，这里我们处于单一掩码生成模式（与通常输出的 3 个掩码形成对比）。

low_res_masks, iou_predictions = sam_model.mask_decoder(
  image_embeddings=image_embedding,
  image_pe=sam_model.prompt_encoder.get_dense_pe(),
  sparse_prompt_embeddings=sparse_embeddings,
  dense_prompt_embeddings=dense_embeddings,
  multimask_output=False,
)

这里的最后一步是将蒙版放大回原始图像大小，因为它们的分辨率很低。我们可以使用Sam.postprocess_masks来实现这一点。我们还希望根据预测的掩码生成二进制掩码，以便我们可以将它们与我们的基本事实进行比较。为了不破坏反向传播，使用torch功能很重要。

upscaled_masks = sam_model.postprocess_masks(low_res_masks, input_size, original_image_size).to(device)

from torch.nn.functional import threshold, normalize

binary_mask = normalize(threshold(upscaled_masks, 0.0, 0)).to(device)

最后我们可以计算损失并运行优化步骤：

loss = loss_fn(binary_mask, gt_binary_mask)
optimizer.zero_grad()
loss.backward()
optimizer.step()

通过在多个时期和批次上重复此操作，我们可以微调 SAM 解码器。

保存检查点并从中启动模型

一旦我们完成训练并对性能提升感到满意，我们可以使用：

torch.save(model.state_dict(), PATH)

保存调优模型的状态字典。当我们想要对与我们用于微调模型的数据相似的数据执行推理时，我们可以加载这个状态字典。

您可以在此处找到 Colab Notebook，其中包含微调 SAM 所需的所有代码。如果您想要一个开箱即用的完整解决方案，请继续阅读！

下游应用程序的微调

虽然 SAM 目前不提供开箱即用的微调，但我们正在构建与 Encord 平台集成的自定义微调器。如本文所示，我们微调解码器以实现此目的。这可以在 Web 应用程序中作为开箱即用的一键式程序使用，其中会自动设置超参数。

原SAM 预测结果：

由模型的微调版本生成的掩码：

我们可以看到这个面具比原来的面具更紧。这是对邮票验证数据集中的一小部分图像进行微调，然后在以前未见过的示例上运行调优模型的结果。通过进一步的训练和更多的例子，我们可以获得更好的结果。