TorchScript 导出与模型部署
模型部署需要将训练好的 PyTorch 模型转换为可在生产环境高效运行的格式。TorchScript 和 ONNX 是两种主流方案。
TorchScript 导出
python
# Tracing 模式
model = MyModel().eval()
example_input = torch.randn(1, 3, 224, 224)
scripted = torch.jit.trace(model, example_input)
scripted.save("model.pt")
# Scripting 模式（支持控制流）
scripted = torch.jit.script(model)
scripted.save("model.pt")
# ONNX 导出
torch.onnx.export(
model, example_input, "model.onnx",
input_names=["input"], output_names=["output"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)Tracing vs Scripting
Tracing 记录一次前向传播的操作，不支持数据依赖的控制流；Scripting 分析 Python 源码，支持 if/for 等控制流。优先使用 Tracing，控制流复杂时用 Scripting。