使用全新 `torch.export` API 简化 PyTorch 模型部署
在启动一个新的人工智能或机器学习项目时，大部分焦点自然会集中在宏伟的任务上：整理海量数据集、设计复杂的模型以及获取强大的 GPU 集群进行训练。然而，往往是看似微小的细节会成为意想不到的绊脚石，导致令人沮丧的错误和严重的生产延迟。一个典型的例子就是将训练好的模型从开发环境移交到推理环境。虽然这一步骤可能看起来很简单，但运行时库、硬件配置和版本控制的差异，都可能使其变成一个相当大的难题。开发人员必须确保模型的定义及其训练权重能够正确加载，最重要的是，其行为保持不变。
传统上，针对这种关键模型捕获和部署，主要采用了两种方法。第一种也是最简单的方法，是只使用 torch.save 保存模型的权重。这种方法提供了最大的灵活性，允许在推理环境中进行特定于机器的优化。然而，它要求在部署环境中明确重新定义模型的架构，这可能会带来版本控制的噩梦和依赖不匹配，尤其是在对运行时库控制有限的受限环境中。定义和权重的分离常常成为“丑陋错误”的温床，需要严格的版本管理。
多年来，更全面的解决方案是 TorchScript，它将模型定义和权重捆绑到可序列化的图表示中。TorchScript 提供了两种不同的功能：torch.jit.script 和 torch.jit.trace。Scripting 对源代码执行静态分析，能够捕获条件控制流和动态输入形状等复杂元素。相比之下，Tracing 记录模型在样本输入上的实际执行路径，使其不易出现某些故障，但无法处理动态行为。通常，需要两者的结合，但即便如此，TorchScript 仍经常难以处理复杂模型，需要费力且侵入性的代码重写以确保兼容性。我们自己对 HuggingFace 图像到文本生成模型的实验就证明了这一局限性：虽然固定输入的编码器可以被追踪，但动态形状的解码器在不大量修改底层库代码的情况下始终无法进行脚本化处理。
torch.export 应运而生，它是 PyTorch 用于模型捕获的全新、更强大的解决方案。与 torch.jit.trace 类似，torch.export 通过追踪模型的执行来操作。然而，它通过支持动态性和条件控制流，显著改进了其前身，克服了 TorchScript 许多历史局限性。其输出是一种中间图表示，称为 Export IR，可以作为独立的 PyTorch 程序加载和运行，且依赖项极少。torch.export 的一个关键优势是其与 torch.compile 的兼容性，允许在推理环境中进行进一步的即时优化，这是 TorchScript 模型不具备的能力。此功能由 PyTorch 图编译解决方案的核心组件 Torch Dynamo 提供支持。
尽管功能强大，torch.export 仍然是一个原型功能，并带来了一系列挑战。一个常见的障碍是“图中断”（graph break），当导出函数遇到无法追踪的 Python 代码时就会发生。与模型编译不同，PyTorch 可能会回退到 eager execution，而 torch.export 严格禁止图中断，要求开发人员重写代码以绕过它们。调试导出的图也可能很棘手；虽然它们表现得像标准的 torch.nn.Module 对象，但传统调试器无法步入其编译后的 forward 函数。问题通常出现在导出环境中的变量在不经意间作为常量“硬编码”到图中，从而导致在不同环境中的运行时错误。例如，我们导出的解码器最初在 GPU 上失败，原因是其导出环境中的 CPU 设备引用被硬编码，需要手动“猴子补丁”（monkey-patching）HuggingFace 库才能解决。虽然这对于我们的玩具模型有效，但对于生产系统而言，未经大量测试的此类侵入性修改是不可取的。
在我们的示例模型上进行测试时，torch.export 成功捕获了编码器和解码器，没有遇到图中断，这是对 TorchScript 的重大改进。将更正后的导出模型部署到 Amazon EC2 实例显示，与原始模型相比，推理时间适度加快了 10.7%。有趣的是，在这种特定情况下，将 torch.compile 应用于导出的模型，虽然前景广阔，但意外地增加了执行时间，这凸显了需要仔细调整编译参数。
总而言之，torch.export 代表着 PyTorch 模型部署的重大飞跃。它对复杂模型表现出卓越的支持，能够捕获以前困扰 TorchScript 的架构。导出的模型高度可移植，无需大量包依赖即可独立执行，并通过 torch.compile 与强大的机器特定优化兼容。然而，作为一个快速发展的原型，它目前仍存在局限性，包括将意外的环境特定值硬编码到图中的可能性，以及一套尚处于萌芽阶段的调试工具。尽管存在这些不足，torch.export 比以前的解决方案有了实质性改进，为简化 AI 模型开发的最后关键一英里带来了巨大希望。