pytorch-deployment

用途

本技能是 PyTorch 的进阶子技能，专注于模型的生产化与部署。其目标场景是在高性能环境中去除 Python 依赖，把模型序列化为可被 C++ 加载、可为边缘设备优化，或可在高吞吐推理引擎（如 TensorRT）中执行的格式。

适用时机包括：

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

主要能力

脚本化与追踪

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

ONNX 桥接

ONNX（Open Neural Network Exchange）是跨平台格式。导出为 ONNX 的模型可由 Microsoft 的 ONNX Runtime 运行，其推理速度通常快于标准 PyTorch。

量化

将权重从 float32（4 字节）降低为 int8（1 字节）。模型体积缩小为原来的 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'}})

关键规则

应当（DO）：

导出前调用 model.eval()——这会冻结 BatchNorm 与 Dropout 层；遗漏会导致预测结果错误。
使用 torch.no_grad()——始终用 no_grad 上下文包裹导出逻辑，避免保存不必要的梯度追踪元数据。
在 ONNX 中定义 dynamic_axes——若模型需处理不同 batch size 或图像分辨率，必须在导出时指定。
校验导出精度——始终用 torch.allclose() 比较原始 Python 模型与导出模型的输出。
Python 部署使用 torch.compile——若在 Python 内部署，使用 torch.compile（PyTorch 2.0+）而非 JIT，以获得更好性能。

不应（DON'T）：

不要对有 if/else 的模型使用 JIT Tracing——追踪器只会捕获示例运行时走到的分支。
（通常）不要把预处理包含进模型——将图像缩放/归一化保留在核心模型之外以获得更好灵活性，除非使用 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;
}

实践工作流

为移动端优化（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")

通过 ONNX Runtime 部署

import onnxruntime as ort

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

常见陷阱与解决办法

JIT 中的“Missing Attribute”错误

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.

适用场景与限制

适用场景：生产 Web 服务、C++ 应用嵌入、移动端与边缘设备推理、硬件后端加速、跨版本可复现性保障。

限制：资料未提供具体的性能基准数值、各后端版本兼容性细节，也未提供收费、价格或平台支持信息，相关内容暂未提供。