PyTorch可视化工具：GPU内存分配

本文来源机器之心编辑：杜伟

想要了解自己的 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.svgonly_before = []only_after = [140636932014080, 140636827156480, 140634912456704, 140634839056384, 140634843250688, 140634841153536, 140634866319360, 140634811793408, 140634845347840, $ 140636806184960, 140636778921984, 140634878902272]