来自｜机器之心 编辑｜杜伟想要了解自己的 PyTorch 项目在哪些地方分配 GPU 内存以及为什么用完吗？不妨试试这个可视化工具。近日，PyTorch 核心开发者和 FAIR 研究者 Zachary DeVito 创建了一个新工具（添加实验性 API），通过生成和可视化内存快照（memory snapshot）来可视化 GPU 内存的分配状态。这些内存快照记录了内存分配的堆栈跟踪以及内存在缓存分配器状态中的位置。接下来，通过将这些内存快照可视化为火焰图（flamegraphs），内存的使用位置也就能一目了然地看到了。图灵奖得主 Yann Lecun 也转推了这个工具。Twitter@Zachary DeVito下面我们来看这个工具的实现原理（以第一人称「我们」描述）。生成快照首先，我们必须记录每次分配的堆栈帧信息。记录这些堆栈跟踪的速度非常快（每次分配约 1 us，正常的 PyTorch 内核调用需要至少 8 us），但我们默认
………………………………
