企业官网建设流程全解析

PyTorch Lightning 生产部署实战：使用 torch.export 导出与运行 LightningModule（进阶指南）

【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

本篇技术指南聚焦于 PyTorch Lightning 项目（gh_mirrors/py/pytorch-lightning）中面向企业级生产环境的模型导出方案：使用 PyTorch 官方推荐的torch.export.export()API 将任意LightningModule捕获为干净的中间表示（IR）并保存、加载、运行。读完本文，你将掌握 Lightning 模型从训练态到生产部署态的完整导出链路，包括基础模型导出、predict_step等自定义方法的导出，以及它与to_onnx、to_tensorrt等 Lightning 内置导出 API 的定位差异与选型依据。

为什么生产部署要选择 torch.export

在生产环境中，我们需要的不是"能跑"，而是"可靠、可优化、可跨平台"。torch.export是 PyTorch 官方推荐的模型捕获（capture）方式，它通过一次完整的图捕获（graph capture）将模型转换为一个干净、结构化的中间表示，并附带较强的健全性保证（soundness guarantees）——即导出的程序能够忠实反映原始模型的语义，不会出现torch.jit.trace时代常见的"静默行为漂移"问题。

这份干净的 IR 具备两个直接价值：

• 推理优化：导出的ExportedProgram可以被 AOTInductor、TensorRT 等后端进一步编译优化；
• 跨平台部署：IR 不绑定 Python 执行环境，可配合torch.export.save/load实现模型在独立生产环境中的序列化与恢复。

对于 Lightning 用户来说，最核心的一点是：任意LightningModule都可以直接通过torch.export.export()导出，无需改写网络结构，也无需依赖 Lightning 运行时。这与 Lightning 生态中其他导出路径（见下文"与 Lightning 内置导出 API 的定位对比"）形成互补。

环境前提：PyTorch 版本要求

torch.export相关能力由 PyTorch 提供。当前仓库的依赖声明位于 requirements/pytorch/base.txt，其中约束了：

torch >=2.6.0, <2.14.0

原文档同时强调：建议安装最新受支持的 PyTorch 版本再使用本特性，以免某些导出能力受到版本限制。在动手之前，请先确认你的运行环境满足该版本区间。

第一步：导出基础 LightningModule

导出一个 Lightning 模型与导出一个普通nn.Module在 API 层面完全一致：实例化模型、构造示例输入、调用torch.export.export()，最后用torch.export.save()落盘。

以下示例完整取自 docs/source-pytorch/deploy/production_advanced_2.rst，模型是一个典型的 Lightning 分类器——单层线性层加 ReLU 激活：

import torch from torch.export import export class SimpleModel(LightningModule): def __init__(self): super().__init__() self.l1 = torch.nn.Linear(in_features=64, out_features=4) def forward(self, x): return torch.relu(self.l1(x.view(x.size(0), -1))) # 创建模型与示例输入 model = SimpleModel() example_input = torch.randn(1, 64) # 导出模型 exported_program = export(model, (example_input,)) # 保存，供生产环境使用 torch.export.save(exported_program, "model.pt2")

几个要点：

• example_input用于驱动图捕获，其形状与 dtype 会被固化进导出结果中，因此请使用与线上推理一致的输入规格构造示例张量；
• 导出结果是ExportedProgram对象，torch.export.save将其序列化到.pt2文件（这是 torch.export 生态约定的后缀，原文档示例即采用model.pt2）；
• 若模型内部依赖training状态（如 Dropout、BatchNorm），建议在导出前调用model.eval()，确保捕获到的是推理语义的图。这与 Lightning 内置to_torchscript的"导出后置于 eval 模式"的行为保持一致（见 src/lightning/pytorch/core/module.py 中相关实现）。

第二步：在生产环境加载并运行

导出完成后，生产侧只需要 PyTorch 与.pt2文件，不再需要训练代码甚至 Lightning。加载与推理同样简单：

inp = torch.rand(1, 64) loaded_program = torch.export.load("model.pt2") output = loaded_program.module()(inp)

执行流程拆解：

1. torch.export.load("model.pt2")反序列化出ExportedProgram；
2. loaded_program.module()取出可调用的模块实例；
3. 直接以与导出时一致的输入形状调用，获得输出张量。

需要留意，导出的 IR 是静态图，调用时输入形状需与导出时的示例输入匹配（如这里的(1, 64)），动态形状场景需要额外的动态维度标注支持。

进阶：导出 LightningModule 的特定方法（以 predict_step 为例）

torch.export.export()默认导出的是模型的forward。但在推理管线中，我们往往需要导出的是 Lightning 的predict_step钩子——它可能包含了推理前处理、多次采样取均值等"生产逻辑"，而这些逻辑并不存在于forward中。

LightningModule在 src/lightning/pytorch/core/module.py 中提供了predict_step的默认实现（委托给forward），子类可覆写它以注入自定义推理逻辑。原文档给出的方案是：用一个 lambda 包装器把predict_step的调用语义原样捕获进导出图。

以下示例演示了导出包含 Monte Carlo Dropout 推理逻辑的模型——predict_step会开启 Dropout 并多次前向取平均，以实现不确定性估计：

class LitMCdropoutModel(L.LightningModule): def __init__(self, model, mc_iteration): super().__init__() self.model = model self.dropout = nn.Dropout() self.mc_iteration = mc_iteration def predict_step(self, batch, batch_idx): # 开启 Monte Carlo Dropout self.dropout.train() # 对 `self.mc_iteration` 次迭代取平均 pred = [self.dropout(self.model(x)).unsqueeze(0) for _ in range(self.mc_iteration)] pred = torch.vstack(pred).mean(dim=0) return pred model = LitMCdropoutModel(...) example_batch = torch.randn(32, 10) # 示例输入 # 导出 predict_step 方法 exported_program = torch.export.export( lambda batch, idx: model.predict_step(batch, idx), (example_batch, 0) ) torch.export.save(exported_program, "mc_dropout_model.pt2")

这里的技巧在于：

• predict_step的签名是(batch, batch_idx)，与forward不同，因此需要以 lambda 显式包装；
• 导出的输入元组必须同时给出batch（示例批次）与batch_idx（这里用整数0占位）；
• mc_iteration等超参数在构造时已固化进模型，因此多次采样逻辑会被完整捕获进图中，生产侧调用者无需感知。

这个模式可以推广到任意自定义方法：凡是需要在推理阶段执行、且不属于forward的逻辑（如后处理、集成推理、去归一化），都可以用同样的 lambda 包装方式导出。

与 Lightning 内置导出 API 的定位对比

虽然torch.export是本次主题的主角，但理解它在 Lightning 导出家族中的位置有助于正确选型。LightningModule在 src/lightning/pytorch/core/module.py 中维护了多条导出路径：

从源码注释可以看到，TorchScript 已被 PyTorch 官方弃用，to_torchscript在 v2.7 起标记弃用并将于 v2.8 移除，其替代方案正是本文的torch.export.export()。因此对新项目而言，选择 torch.export 不仅是为了更好的 IR，也是跟随上游演进方向的必然选择。

三条导出路径在仓库测试目录中均有对应的回归验证：tests/tests_pytorch/models/test_onnx.py、tests/tests_pytorch/models/test_torchscript.py、tests/tests_pytorch/models/test_torch_tensorrt.py，可作为阅读导出行为细节的参考。

生产部署要点小结

1. 导出时机：训练收敛后、评估通过后再导出，避免导出未收敛权重；导出前将模型置于eval()模式（MC Dropout 等特殊场景除外）。
2. 输入规格：示例输入的 shape/dtype 决定线上推理的固定输入形状；如果线上输入维度可变，请确认所用 PyTorch 版本对动态形状的支持程度。
3. 逻辑归属：推理期需要的自定义逻辑（前/后处理、多次采样）应放进predict_step或专用方法，再以 lambda 包装导出，而不是混入forward。
4. 运行时依赖：.pt2文件 + PyTorch 即可运行，生产环境无需 Lightning 依赖——这与同系列文档 docs/source-pytorch/deploy/production_intermediate.rst 中"去掉 Lightning 依赖、纯 PyTorch 推理"的指导思想一脉相承。
5. 版本锁定：以 requirements/pytorch/base.txt 的torch >=2.6.0为基线，并尽量使用最新受支持的 PyTorch 版本，以获得最完整的torch.export能力。

通过本文的导出、保存、加载三步流程，你可以将任意LightningModule（包括带自定义predict_step的复杂模型）转化为可独立运行的生产推理单元；再结合 Lightning 的 ONNX / TensorRT 路径，即可构建覆盖 CPU 服务、跨平台推理、GPU 编译加速的企业级部署矩阵。

【免费下载链接】pytorch-lightningPretrain, finetune ANY AI model of ANY size on 1 or 10,000+ GPUs with zero code changes.项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning