使用 PyTorch 加速生成式 AI：Segment Anything 2 - 低延迟、快速冷启动的高速推理

作者：PyTorch 团队

本文是我们系列博客第一篇的后续，该系列博客专注于如何使用纯原生 PyTorch 加速生成式 AI 模型，并侧重于延迟和弹性扩展性。我们使用 torch.compile 和 torch.export 创建了高度优化、低延迟的 SAM2 版本，可在新实例上快速扩展。

通过利用 AOTInductor (AOTI) 经由 torch.export 进行的提前编译、降低精度、批量处理提示和 GPU 预处理，我们观察到与常规 eager 模式 PyTorch 相比，p90 执行延迟和队列时间最多提高了 13 倍。

我们计算了最终结果，并在 Modal 的自动扩展云基础设施上进行了实际部署演示了改进效果。

任务

第一篇文章专注于处理每张图像少量变化的提示（兴趣点）。这些点代表了地面真实掩码的中心点。对于本文，我们将重点放在更广泛的任务集上：单提示分割 (SPS)、多提示分割 (MPS) 和自动掩码生成 (AMG)，后者为输入图像生成完整的掩码集，无需给定提示集。第一篇文章仅关注 MPS。

图像中的小星星代表用户提示。对于 AMG，没有提示，掩码是从初始候选提示（猜测）的密集网格中启发式筛选出来的。对于 SPS 和 MPS，用户提示是从 AMG 掩码的中心点导出的。对于 SPS，我们选择面积最大的掩码。

请注意，SAM2 使用与 SAM1 不同的骨干网络。特别是，本文仅考虑最大且最准确的 sam2.1_hiera_large 骨干网络。

我们将重现结果所需的脚本汇集到 torchao 的示例文件夹中，并将 torchao 中对 SAM2 模型的更改中较稳定的部分逐步向上游提交到主 SAM2 仓库。因此，如果您有兴趣查看前沿变体或想贡献实验性功能，请随时联系 torchao 仓库和团队。对于更稳定和最新的模型版本，请直接访问 SAM2。

概述

我们将此处提出的更改分为两类。“快速”更改仅限于不影响模型精度的技术。“狂暴”更改通过利用低精度数据类型等近似方法，牺牲部分数值精度以换取额外速度。

近似方法可能会稍微降低精度指标，以支持显著提高的性能，同时仍能通过基于平均交并比 (mIoU) 的端到端检查。

为了衡量性能提升，我们处理了从 SAM2 验证数据集中随机选择的 1000 张图像。我们关注每张图像的 p50 和 p90 延迟。为了衡量精度，我们考虑 mIoU。最值得注意的是，对于 AMG 任务，我们还定义了一个失败计数指标。如果掩码数量不同，我们认为比较失败。事实证明这是一个相当不稳定的量，我们可以看到其他任务不像 AMG 那样对微小的数值变化敏感。

设置

我们在常规的 H100 开发服务器上运行离线实验，这是一台性能相当强大的机器。

然而，我们尝试在实际约束下看待这些任务。特别是，我们希望模拟服务器端推理环境。这意味着我们不使用 DataLoader 来隐藏图像预处理或解码例程的延迟。

对于延迟计算，我们包括了解码、分割以及将掩码转换为运行长度编码掩码字典的过程。或者换句话说，我们不包括将图像加载到内存中的主机字节数组以及将生成的字典存储为磁盘上的 json 文件。这是为了模拟更实际的设置。

更具体地说，考虑下面用于我们测量中包含的例程的代码。对于任何任务，gen_masks 代码会生成一个批处理的布尔 Tensor 位掩码，表示相应的对象掩码。然后我们将此位掩码压缩成运行长度编码 (rle) 格式，该格式可以更有效地将结果从远程服务器传回。

image_tensors = decode_img_bytes(...)
masks = gen_masks(image_tensors, ...)
rle_dicts = [rle_dict_from_masks(m) for m in masks]

优化

ao: eager 代码优化

这项工作最有效的工具是 PyTorch autograd 分析器与 record_function 的结合使用。为了构建这个软件，我们反复使用分析器来观察程序并确认任何更改的有效性。同样重要的是要注意分析器本身有开销。收集的数据越多，例如堆栈跟踪，引入的开销就越大，这可能会扭曲收集到的跟踪。但它非常适合查找同步点、内核之间的间隔以及耗时较长的 GPU 内核。

GPU 跟踪有助于您了解不一定容易通过编译解决的瓶颈。我们发现自动掩码生成 (AutomaticMaskGeneration) 特别受用于存储掩码的数据结构以及用于将掩码转换为运行长度编码压缩格式的例程的支配。我们还发现 AMG 性能的很大一部分是由作为单个批次创建的大量掩码主导的。有时可以通过重新排序操作在后处理阶段更早地将候选掩码过滤到更少的候选。这反过来又显著加快了后续操作。

为了确认我们实现的准确性，我们首先在不改变任何设置并使用 float32 精度的情况下进行比较。我们看到，在使用完全相同的设置时，mIoU 保持不变，并且掩码完美匹配。这意味着这些 eager 模式更改没有影响这些任务的准确性。

AMG

ao: 批量处理提示

我们能够应用的另一个无损性能优化是批量处理用户输入提示计算。当在 H100 等服务器级 GPU 上以批量大小 1 优化延迟时，我们通常会剩下大量备用内存。通过一次处理更多兴趣点（也称为用户提示），我们可以轻松地牺牲内存来换取更高的性能。请记住，SAM2 分为两部分：首先是骨干网络（图像编码器），其次是基于一组用户提示/兴趣点的掩码预测和解码。我们在第二部分中可能会期望更大或甚至可变数量的输入，并且正是在这第二部分中我们应用了批量处理。

这导致内存大幅增加，但也带来了更好的延迟。基线在循环中为每个提示生成一个掩码。对于 AMG，基线一次处理 64 个提示，需要做的只是将其更改为 1024，这是生成的候选提示的数量。对于 SPS，我们一次处理一个提示，但为了完整性，它仍然包含在下方。

AMG

SPS

MPS

作为技术旁注：最值得注意的是，为了启用 MPS 的批量处理并避免对手写代码库进行大量重写以同时支持多个提示，我们使用了一个名为 MapTensor 的 Tensor 子类。MapTensor 允许我们传递一批 N 个提示，但让它声称批量大小为 1。然后任何操作都会自动广播到包装的 Tensor，并传播到模型的预测部分。这之所以有效，是因为单个提示预测相互独立。这与 torch.vmap 非常相似。

center_points_torch = to_map_tensor(center_points_torch)
center_points_label_torch = to_map_tensor(center_points_label_torch)
masks, scores, _ = mask_generator.predictor.predict(
    point_coords=center_points_torch,
    point_labels=center_points_label_torch,
    multimask_output=True,
    return_logits=False,
    return_type="torch",
)
# Unwrapping MapTensor
masks = masks.elems
scores = scores.elems

fast: 全图编译

与我们第一篇文章一样，我们首先删除 GPU 同步和图中断，以利用适当位置使用 max-autotune 内核的全图编译模型代码。经过一些重写后，我们能够编译图像编码器和掩码预测。

我们运行了两次实验，以了解编译带来的开销。第一次是在 TORCHINDUCTOR_CACHE_DIR 为空的环境中运行，然后再次运行，同时利用前一次运行产生的工件。特别是，自动调优可能需要很长时间，并且会在首次在原始环境中调用时发生。我们将第二次运行称为“预热”。首次迭代通常会因各种其他相关的初始化过程而变慢，但即使使用现有缓存并再次输入完全相同的形状，编译也会显著增加其时间。话虽如此，在预热环境中，首次调用的几秒钟开销通常仍然可以接受。

这些缺点大多可以缓解，并且编译会显著改善延迟并减少内存。

AMG

使用自动掩码分割时，每个掩码生成的掩码数量可能会略有不同。模型为每个对象生成的掩码数量存在模糊性。例如，一辆汽车可能被细分为框架、窗户和车门，或者作为一个整体对待。当修改导致掩码数量改变时，我们认为比较失败，并且仅对完全匹配的掩码计算 mIoU。这不适用于其他任务。我们发现生成的掩码数量对微小的数值变化非常敏感。其他任务使用相同的代码，特别是 MPS 可以帮助我们进一步验证正确性。

SPS

MPS

furious: TF32、float16 和 GPU 预处理

我们发现使用 float16 对于模型的几个重要子组件来说是正确的精度级别。特别是，图像编码器和掩码解码器的权重可以完全转换为 float16。我们还可以对剩余的 float32 矩阵操作使用 TensorFloat32 精度。应该有可能进一步降低精度，我们可能会在以后的文章中讨论这个问题。在狂暴模式下，我们还将图像归一化等图像预处理移到 GPU 上。我们不能使用 GPU 解码 (nvJPEG) 例程，因为差异太大，模型在 mIoU 上会遭受严重退化，因此图像解码仍然在 CPU 上进行。

AMG

这导致 AMG 任务的 mIoU 显著下降，但不影响其他任务。经过深入调查，我们仍然将其归因于数值不稳定和操作的重新排序。需要做更多的工作来进一步研究这一点，并且在较低精度下运行 AMG 任务可能没有意义。然而，其他任务在延迟方面受益巨大，而 mIoU 变化很小。

SPS

MPS

AOTInductor (AOTI) 经由 torch.export 进行的提前编译

弹性扩展时，通常无法容忍长时间的启动时间。这意味着第一次迭代不能慢，而必须快速提供结果。这时 torch.compile 当前的编译开销可能会成为障碍。为了解决这个问题，我们可以使用 AOTInductor (AOTI) 经由 torch.export 进行的提前编译。AOTI 允许我们在代表性输入上编译模型，并将生成的代码存储在一个加载和运行速度很快的二进制文件中。

AOTI 经由 torch.export 是一个新功能，目前我们还不能导出所有可编译的内容。我们已经能够导出所有任务的图像编码器，但由于提示可变，我们只能导出 AMG 和 SPS 任务的掩码预测。torch.export 也支持动态形状，但我们需要投入更多时间来为此准备代码。

AMG: AO + 批量处理 + furious

SPS: AO + furious

请注意，加载导出的模型会显著增加内存。它可能只增加了峰值内存利用率，因为在加载导出的模型之前确实需要延迟初始化，以避免一次性在内存中存在两倍的权重。这是我们可以解决的问题，但当前的内存消耗远未达到上限。我们没有看到其他任务的内存增加，因为 AMG 和 MPS 的峰值内存主要由处理批量掩码决定。减少内存的一种方法可能是在早期阶段就使用 rle 格式（或其他稀疏格式）处理掩码，但目前来看，考虑到当前的内存消耗和对延迟的关注，没有理由这样做。

MPS: AO + 批量处理 + furious

单独使用导出似乎不会从广泛的预热中受益，并且可以在全新的 inductor 缓存目录中运行。但同样，我们不会清除 CUDA 缓存或其他缓存。在关于 Modal 的部分，我们在原始环境中运行其中一些实验。

在新进程中仅处理 1000 张图像时，使用导出确实值得，可以节省编译和其他冷启动开销。

彩蛋：更多 GPU 预处理

此时，延迟已经相当低。特别是对于 SPS 和 MPS 任务，我们处理的时间约为 30ms 到 40ms。让我们再次回顾设置部分的伪代码。

image_tensors = decode_img_bytes(...)
masks = gen_masks(image_tensors, ...)
rle_dicts = [rle_dict_from_masks(m) for m in masks]

进一步分析表明，此时 decode_img_bytes 代码大约需要 10ms。特别是，它使用 torchvision 的 ToTensor 转换将 numpy Tensor 转换为缩放的 float32 torch.Tensor。传递给 ToTensor 的字节已经解码并转换为 numpy ndarray。通过稍微重写 ToTensor，使用 torchvision 的 v2 API，并将 uint8 解码的小整数 Tensor 先移到 GPU 再进行缩放，我们可以在延迟上再节省 10ms。如果我们的分析中不包括 decode_img_bytes，我们就会错过这个对服务器端推理有实际影响的机会。

image_tensor = torch.from_numpy(image_tensor)
image_tensor = image_tensor.permute((2, 0, 1))
image_tensor = image_tensor.cuda()
image_tensor = v2.ToDtype(torch.float32, scale=True)( image_tensor)

特别注意，使用固定内存 (pinned memory) 执行异步数据传输在此处不适用，因为将 Tensor 移动到固定内存所需的时间不值得为了此数据移动而获得的异步性提升。对于未来的工作，我们可能希望通过使用更高级的直接内存传输技术来进一步探索这里的改进。

AMG: AO + 批量处理 + furious

SPS: AO + furious

MPS: AO + 批量处理 + furious

这个小改动对 SPS 和 MPS 任务产生了显著影响。

在 Modal 上部署

最后，我们将优化后的推理部署到无服务器基础设施提供商 Modal 上，以证明这些优化的优势可以在更实际的部署环境中实现。

特别是，通过 torch.export 进行编译和 AOTI 需要额外的工作。在简单部署中，这项工作可能会添加到每一次推理执行中，增加的延迟可能会使更快的模型带来的任何改进相形见绌。这对于弹性或自动扩展基础设施尤其具有挑战性，因为推理服务的副本需要定期自动创建和销毁。

我们在 torchao 仓库中共享了一个部署脚本（cli_on_modal.py），以演示一种弹性部署模式。我们提前构建导出的模型，然后将其上传到分布式存储。与 eager 执行相比，当副本启动时，这会增加一些额外的工作，因为它们需要通过网络读取这些数据，但这远比编译或导出所需的成本低。

我们使用大规模批量推理工作负载对该部署进行了基准测试：发送 1000 张图像进行并发处理。该部署在高峰时可扩展到十个 GPU 上的十个副本，并在不活动时缩减到零个 GPU。

首先，让我们看一下执行延迟。

我们注意到 Modal 和离线环境下的执行延迟相当接近，特别是相对于基线而言，这表明离线优化部署是直接优化部署的一个合理替代方案。

除了执行延迟外，我们的批量工作负载还存在队列时间，因为副本数量少于输入数量，因此一些输入需要排队等待。

尽管基础设施提供的排队系统没有改变，但当我们使用优化后的模型时，队列延迟也随之降低——在 p90 的情况下降低了 2 到 12 倍。这是因为当我们更快地完成先前的输入（由于执行延迟降低）时，我们可以更早地拉取下一个输入（从而减少它们的排队时间）。

如果您对进一步优化 SAM2 推理或部署感兴趣，请随时通过 torchao 仓库联系我们！

结论

我们使用纯 PyTorch 重写了 Meta 原来的 SAM2，损失极小的精度，并高度关注延迟。我们将优化后的推理部署到无服务器基础设施提供商 Modal 上，以证明这些优化的优势可以在更实际的部署环境中实现。

通过利用 AOTInductor (AOTI) 经由 torch.export 进行的提前编译、降低精度、批量处理提示和 GPU 预处理，我们观察到与常规 eager 模式 PyTorch 相比，p90 执行延迟和队列时间最多提高了 13 倍。

对于弹性或自动扩展基础设施，推理服务的副本需要定期自动创建和销毁，简单部署 torch.compile 可能会给推理执行增加工作，这使得更快的模型带来的任何改进都相形见绌。通过利用 AOTInductor (AOTI) 经由 torch.export 进行的提前编译，我们能够提前上传导出的模型并通过网络读取这些数据，这使我们能够在工作量没有显著增加的情况下获得编译带来的好处。

有关如何重现此博客文章中数据的更多详细信息，请查看 torchao 的 experiments 文件夹。如果您遇到任何技术问题，请随时联系我们或提交一个 issue。