LeCun转推,PyTorch GPU内存分配有了火焰图可视化工具

2022 年 10 月 24 日 极市平台
↑ 点击 蓝字  关注极市平台

来源丨机器之心
编辑丨杜伟

极市导读

 

想要了解自己的 PyTorch 项目在哪些地方分配 GPU 内存以及为什么用完吗?不妨试试这个可视化工具。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

近日,PyTorch 核心开发者和 FAIR 研究者 Zachary DeVito 创建了一个新工具(添加实验性 API),通过生成和可视化内存快照(memory snapshot)来可视化 GPU 内存的分配状态。这些内存快照记录了内存分配的堆栈跟踪以及内存在缓存分配器状态中的位置。

接下来,通过将这些内存快照可视化为火焰图(flamegraphs),内存的使用位置也就能一目了然地看到了。

图灵奖得主 Yann Lecun 也转推了这个工具。


Twitter@Zachary DeVito
下面我们来看这个工具的实现原理(以第一人称「我们」描述)。

生成快照

首先,我们必须记录每次分配的堆栈帧信息。

记录这些堆栈跟踪的速度非常快(每次分配约 1 us,正常的 PyTorch 内核调用需要至少 8 us),但我们默认将其关闭。而启用之后,我们可以分配一些内存并拍摄快照。

快照记录了整个分配器的状态,如下所示。

快照是具有以下结构的 Segment 字典列表。

Segments 是直接从 cudaMalloc 请求并由分配器缓存的内存。因此,我们可以只使用这些 segments 中的一部分,缓存分配器将它们分为一个或更多个 Block。所有的块始终保持相同的分配状态。同时,使用_record_memory_history,每个块还将记录一个 History 对象,该对象会记住块中最后一次分配的位置,包括作为 Frames 列表的堆栈跟踪。

对于 active_allocated 块,它其中存在的内容和当前分配的内容将有一个历史记录。对于 inactive 块,可能会有多个条目来记录块内存中最后存在的内容。可能不止一个条目的原因在于分配器在空闲时会合并分割块,并记录下两次拆分的历史。为了避免出现大量的历史记录,我们只保留不与任何更新块重合的块的历史记录。

保存快照

快照因自身设计而可以之后离线查看。

文件_memory_viz.py 可以直接用作交互式命令来处理保存的快照。


可视化快照

_memory_viz.py 工具也可以生成内存的可视化火焰图。

可视化图将分配器中所有的字节(byte)按不同的类来分割成段,如下图所示(原文为可交互视图)。

火焰图可视化是一种将资源(如内存)使用划分为不同类的方法,然后可以进一步细分为更细粒度的类别。

memory 视图很好地展现了内存的使用方式。但对于具体地调试分配器问题,首先将内存分类为不同的 Segment 对象是有用的,而这些对象是分配轨迹的单个 cudaMalloc 段。


比较快照

该可视化器还可以生成显示在两个快照之间添加和删除的段的可视化。例如,我们可以使用更大的输入重新运行模型,并查看分配器如何为更大的临时对象请求更多内存。


比较视图仅显示新段,这有助于找出哪些代码路径提示分配更多内存。

$ python _memory_viz.py compare snapshot.pickle snapshot2.pickle  -o compare.svgonly_before = []only_after = [140636932014080, 140636827156480, 140634912456704, 140634839056384, 140634843250688, 140634841153536, 140634866319360, 140634811793408, 140634845347840, $ 140636806184960, 140636778921984, 140634878902272]

原文链接:https://zdevito.github.io/2022/08/16/memory-snapshots.html

△点击卡片关注极市平台,获取 最新CV干货


极市干货
算法竞赛:往届获奖方案总结以及经验详解|ACCV2022国际细粒度图像分析挑战赛
技术综述 BEV 学术界和工业界方案、优化方法与tricks综述 PyTorch下的可视化工具(网络结构/训练过程可视化)
极视角动态:极视角与华为联合发布基于昇腾AI的「AICE赋能行业解决方案」算法误报怎么办?自训练工具使得算法迭代效率提升50%!

CV技术社群邀请函 #




△长按添加极市小助手
添加极市小助手微信(ID : cvmart2)

备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳)


即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群


极市&深大CV技术交流群已创建,欢迎深大校友加入,在群内自由交流学术心得,分享学术讯息,共建良好的技术交流氛围。


点击阅读原文进入CV社区

收获更多技术干货

登录查看更多
0

相关内容

【经典书】现代C语言教程,408页pdf手写教你写C代码
专知会员服务
62+阅读 · 2022年8月5日
【经典书】多处理器编程艺术,562页pdf
专知会员服务
66+阅读 · 2022年7月10日
代码注释最详细的Transformer
专知会员服务
110+阅读 · 2022年6月30日
【2022新书】Python手册,275页pdf
专知会员服务
179+阅读 · 2022年3月18日
【开放书】Python + Matplotlib可视化指南,249页pdf
专知会员服务
97+阅读 · 2021年11月17日
【资源】100+本免费数据科学书
专知会员服务
107+阅读 · 2020年3月17日
Yann Lecun 纽约大学《深度学习(PyTorch)》课程(2020)PPT
专知会员服务
179+阅读 · 2020年3月16日
三行代码可视化神经网络特征图
极市平台
1+阅读 · 2022年1月4日
实践教程 | 浅谈 PyTorch 中的 tensor 及使用
极市平台
1+阅读 · 2021年12月14日
PyTorch模型训练特征图可视化(TensorboardX)
极市平台
33+阅读 · 2019年6月29日
用PyTorch做物体检测和追踪
AI研习社
12+阅读 · 2019年1月6日
Pytorch作者Adam Paszke-53页Pytorch教程PPT
专知
34+阅读 · 2018年10月27日
可视化循环神经网络的注意力机制
论智
22+阅读 · 2018年9月23日
教程 | 从头开始了解PyTorch的简单实现
机器之心
20+阅读 · 2018年4月11日
国家自然科学基金
1+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
2+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2009年12月31日
国家自然科学基金
0+阅读 · 2009年12月31日
Arxiv
0+阅读 · 2022年12月15日
Arxiv
0+阅读 · 2022年12月15日
Arxiv
17+阅读 · 2021年3月29日
Arxiv
23+阅读 · 2021年3月4日
VIP会员
相关VIP内容
【经典书】现代C语言教程,408页pdf手写教你写C代码
专知会员服务
62+阅读 · 2022年8月5日
【经典书】多处理器编程艺术,562页pdf
专知会员服务
66+阅读 · 2022年7月10日
代码注释最详细的Transformer
专知会员服务
110+阅读 · 2022年6月30日
【2022新书】Python手册,275页pdf
专知会员服务
179+阅读 · 2022年3月18日
【开放书】Python + Matplotlib可视化指南,249页pdf
专知会员服务
97+阅读 · 2021年11月17日
【资源】100+本免费数据科学书
专知会员服务
107+阅读 · 2020年3月17日
Yann Lecun 纽约大学《深度学习(PyTorch)》课程(2020)PPT
专知会员服务
179+阅读 · 2020年3月16日
相关资讯
三行代码可视化神经网络特征图
极市平台
1+阅读 · 2022年1月4日
实践教程 | 浅谈 PyTorch 中的 tensor 及使用
极市平台
1+阅读 · 2021年12月14日
PyTorch模型训练特征图可视化(TensorboardX)
极市平台
33+阅读 · 2019年6月29日
用PyTorch做物体检测和追踪
AI研习社
12+阅读 · 2019年1月6日
Pytorch作者Adam Paszke-53页Pytorch教程PPT
专知
34+阅读 · 2018年10月27日
可视化循环神经网络的注意力机制
论智
22+阅读 · 2018年9月23日
教程 | 从头开始了解PyTorch的简单实现
机器之心
20+阅读 · 2018年4月11日
相关基金
国家自然科学基金
1+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2013年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2012年12月31日
国家自然科学基金
2+阅读 · 2012年12月31日
国家自然科学基金
0+阅读 · 2011年12月31日
国家自然科学基金
0+阅读 · 2009年12月31日
国家自然科学基金
0+阅读 · 2009年12月31日
Top
微信扫码咨询专知VIP会员