torch.export
💡
原文英文，约600词，阅读约需3分钟。
📝
内容提要
torch.export()用于提前编译Python可调用对象，生成功能图。它利用TorchDynamo进行字节码追踪，AOT Autograd功能化图形，并通过torch.fx.graph提供图形表示。ExportedProgram包含GraphModule和图签名，支持动态形状和序列化。使用export_for_training()可处理非功能操作，但存在无法追踪的代码和缺失的假内核等限制。
🎯
关键要点
- torch.export()用于提前编译Python可调用对象，生成功能图。
- torch.export()内部使用TorchDynamo进行字节码追踪，AOT Autograd功能化图形，并通过torch.fx.graph提供图形表示。
- ExportedProgram包含GraphModule和图签名，支持动态形状和序列化。
- GraphModule将每个指令编译为低级ATen操作，ATen是PyTorch的核心张量库。
- ExportedProgram的图签名是功能性的，给定相同输入总会产生相同输出。
- 严格模式和非严格模式都将模型编译为torch.fx.Graph，非严格模式需要Python运行时，严格模式通过torch.Dynamo检查字节码。
- 使用export_for_training()可以处理非功能操作，适用于训练场景。
- 支持动态形状，使用Dim()生成编译期间的范围约束。
- torch.export.save()用于保存模型，torch.export.load()用于加载模型。
- 在导出期间，某些值（如输入形状和Python原始类型）被固定为常量，简化图形。
- torch.export存在一些限制，如无法追踪的代码和缺失的假内核可能导致失败。
➡️
