当前位置：

首页 > 业界资讯 > PyTorch团队重写「分割一切」模型，比原始实现快8倍

PyTorch团队重写「分割一切」模型，比原始实现快8倍

作者：OpenWorld 时间：2023-11-22 来源：正软商城浏览：3

我们该如何优化Meta的「分割一切」模型，PyTorch团队撰写的这篇博客由浅入深的帮你解答。从年初到现在，生成式AI发展迅猛。但很多时候，我们又不得不面临一个难题：如何加快生成式AI的训练、推理等，尤其是在使用PyTorch的情况下。本文PyTorch团队的研究者为我们提供了一个解决方案。文章重点介绍了如何使用纯原生PyTorch加速生成式AI模型，此外，文章还介绍了PyTorch新功能，以及如何组合这些功能的实际示例。结果如何呢？PyTorch团队表示，他们重写了Meta的「分割一切」(SAM)模型，

本文目录

我们该如何优化 Meta 的「分割一切」模型，PyTorch 团队撰写的这篇博客由浅入深的帮你解答。

从年初到现在，生成式 AI 发展迅猛。但很多时候，我们又不得不面临一个难题：如何加快生成式 AI 的训练、推理等，尤其是在使用 PyTorch 的情况下。

本文 PyTorch 团队的研究者为我们提供了一个解决方案。文章重点介绍了如何使用纯原生 PyTorch 加速生成式 AI 模型，此外，文章还介绍了 PyTorch 新功能，以及如何组合这些功能的实际示例。

结果如何呢？PyTorch 团队表示，他们重写了 Meta 的「分割一切」 (SAM) 模型，从而使代码比原始实现快 8 倍，并且没有损失准确率，所有这些都是使用原生 PyTorch 进行优化的。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

博客地址：https://pytorch.org/blog/accelerating-generative-ai/

看完本文，你将了解到：

Torch.compile：PyTorch 模型编译器， PyTorch 2.0 加入了一个新的函数，叫做 torch.compile ()，能够通过一行代码对已有的模型进行加速；
GPU 量化：通过降低运算精度来加速模型；
SDPA（Scaled Dot Product Attention ）：内存高效的注意力实现方式；
半结构化 (2:4) 稀疏性：一种针对 GPU 优化的稀疏内存格式；
Nested Tensor：Nested Tensor 把 {tensor, mask} 打包在一起，将非均匀大小的数据批处理到单个张量中，例如不同大小的图像；
Triton 自定义操作：使用 Triton Python DSL 编写 GPU 操作，并通过自定义操作符注册轻松将其集成到 PyTorch 的各种组件中。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

^{PyTorch 原生特性所带来的吞吐量增加以及减少的内存开销。}

SAM 由 Meta 提出，关于这项研究的更多内容请参考「CV 不存在了？Meta 发布「分割一切」AI 模型，CV 或迎来 GPT-3 时刻」。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

接下来，文章介绍了 SAM 优化过程，包括性能分析、瓶颈识别，以及如何将这些新功能整合进 PyTorch 以解决 SAM 面临的这些问题。除此以外，本文还介绍了 PyTorch 的一些新特性：torch.compile、SDPA、Triton kernels、Nested Tensor 以及 semi-structured sparsity（半结构化稀疏）。

本文内容逐层深入，文章的最后会介绍快速版 SAM，感兴趣的小伙伴可以去 GitHub 上下载，此外，本文还通过 Perfetto UI 对这些数据进行了可视化，以此来阐释 PyTorch 每项特性的应用价值。

GitHub 地址：https://github.com/pytorch-labs/segment-anything-fast

对分割一切模型 SAM 的重写

该研究表示，本文利用的 SAM 基线数据类型为 float32 dtype、batch 大小为 1，使用 PyTorch Profiler 查看内核跟踪的结果如下：

PyTorch团队重写「分割一切」模型，比原始实现快8倍

本文发现 SAM 有两个地方可以优化：

第一个是对 aten::index 的长调用，这是由张量索引操作（例如 []）产生的底层调用导致的。然而实际上 GPU 花费在 aten::index 上的时间相对较低，原因在于 aten::index 在启动两个内核的过程中，两者之间发生了阻塞 cudaStreamSynchronize。这意味着 CPU 会等待 GPU 完成处理，直到启动第二个内核。因而为了优化 SAM，本文认为应该致力于消除导致空闲时间的阻塞 GPU 同步。

第二个是 SAM 在矩阵乘法中花费了大量的 GPU 时间（上图中的深绿色），这在 Transformers 中很常见。如果能够减少 SAM 模型在矩阵乘法上花费的 GPU 时间，我们就可以显着加快 SAM 的速度。

接下来本文用 SAM 的吞吐量 (img/s) 和内存开销 (GiB) 来建立基线。之后就是优化过程了。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

Bfloat16 半精度（加上 GPU 同步和批处理）

为了解决上述问题，即让矩阵乘法花费的时间更少，本文转向 bfloat16。Bfloat16 是常用的半精度类型，通过降低每个参数和激活的精度，能够节省大量的计算时间和内存。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

^{用 bfloat16 替换 padding 类型}

此外，为了移除 GPU 同步，本文发现有两个位置可以优化。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

具体来说（参考上图更容易理解，出现的变量名都在代码中），该研究发现在 SAM 的图像编码器中，有充当坐标缩放器（coordinate scalers）的变量 q_coords 和 k_coords，这些变量都是在 CPU 上分配和处理的。然而，一旦这些变量被用来在 rel_pos_resized 中建立索引，这些索引操作就会自动的将这些变量移动到 GPU 上，这种复制会导致 GPU 同步。为了解决上述问题，该研究注意到可以使用 torch.where 重写这部分内容来解决问题，如上所示。

内核跟踪

在应用了这些更改之后，本文注意到单个内核调用之间有着显著的时间间隔，尤其在小批量（这里为 1）时更为突出。为了更深入的了解这一现象，本文开始对批大小为 8 的 SAM 推理进行性能分析：

PyTorch团队重写「分割一切」模型，比原始实现快8倍

在查看每个内核所花费的时间时，本文观察到 SAM 的大部分 GPU 时间都花费在逐元素内核（elementwise kernels）和 softmax 操作上。

现在可以看到矩阵乘法的相对开销小了很多。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

将 GPU 同步和 bfloat16 优化结合在一起，SAM 性能提高了 3 倍。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

Torch.compile（+graph breaks 和 CUDA graphs）

本文发现在深入研究 SAM 的过程中有很多小的操作，他们认为使用编译器来融合操作有很大的好处，因而 PyTorch 对 torch.compile 做了以下优化：

将 nn.LayerNorm 或 nn.GELU 等操作序列融合成一个单一的 GPU 内核；
融合紧跟在矩阵乘法内核之后的操作，以减少 GPU 内核调用的数量。

通过这些优化，该研究减少了 GPU 全局内存往返次数（roundtrips），从而加快了推理速度。我们现在可以在 SAM 的图像编码器上尝试 torch.compile。为了最大限度地提高性能，本文使用了一些高级编译技术：

PyTorch团队重写「分割一切」模型，比原始实现快8倍

内核跟踪

PyTorch团队重写「分割一切」模型，比原始实现快8倍

结果显示，torch.compile 工作得很好。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

可以观察到 softmax 占了很大一部分时间，然后是各种 GEMM 变体。以下测量的是批大小为 8 及以上的变化。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

SDPA: scaled_dot_product_attention

接下来，本文又对 SDPA（scaled_dot_product_attention）进行了实验，研究的重点是注意力机制。一般来讲，原生注意力机制在时间和内存上随序列长度呈二次方扩展。PyTorch 的 SDPA 操作基于 Flash Attention、FlashAttentionV2 和 xFormer 的内存高效注意力原理构建，可以显着加快 GPU 注意力。与 torch.compile 相结合，这个操作允许在 MultiheadAttention 的变体中表达和融合一个共同的模式。经过一小部分更改后，现在模型可以使用 scaled_dot_product_attention。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

内核跟踪

现在可以看到内存高效的注意力内核占用了 GPU 上大量的计算时间：

PyTorch团队重写「分割一切」模型，比原始实现快8倍

使用 PyTorch 的原生 scaled_dot_product_attention，可以显著增加批处理大小。下图为批大小为 32 及以上的变化。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

之后，该研究又实验了 Triton，NestedTensor 、批处理 Predict_torch， int8 量化，半结构化 (2:4) 稀疏性等操作。

例如本文使用自定义 positional Triton 内核，观察到批大小为 32 的测量结果。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

使用 Nested Tensor，批大小为 32 及以上的变化。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

添加量化后，批大小为 32 及以上变化的测量结果。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

文章的最后是半结构化稀疏性。该研究表示，矩阵乘法仍然是需要面对的一个瓶颈。解决的办法是使用稀疏化来近似矩阵乘法。通过稀疏矩阵（即将值归零）可以使用更少的位来存储权重和激活张量。该研究将张量中哪些权重设置为零的过程称为剪枝。剪枝掉较小的权重可以潜在地减小模型大小，而不会显着损失准确率。

剪枝的方法多种多样，从完全非结构化到高度结构化。虽然非结构化剪枝理论上对精度的影响最小，但 GPU 在进行大型密集矩阵乘法方面尽管非常高效，然而在稀疏情况下可能还会遭受显着的性能下降。PyTorch 最近支持的一种剪枝方法旨在寻求平衡，称为半结构化（或 2:4）稀疏性。这种稀疏存储将原始张量减少了 50%，同时产生密集张量输出。参见下图的说明。

PyTorch团队重写「分割一切」模型，比原始实现快8倍

为了使用这种稀疏存储格式和相关的快速内核，接下来要做的是剪枝权重。本文在 2：4 的稀疏度下选择最小的两个权重进行剪枝，将权重从默认的 PyTorch（“strided”）布局更改为这种新的半结构化稀疏布局很容易。要实现 apply_sparse (model)，只需要 32 行 Python 代码：

PyTorch团队重写「分割一切」模型，比原始实现快8倍

在 2:4 的稀疏度下，本文观察到 vit_b 和批大小为 32 时的 SAM 峰值性能：

最后，一句话总结这篇文章：本文介绍了迄今为止在 PyTorch 上最快的 Segment Anything 实现方式，借助官方发布的一系列新功能，本文在纯 PyTorch 中重写了原始 SAM，并且没有损失准确率。

感兴趣的读者可以查看原博客了解更多内容。

^{参考链接：https://pytorch.org/blog/accelerating-generative-ai/}

本站声明：本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系bd@zhengruan.com

作者最新文章

PDF图片太多怎么批量压缩？PDF文档如何压缩变小？ 2026-10-27 11:48

花4700元买RTX 5070：收到却是一堆20年前的电子垃圾！ 2026-10-26 16:25

Aperty人像精修指南：去除面部红晕斑点与调整阴影 2026-10-22 21:53

即梦AI怎么导出到剪映继续剪辑？ 2026-10-22 19:49

PDF转TXT小说怎么不乱码？手机PDF转TXT用哪个？ 2026-10-17 16:32

业界资讯

上一篇：如何设计可扩展的Java框架：经验与建议

下一篇： Golang开发注意事项：如何避免内存溢出问题

相关文章更多

PyTorch中怎么加载本地权重文件_使用torch.load读取pth文件

使用torch.load加载.pth文件时常见问题包括路径错误、文件损坏、键名不匹配及设备映射不当。需通过os.path.exists检查路径，用map_location指定设备，并清洗键名前缀差异。确保文件完整性和模型结构一致，避免加载失败。

2026-07-19 2 RiverSoul

PyTorch 0.4 Windows支持，API全面升级

前言昨日，PyTorch官方在GitHub上发布了0.4.0版本，备受期待的Windows支持终于实现！现在，你可以轻松加入深度学习的队伍了！PyTorch官方在GitHub上发布了0.4.0版本，这一版本的最大亮点是正式支持Windows系统。新版本对API进行了彻底的重构，是今年最重要的更新之一。GitHub发布地址：https://github.com/pytorch/pytorch/releasesPyTorch官网：http://pytorch.org/此

2025-09-14 4 BlueNight

Llama3训练每3小时崩一次？豆包大模型、港大团队为脆皮万卡训练提效

伴随大模型迭代速度越来越快，训练集群规模越来越大，高频率的软硬件故障已经成为阻碍训练效率进一步提高的痛点，检查点（Checkpoint）系统在训练过程中负责状态的存储和恢复，已经成为克服训练故障、保障训练进度和提高训练效率的关键。近日，字节跳动豆包大模型团队与香港大学联合提出了ByteCheckpoint。这是一个PyTorch原生，兼容多个训练框架，支持Checkpoint的高效读写和自动重新切分的大模型Checkpointing系统，相比现有方法有显著性能提升和易用性优势。本文介绍了大模型训练提效中C

2025-07-11 4 吹吹风会发光

字节跳动

TPAMI 2024 | ProCo: 无限contrastive pairs的长尾对比学习

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com本论文第一作者杜超群是清华大学自动化系2020级直博生。导师为黄高副教授。此前于清华大学物理系获理学学士学位。研究兴趣为不同数据分布上的模型泛化和鲁棒性研究，如长尾学习，半监督学习，迁移

2025-06-13 6 吹吹风会发光

亚马逊云创新「神经稀疏检索」：仅需要文本匹配就能实现语义搜索

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com本文作者是来自OpenSearch中国研发团队的机器学习负责人杨扬博士以及机器学习工程师耿志超和管聪。OpenSearch是一个由亚马逊云科技发起的纯开源搜索和实时分析引擎项目。目前软件

2025-06-12 7 RiverSoul

将图像自动文本化，图像描述质量更高、更准确了

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com皮仁杰：香港科技大学博士三年级学生，师从张潼教授和周晓方教授。于2024年获得苹果奖学金。目前主要研究方向是多模态大语言模型和数据为中心的AI。张鉴殊：武汉大学本科三年级学生，目前在张潼

2025-05-24 5 RiverSoul

告别逐一标注，一个提示实现批量图片分割，高效又准确

仅需一个任务描述，即可一键分割所有图片！SegmentationAnythingModel(SAM)的提出在图像分割领域引起了巨大的关注，其卓越的泛化性能引发了广泛的兴趣。然而，尽管如此，SAM仍面临一个无法回避的问题：为了使SAM能够准确地分割出目标物体的位置，每张图片都需要手动提供一个独特的视觉提示。如下图所示，即使点击的是同一物体（图(b)-(d)），微小位置变化都会导致分割结果的显著差异。这是因为视觉提示语义信息的缺乏，即使提供了想要分割的目标物体上，仍然可能引发歧义。框提示和涂鸦提示（图(e)(

2025-05-23 5 慢热型

CVPR 2024 | 分割一切模型SAM泛化能力差？域适应策略给解决了

第一个针对「SegmentAnything」大模型的域适应策略来了！相关论文已被CVPR2024接收。引言大语言模型（LLMs）的成功激发了计算机视觉领域探索分割基础模型的兴趣。这些基础分割模型通常通过PromptEngineer来进行zero/few图像分割。其中，SegmentAnythingModel（SAM）是最先进的图像分割基础模型。图SAM在多个下游任务上表现不佳但是最近的研究表明，SAM在多种下游任务中并非具有很强的鲁棒性与泛化性，例如在医学图像、伪装物体、添加干扰的自然图像等领域表现较差。

2025-05-14 4 NorthPath

ECCV 2024 | 提升GPT-4V、Gemini检测任务性能，你需要这种提示范式

AIxiv专栏是本站发布学术、技术内容的栏目。过去数年，本站AIxiv专栏接收报道了2000多篇内容，覆盖全球各大高校与企业的顶级实验室，有效促进了学术交流与传播。如果您有优秀的工作想要分享，欢迎投稿或者联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com本文作者来自浙江大学、上海人工智能实验室、香港中文大学、悉尼大学和牛津大学。作者列表：吴逸璇，王逸舟，唐诗翔，吴文灏，贺通，WanliOuyang，PhilipTorr，JianWu。其

2025-05-13 5 SoftHope

贾佳亚团队联手剑桥清华等共推评测新范式一秒侦破大模型“高分低能”

颠覆过往大模型评测标准，最新、最全、最权威的测评数据集MR-Ben来了！这是继今年4月发布堪称GPT-4+DALL-E-3的王炸产品超强视觉语言模型Mini-Gemini后，港中文贾佳亚团队再次提出的极具代表性的作品。在MR-Ben的“监督”下，大模型不仅要像学生那样会答题，还要像老师那样会阅卷，真实的推理能力无所遁形。MR-Ben细致地评测了不少国内外一线的开源和闭源模型，如GPT4-Turbo、Cluade3.5-Sonnet、Mistral-Large、Zhipu-GLM4、

2025-05-11 6 EasyWind

热门文章更多

Photoshop蒙版换证件照背景颜色怎么改？实用操作步骤详解 AutoCAD绘制机械零件三维实体详细教程 AE教程书籍挑选指南：零基础、动效与合成方向实战标准 Photoshop图层合并快捷键是什么？合并盖印及分开图层方法 Photoshop日系调色教程：小清新风格参数与思路

精品专题更多

装机必备

正软商城装机必备专区，精选办公、浏览器、安全防护、影音播放、压缩解压、设计创作和系统工具等电脑常用正版软件，帮助用户快速完成新电脑软件配置。

Windows

正软商城Windows软件专区，汇集适用于Windows电脑的办公、设计、安全防护、影音播放、开发工具和系统优化软件，提供软件介绍、系统要求、正版授权及购买下载服务。

PDF教程

正软商城PDF教程频道提供PDF编辑、转换、合并、拆分、压缩及格式处理方法，同时介绍常用PDF软件和工具的使用技巧。

Mac软件更多