LeCun转推！PyTorch有了火焰图可视化工具

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—> CV 微信技术交流群

转载自：机器之心  |  编辑：杜伟

想要了解自己的 PyTorch 项目在哪些地方分配 GPU 内存以及为什么用完吗？不妨试试这个可视化工具。

近日，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.svg
only_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 微信技术交流群

CVPR/ECCV 2022论文和代码下载

后台回复：CVPR2022，即可下载CVPR 2022论文和代码开源的论文合集

后台回复：ECCV2022，即可下载CVPR 2022论文和代码开源的论文合集

后台回复：Transformer综述，即可下载最新的3篇Transformer综述PDF

目标检测和Transformer交流群成立
扫描下方二维码，或者添加微信：CVer222，即可添加CVer小助手微信，便可申请加入CVer-目标检测或者Transformer 微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer等。
一定要备注：研究方向+地点+学校/公司+昵称（如目标检测或者Transformer+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer222，进交流群
CVer学术交流群（知识星球）来了！想要了解最新最快最好的CV/DL/ML论文速递、优质开源项目、学习教程和实战训练等资料，欢迎扫描下方二维码，加入CVer学术交流群，已汇集数千人！

▲扫码进群
▲点击上方卡片，关注CVer公众号

整理不易，请点赞和在看