LeCun转推，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.svg

only_before = []

only_after = [140636932014080, 140636827156480, 140634912456704, 140634839056384, 140634843250688, 140634841153536, 140634866319360, 140634811793408, 140634845347840, $ 140636806184960, 140636778921984, 140634878902272]

原文链接：