PyTorch模型部署：从训练到生产环境的完整流程

问题背景

在实际项目中，我们遇到一个典型的PyTorch模型部署难题：训练好的ResNet50模型在生产环境中推理速度无法满足实时性要求。经过排查发现，主要瓶颈在于模型转换和推理引擎选择。

解决方案

步骤1：模型优化与量化

import torch
import torch.quantization

# 加载训练好的模型
model = torch.load('resnet50.pth')
model.eval()

# 设置量化配置
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model_fused = torch.quantization.fuse_modules(model, [['conv1', 'bn1', 'relu']])
model_quantized = torch.quantization.prepare(model_fused)

# 运行量化校准
for data, _ in calib_loader:
    model_quantized(data)

# 转换为量化模型
model_prepared = torch.quantization.convert(model_quantized)

步骤2：模型转换为ONNX格式

# 导出ONNX模型
input_tensor = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model_prepared,
    input_tensor,
    "resnet50_quantized.onnx",
    export_params=True,
    opset_version=11,
    do_constant_folding=True
)

步骤3：使用ONNX Runtime进行部署

import onnxruntime as ort

# 加载模型
session = ort.InferenceSession("resnet50_quantized.onnx")
input_name = session.get_inputs()[0].name

# 性能测试
import time
start_time = time.time()
for _ in range(100):
    result = session.run(None, {input_name: input_tensor.numpy()})
end_time = time.time()
print(f"平均推理时间: {(end_time-start_time)/100*1000:.2f}ms")

性能对比数据

方法	模型大小	推理时间(ms)	精度损失
原始FP32	97MB	156ms	0%
量化模型	24MB	45ms	0.8%
ONNX + ORT	24MB	38ms	1.2%

实践建议

建议在生产环境中使用ONNX Runtime进行部署，推理速度提升明显
量化策略要平衡精度与性能
部署前务必进行充分的性能测试

BadApp · 2026-01-08T10:24:58

模型部署不是训练结束的终点，而是真正考验工程能力的起点。我之前也踩过坑，以为量化后速度就上去了，结果忽略了推理框架的选择，最后还是得用TensorRT才能满足线上低延迟要求。

代码工匠 · 2026-01-08T10:24:58

ONNX转换一定要注意输入输出的shape和类型，我之前因为没设置好输入tensor的name，导致部署时直接报错。建议导出前加个print确认一下模型结构，别等到生产环境才发现问题。

CoolCharlie · 2026-01-08T10:24:58

量化虽然能减小模型体积、提升推理速度，但别盲目追求极致量化，比如INT8可能在某些场景下反而拖慢性能。我有个项目最后用的是动态量化，在精度和速度之间找到了平衡点。

KindSilver · 2026-01-08T10:24:58

生产环境部署前一定要做充分的压测，包括不同batch size下的延迟、内存占用等。我在一次上线中发现模型在单线程下跑得挺快，但多线程并发时反而卡死，后来才发现是ONNX Runtime配置没调好

PyTorch模型部署：从训练到生产环境的完整流程