ExecuTorch 中的内存规划检查

在 ExecuTorch 的内存规划 (Memory Planning) 阶段之后，内存分配信息会存储在ExportedProgram 的节点上。这里，我们介绍一个用于检查内存分配和可视化所有活跃张量对象的工具。

用法

用户在调用 to_executorch() 后应添加此代码，它将把存储在节点上的内存分配信息写入文件路径 “memory_profile.json”。该文件与 Chrome trace viewer 兼容；有关解释结果的更多信息，请参阅下文。

from executorch.util.activation_memory_profiler import generate_memory_trace
generate_memory_trace(
    executorch_program_manager=prog,
    chrome_trace_filename="memory_profile.json",
    enable_memory_offsets=True,
)

• prog 是 ExecuTorchProgramManager 的一个实例，由 to_executorch() 返回。

• 将 enable_memory_offsets 设置为 True 以显示每个张量在内存空间中的位置。

Chrome 跟踪

打开一个 Chrome 浏览器标签页并导航到 chrome://tracing/。上传生成的 .json 文件以查看。例如 MobileNet V2 模型

请注意，由于我们在此复用 Chrome trace 工具，因此在此上下文中的坐标轴可能与您之前遇到的其他 Chrome trace 图表具有不同的含义

• 水平轴虽然标记为秒 (s)，但实际代表兆字节 (MBs)。

• 垂直轴具有两层层次结构。第一层，“pid”，代表内存空间。对于 CPU，所有内容都分配在一个“空间”中；其他后端可能具有多个。在第二层，每一行代表一个时间步。由于节点将按顺序执行，每个节点代表一个时间步，因此行数与节点数相同。

延伸阅读

• 内存规划