展开菜单
首页 精品内容 本月促销 装机必备 Windows PDF教程 macOS软件 IOS软件 Android AI 专题
全部分类

当前位置:

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

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

我们该如何优化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倍

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
作者最新文章
业界资讯
相关文章 更多
精品专题 更多
装机必备

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

Windows

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

PDF教程

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

macOS软件

正软商城macOS软件专区,精选适用于Mac电脑的办公、设计、影音、效率、开发和系统工具,提供软件功能介绍、macOS兼容版本、正版授权及购买下载服务。

IOS软件

正软商城iOS软件专区,精选适用于iPhone和iPad的办公、学习、影音、设计、效率及AI应用,提供功能介绍、适用设备、系统要求和正版获取方式等信息。

AI

正软商城AI软件专区,汇集AI写作、AI绘画、AI视频、AI办公、AI编程、AI翻译、智能客服和数据分析等人工智能工具,提供功能介绍、适用平台、收费方式及正版购买信息。

Mac软件 更多
WPS Office macOS版
WPS Office macOS版

WPS Office是一款面向macOS用户的综合办公套件,将文字处理、电子表格、演示文稿、PDF工具、模板与云协作整合在同一应用中。它兼容常见Microsoft Office格式,支持Intel及Apple芯片Mac,基础功能可免费使用,

Shapr3D macOS版
Shapr3D macOS版

Shapr3D是一款面向工业设计、机械工程、建筑概念和三维打印工作流的CAD软件。Mac版采用Parasolid建模内核,支持草图约束、实体建模、工程图、可视化渲染及常见CAD格式交换,并可通过账户在多台设备之间同步项目。

REAPER macOS版
REAPER macOS版

REAPER是Cockos开发的数字音频工作站,提供多轨音频与MIDI录制、剪辑、处理、混音和母带制作工具。Mac版兼容Intel与Apple芯片,支持AU、VST、VST3、CLAP等插件格式,并提供高度可定制的工作流程。

Ableton Live macOS版
Ableton Live macOS版

Ableton Live 是面向音乐制作人与现场表演者的数字音频工作站,提供编曲视图、独具特色的现场视图、音频录制、MIDI创作、实时变速、乐器及效果器。Mac版原生支持Apple芯片,并可连接音频接口、MIDI控制器和第三方插件。

Adobe After Effects macOS版
Adobe After Effects macOS版

Adobe After Effects 是面向动态图形、影视特效与视频合成的专业创作软件。它提供关键帧动画、遮罩抠像、运动跟踪、三维空间、文字动画和表达式等工具,并可与 Premiere Pro、Photoshop、Illustrator

Adobe Audition macOS版
Adobe Audition macOS版

Adobe Audition 是面向播客、影视后期、音乐制作及内容创作者的专业音频工作站,提供波形精修、多轨录制与混音、频谱编辑、降噪修复、响度匹配和效果处理工具,并可与 Adobe Premiere 协同完成视频声音制作。

Adobe Bridge macOS版
Adobe Bridge macOS版

Adobe Bridge 是面向摄影师、设计师及内容团队的免费数字资产管理工具,可集中预览照片、视频和 Adobe 项目文件,并利用评级、标签、关键词、元数据、筛选与收藏集完成分类检索。它还支持批量重命名、格式导出、相机导入及 Camera

Adobe Illustrator macOS版
Adobe Illustrator macOS版

Adobe Illustrator 是 Adobe 推出的专业矢量图形设计软件,适合在 Mac 上制作标志、图标、插画、包装、信息图及印刷版面。它以路径和锚点构建可无损缩放的图形,并提供钢笔、形状、文字、图像描摹、渐变、画板及生成式功能。

Adobe InDesign macOS版
Adobe InDesign macOS版

Adobe InDesign 是面向印刷品与数字出版物的专业版面设计软件,可在 Mac 上制作书籍、杂志、宣传册、海报、报告及交互式文档。它提供精细的文字样式、网格、母版、长文档管理和印前输出能力。

Adobe Lightroom Classic macOS版
Adobe Lightroom Classic macOS版

Adobe Lightroom Classic 是面向桌面摄影工作流的照片管理与后期处理软件,可将原始照片保存在本地硬盘,通过目录、关键词、评分和收藏夹高效整理图库,并提供 RAW 冲印、蒙版、镜头校正、批量同步及多格式导出等功能。

Adobe Media Encoder macOS版
Adobe Media Encoder macOS版

Adobe Media Encoder 是 Adobe 推出的专业音视频编码与转码工具,可通过队列、预设和监视文件夹批量完成格式转换、代理文件创建及多平台成片输出,并与 Premiere Pro、After Effects 等软件紧密协作。

Adobe Photoshop macOS版
Adobe Photoshop macOS版

Adobe Photoshop 是面向摄影师、设计师和内容创作者的专业图像处理软件,提供图层、蒙版、选区、修复、调色、文字排版、智能对象及生成式编辑等能力。Mac 版支持 Intel 与 Apple Silicon 处理器。

WINDOWS 更多
Adobe Photoshop Windows版
Adobe Photoshop Windows版

Adobe Photoshop 是面向摄影师、设计师和内容创作者的专业图像处理工具,可在 Windows 上完成照片修饰、选区抠图、图层合成、文字排版、绘画及生成式编辑。软件采用订阅授权,由 Creative Cloud 管理安装、更新与许

Adobe InDesign Windows版
Adobe InDesign Windows版

Adobe InDesign 是面向印刷与数字出版的专业页面排版工具,可制作书籍、杂志、宣传册、海报、交互式 PDF 和电子书。它提供精细的文字样式、多页面管理、印前检查及多格式输出能力,并能与 Photoshop、Illustrator

Adobe Illustrator Windows版
Adobe Illustrator Windows版

Adobe Illustrator 是用于绘制矢量图形、插画、图标、徽标和版式的专业设计工具。它以路径和锚点为核心,作品放大后仍能保持清晰,并提供画板、文字、渐变、图像描摹及精确变换等功能,适合品牌视觉、印刷物料、数字界面和商业插画创作。

Adobe Bridge Windows版
Adobe Bridge Windows版

Adobe Bridge 是 Adobe 推出的创意资产管理工具,可集中预览、分类和查找照片、视频、PDF及多种设计文件,并支持元数据、关键词、评级、批量重命名和格式导出。它可与 Photoshop、Illustrator 等应用衔接,

Affinity Photo Windows版
Affinity Photo Windows版

Affinity Photo 是一款面向摄影师、修图师和视觉设计人员的桌面图像编辑工具,提供 RAW 显影、多图层合成、蒙版、实时滤镜、批处理和色彩管理等功能。原独立产品已停止销售,其照片编辑能力已延续至免费的新 Affinity 应用;

Adobe Lightroom Classic Windows版
Adobe Lightroom Classic Windows版

Adobe Lightroom Classic 是面向 Windows 桌面摄影工作流的照片管理与无损编辑工具,可通过目录集中整理大型图库,并完成 RAW 显影、局部蒙版、批量调整、打印和多格式导出。它强调本地文件管理与精细控制,

Figma Windows版
Figma Windows版

Figma 是面向产品设计与开发团队的云端协作平台,可在 Windows 桌面客户端中完成界面绘制、组件管理、交互原型、评论评审和开发交付。文件保存在云端,多名成员能够同步编辑并共享设计系统;官方同时提供免费 Starter 方案及按席位计

Affinity Designer Windows版
Affinity Designer Windows版

Affinity Designer 是面向矢量图形、数字插画和界面设计的专业创作工具,可在矢量、像素与导出工作区之间快速切换。它提供曲线编辑、形状构建、画板、排版、图层和专业色彩管理,并支持常见设计文件的导入及多种印刷、网页格式输出。

Capture One Pro Windows版
Capture One Pro Windows版

Capture One Pro 是面向摄影师和影像团队的专业桌面照片处理工具,提供相机定制色彩配置、RAW转换、图层与蒙版、联机拍摄、目录及会话管理等功能。它适合棚拍、人像、商业产品、婚礼和风光摄影,

Adobe Acrobat Pro Windows版
Adobe Acrobat Pro Windows版

Adobe Acrobat Pro 是面向专业 PDF 工作的综合解决方案,可在 Windows 中创建、编辑、转换、比较、整理和保护文档,并支持表单、电子签名、扫描件文字识别及敏感信息永久遮盖。订阅还包含网页、移动服务和云端协作能力,

Adobe After Effects Windows版
Adobe After Effects Windows版

Adobe After Effects 是面向动态图形、动画与影视视觉特效制作的专业工具,可完成关键帧动画、图层合成、遮罩跟踪、抠像、文字动画和二维三维元素整合。Windows 版通过 Adobe Creative Cloud 安装及更新,

DaVinci Resolve Windows版
DaVinci Resolve Windows版

DaVinci Resolve是Blackmagic Design推出的一体化影视后期工具,将素材管理、非线性剪辑、专业调色、Fusion视觉特效、Fairlight音频处理和成片交付整合在统一工程中。免费版已覆盖多数个人创作需求,