Lei Mao's Log Book
·
检查和可视化Torch FX图
💡
原文英文，约1200词，阅读约需5分钟。
📝
内容提要
本文介绍了如何使用Torch FX对PyTorch模块进行图形检查和可视化，特别是多层感知器（MLP）。通过符号追踪和导出生成图形，并使用FxGraphDrawer进行可视化。同时展示了记录函数调用和调度的方法，以帮助理解模块的结构和操作。
🎯
关键要点
- PyTorch模块可以包含嵌套模块，理解其结构和操作较为复杂。
- Torch FX图是PyTorch模块的中间表示，可以进行检查和可视化。
- 本文展示了如何使用Torch FX的FxGraphDrawer对多层感知器（MLP）模块进行可视化。
- 通过符号追踪和导出生成图形，并使用FxGraphDrawer进行可视化。
- 使用TorchFunctionMode和TorchDispatchMode记录函数调用和调度，以帮助理解模块的结构和操作。
- Torch FX符号追踪生成的图使用高层torch.nn.module描述模块结构和操作。
- torch.export是用于将模型捕获为ExportedProgram对象的API，支持动态控制流。
- 导出的ATen图使用ATen操作符描述模块结构，提供了张量形状信息。
- Core ATen IR是ATen操作符的核心子集，功能完全且无副作用。
- torch.cond和torch.loop支持动态控制流，导出时会将控制流的分支保存为子图。
- FxGraphDrawer可视化导出模型时，子图需要单独保存以便可视化。
❓
延伸问答
Torch FX图是什么？
Torch FX图是PyTorch模块的中间表示，可以用于检查和可视化模块的结构和操作。
如何使用FxGraphDrawer可视化多层感知器（MLP）模块？
可以通过Torch FX的FxGraphDrawer对MLP模块进行可视化，首先进行符号追踪，然后使用FxGraphDrawer生成图形。
TorchFunctionMode和TorchDispatchMode的作用是什么？
TorchFunctionMode和TorchDispatchMode用于记录函数调用和调度，帮助理解模块的结构和操作。
如何导出PyTorch模型为ATen图？
使用torch.export API可以将PyTorch模型导出为ATen图，需提供示例输入以捕获模型执行过程中的操作。
Core ATen IR与ATen图有什么区别？
Core ATen IR是ATen操作符的核心子集，功能完全且无副作用，而ATen图使用更高层的操作符描述模块结构。
如何处理动态控制流的PyTorch模块？
可以使用torch.cond和torch.loop来处理动态控制流，导出时会将控制流的分支保存为子图。
➡️
