pytorch-deployment

PyTorch - 部署与生产工程

在高性能环境中部署模型，通常意味着移除对 Python 的依赖。本指南介绍如何将模型序列化为可在 C++ 中加载的格式、针对边缘设备进行优化，或在 TensorRT 等高吞吐量推理引擎中运行。

适用场景

• 将模型从 Jupyter Notebook 迁移到生产环境中的 Web 服务器（FastAPI/Go/Rust）。
• 将神经网络嵌入 C++ 应用程序（LibTorch）。
• 在移动设备（iOS/Android）或边缘硬件（NVIDIA Jetson）上运行推理。
• 使用专用硬件后端（OpenVINO、TensorRT）加速推理。
• 确保模型在不同版本的 PyTorch 中具有可复现性。

核心原则

1. 脚本化与追踪

• 追踪：PyTorch 使用“示例数据”运行一次模型，并记录所有操作。速度快，但会忽略 Python 控制流（if、for）。
• 脚本化：PyTorch 编译模块的 Python 源代码。准备时间较长，但会保留逻辑和控制流。

2. ONNX 桥接

ONNX（开放神经网络交换格式）是一种跨平台格式。导出为 ONNX 的模型可由 Microsoft 的 ONNX Runtime 运行，其推理速度通常快于标准 PyTorch。

3. 量化

将权重从 float32（4 字节）缩减为 int8（1 字节）。这会使模型大小缩小 4 倍，并可在 CPU 上将推理速度提升 2–3 倍。

快速参考：导出模式

import torch

model = MyModel().eval()
example_input = torch.randn(1, 3, 224, 224)

# 1. Tracing (Most common)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_jit.pt")

# 2. Scripting (For dynamic logic)
scripted_model = torch.jit.script(model)
scripted_model.save("model_script.pt")

# 3. ONNX Export
torch.onnx.export(model, example_input, "model.onnx", 
                  input_names=['input'], output_names=['output'],
                  dynamic_axes={'input': {0: 'batch_size'}})

重要规则

✅ 应该做

• 导出前调用 model.eval() - 这会冻结 BatchNorm 和 Dropout 层。忘记这一步会导致预测不正确。
• 使用 torch.no_grad() - 始终在 no_grad 上下文中封装导出逻辑，避免保存不必要的梯度跟踪元数据。
• 在 ONNX 中定义 dynamic_axes - 如果模型需要处理不同的批次大小或图像分辨率，必须在导出时指定这些参数。
• 验证导出精度 - 始终使用 torch.allclose() 比较原始 Python 模型和导出模型的输出。
• 使用 torch.compile 进行 Python 部署 - 如果在 Python 中部署，请使用 torch.compile（PyTorch 2.0+），而不是 JIT，以获得更好的性能。

❌ 不要做

• 不要对包含 if/else 的模型使用 JIT 追踪 - 追踪器只会捕获示例运行时所走的分支。
• 通常不要将预处理纳入模型 - 为了提高灵活性，请将图像缩放和归一化留在核心模型之外，除非使用与 TorchScript 兼容的操作。
• 不要忽略量化警告 - 某些层（如自定义激活函数）不支持 int8，会回退到 float32，从而降低收益。

高级优化

训练后量化（静态）

import torch.quantization

# 1. Set backend (x86 or ARM)
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 2. Prepare and Calibrate (Run some data through the model)
model_prepared = torch.quantization.prepare(model)
# ... run calibration loop ...

# 3. Convert
model_int8 = torch.quantization.convert(model_prepared)

LibTorch（C++ 部署）

在 C++ 中加载 TorchScript 模型：

#include <torch/script.h>

int main() {
    // Load model
    torch::jit::script::Module module = torch::jit::load("model_jit.pt");
    
    // Create input tensor
    auto input = torch::randn({1, 3, 224, 224});
    
    // Run inference
    at::Tensor output = module.forward({input}).toTensor();
    std::cout << output.slice(1, 0, 5) << std::endl;
}

实用工作流

1. 针对移动端优化（Lite Interpreter）

对于移动端部署，标准 TorchScript 过于臃肿。请使用“Mobile”优化器。

from torch.utils.mobile_optimizer import optimize_for_mobile
optimized_model = optimize_for_mobile(traced_model)
optimized_model._save_for_lite_interpreter("model_mobile.ptl")

2. 通过 ONNX Runtime 部署

import onnxruntime as ort

session = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])
outputs = session.run(None, {"input": example_input.numpy()})

常见问题与解决方案

JIT 中的“缺少属性”错误

TorchScript 无法识别模型初始化后添加的属性。

# ✅ Solution: Define all needed attributes in __init__ or use @torch.jit.export

动态形状故障

如果模型在计算中使用 x.shape[0]，追踪可能会将该值硬编码。

# ✅ Solution: Use Scripting or ensure calculations use tensor methods 
# like .size(0) which JIT understands.

PyTorch 部署是连接科学与现实世界的桥梁。掌握这些工具，才能确保你的发现不只停留在笔记本中，还能为下一代智能系统提供动力。