PyTorch与ONNX模型导出部署实战

编程小白狼

发布于 2026-11-04 08:55:32

1270

文章被收录于专栏：编程小白狼编程小白狼

从PyTorch训练到ONNX部署，打通AI模型的最后一公里

为什么需要ONNX

作为AI算法工程师，我经常遇到一个尴尬的情况：

PyTorch训练的模型效果很好，但部署到生产环境时遇到了问题。 生产环境可能用的是Java或C++，没有PyTorch运行时，或者对推理性能有严格要求。

ONNX（Open Neural Network Exchange） 就是解决这个问题的——它是一个开放的模型表示格式，可以在不同框架之间自由转换。

一、ONNX核心概念

ONNX的优势

维度	PyTorch原生部署	ONNX部署
框架依赖	需要PyTorch运行时	独立于训练框架
推理性能	一般	更优（算子优化）
跨平台	有限	广泛（CPU/GPU/移动端）
生产支持	需要Python	多种语言Runtime
量化支持	一般	丰富（INT8/FP16）

基本流程

PyTorch模型 (.pth)
     │
     ▼
    torch.onnx.export()
     │
     ▼
  ONNX模型 (.onnx)
     │
     ├──→ ONNX Runtime (CPU/GPU推理)
     ├──→ TensorRT (NVIDIA GPU加速)
     ├──→ OpenVINO (Intel硬件优化)
     └──→ CoreML (Apple设备)

二、PyTorch模型导出为ONNX

基础导出

import torch
import torch.onnx

# 加载模型
model = torch.load("model.pth")
model.eval（)

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)  # 根据模型输入形状调整

# 导出ONNX
torch.onnx.export(
    model,                      # 模型
    dummy_input,                # 示例输入
    "model.onnx",               # 输出文件名
    export_params=True,         # 导出训练参数
    opset_version=18,          # ONNX算子版本
    do_constant_folding=True,   # 常量折叠优化
    input_names=["input"],      # 输入名称
    output_names=["output"],    # 输出名称
    dynamic_axes={
        "input": {0: "batch_size"},   # 动态batch
        "output": {0: "batch_size"},
    },
)

关键参数详解

1. opset_version：ONNX算子集的版本。新版本支持更多算子，但兼容性可能下降。建议使用14-18之间的版本。

2. dynamic_axes：指定哪些维度是动态的。默认所有维度固定，如果部署时输入batch size不固定，必须配置这个。

3. do_constant_folding：常量折叠优化，把计算图中可以预计算的常量部分提前算好，减少推理时的计算量。

常见问题与解决

问题1：算子不支持

错误：RuntimeError: Exporting the operator xxx to ONNX opset version xx is not supported.

解决方案：

# 方案1：升级opset版本
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=18)

# 方案2：使用onnx支持的替代算子
# 在模型定义时用torch.onnx.symbolic_helper注册替代

# 方案3：使用torch.onnx.export的custom_opsets参数
torch.onnx.export(model, dummy_input, "model.onnx",
                  custom_opsets={"custom_domain": 1})

问题2：动态shape导出失败

解决：确保模型中所有操作都支持动态shape。特别是reshape、view等操作，需要确保动态维度正确传递。

# 使用-1而不是具体数值
x = x.reshape(-1, 64, 64)  # ✅ 推荐
# x = x.reshape(batch_size, 64, 64)  # ❌ 不推荐

问题3：控制流（if/for）导出失败

解决：使用torch.jit.script或torch.jit.trace两种方式：

# Tracing方式（推荐）：用示例输入trace一遍
traced_model = torch.jit.trace(model, dummy_input)
torch.onnx.export(traced_model, dummy_input, "model.onnx")

# Script方式：支持控制流
scripted_model = torch.jit.script(model)
torch.onnx.export(scripted_model, dummy_input, "model.onnx")

三、ONNX模型推理

使用ONNX Runtime

ONNX Runtime是微软开源的推理引擎，支持CPU和GPU：

import onnxruntime as ort
import numpy as np

# 创建推理session
session = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]  # GPU优先
)

# 准备输入
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 推理
outputs = session.run([output_name], {input_name: input_data})
print(outputs[0].shape)

环境配置

# CPU版本
pip install onnx onnxruntime

# GPU版本（需要CUDA已安装）
pip install onnx onnxruntime-gpu

# 验证是否启用GPU
python -c "import onnxruntime; print(onnxruntime.get_device())"
# 输出：GPU

性能优化

# 开启ONNX Runtime优化
session_options = ort.SessionOptions()
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session_options.enable_mem_pattern = True
session_options.enable_cpu_mem_arena = True
session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

session = ort.InferenceSession(
    "model.onnx",
    sess_options=session_options,
    providers=["CUDAExecutionProvider"],
)

# 使用IO Binding优化GPU推理
import onnxruntime as ort

io_binding = session.io_binding()
io_binding.bind_cpu_input("input", input_data)
io_binding.bind_output("output")
session.run_with_iobinding(io_binding)
output = io_binding.get_outputs()[0].numpy()

四、ONNX量化

动态量化

from onnxruntime.quantization import quantize_dynamic, QuantType

# INT8动态量化
quantized_model = quantize_dynamic(
    model_input="model.onnx",
    model_output="model_int8.onnx",
    weight_type=QuantType.QInt8,
)

# 效果对比
# 原始：model.onnx - 500MB
# 量化：model_int8.onnx - 125MB（减少75%）
# 速度提升：1.5-2倍

静态量化

from onnxruntime.quantization import quantize_static, QuantType
from onnxruntime.quantization import CalibrationDataReader

# 需要校准数据集
calibration_data = generate_calibration_data()  # 你的校准数据

class DataReader(CalibrationDataReader):
    def __init__(self, data):
        self.data = data
        self.iter = iter(data)
    
    def get_next(self):
        return next(self.iter, None)

quantized_model = quantize_static(
    model_input="model.onnx",
    model_output="model_int8_static.onnx",
    calibration_data_reader=DataReader(calibration_data),
    quant_format=QuantType.QInt8,
)

五、生产级部署架构

使用ONNX Runtime Server

# 启动ONNX Runtime推理服务
onnxruntime_server --model_path model.onnx --address 0.0.0.0 --port 8001

结合FastAPI的高性能服务

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import onnxruntime as ort
import numpy as np

app = FastAPI()

# 启动时加载模型
class ModelServer:
    def __init__(self, model_path):
        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        options.intra_op_num_threads = 4
        self.session = ort.InferenceSession(
            model_path,
            sess_options=options,
            providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
        )
    
    def predict(self, input_data):
        input_name = self.session.get_inputs()[0].name
        output_name = self.session.get_outputs()[0].name
        input_np = np.array(input_data, dtype=np.float32)
        outputs = self.session.run([output_name], {input_name: input_np})
        return outputs[0].tolist()

model_server = ModelServer("model.onnx")

class PredictRequest(BaseModel):
    data: list

class PredictResponse(BaseModel):
    predictions: list

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    try:
        predictions = model_server.predict(request.data)
        return PredictResponse(predictions=predictions)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# 支持批处理
@app.post("/batch_predict")
async def batch_predict(requests: list[PredictRequest]):
    import asyncio
    
    async def single_predict(req):
        return model_server.predict(req.data)
    
    results = await asyncio.gather(*[single_predict(r) for r in requests])
    return {"predictions": results}

# 启动命令
# uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4

六、ONNX模型验证与调试

验证ONNX模型

import onnx

# 加载并验证模型
model = onnx.load("model.onnx")
onnx.checker.check_model(model)
print(f"模型验证通过！")
print(f"输入: {model.graph.input}")
print(f"输出: {model.graph.output}")

精度对比

def compare_accuracy(pytorch_model, onnx_model_path, test_data):
    # PyTorch推理
    with torch.no_grad():
        torch_outputs = pytorch_model(test_data)
    
    # ONNX推理
    session = ort.InferenceSession(onnx_model_path)
    input_name = session.get_inputs()[0].name
    onnx_outputs = session.run(None, {
        input_name: test_data.numpy()
    })
    
    # 比较差异
    diff = torch.abs(torch.tensor(onnx_outputs[0]) - torch_outputs)
    max_diff = diff.max().item()
    mean_diff = diff.mean().item()
    
    print(f"最大差异: {max_diff:.6f}")
    print(f"平均差异: {mean_diff:.6f}")
    
    return max_diff < 1e-3  # 通常差异很小

性能profiling

import time
import onnxruntime as ort

def benchmark(model_path, input_shape, num_runs=100, warmup=10):
    session = ort.InferenceSession(model_path)
    input_data = np.random.randn(*input_shape).astype(np.float32)
    input_name = session.get_inputs()[0].name
    
    # 预热
    for _ in range(warmup):
        session.run(None, {input_name: input_data})
    
    # 基准测试
    start = time.time()
    for _ in range(num_runs):
        session.run(None, {input_name: input_data})
    end = time.time()
    
    avg_time = (end - start) / num_runs * 1000
    print(f"平均推理时间: {avg_time:.2f}ms")
    print(f"吞吐量: {1000 / avg_time:.0f} requests/s")
    
    return avg_time

七、实战案例：图像分类模型部署

完整流程

# 1. 训练或加载PyTorch模型
import torchvision.models as models
model = models.resnet50(pretrained=True)
model.eval（)

# 2. 导出ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "resnet50.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
)

# 3. 优化（可选）
from onnxruntime.transformers import optimizer
optimized_model = optimizer.optimize_model(
    "resnet50.onnx",
    model_type="bert",  # 或使用其他优化器
    num_heads=12,
    hidden_size=768,
)
optimized_model.save_model_to_file("resnet50_optimized.onnx")

# 4. 部署推理服务（见上面的FastAPI示例）

print("✅ 模型导出和优化完成！")
print("原始模型: resnet50.onnx")
print("优化模型: resnet50_optimized.onnx")

八、常见问题排查

问题	原因	解决
导出时算子不支持	ONNX opset版本低	升级opset或注册自定义算子
推理结果差异大	动态shape处理不当	检查reshape/view等操作
GPU推理失败	CUDA版本不匹配	检查onnxruntime-gpu和CUDA版本
推理速度慢	模型没优化	开启graph_optimization_level
模型文件太大	包含大量参数	使用量化（INT8/FP16）

总结

ONNX + ONNX Runtime是将PyTorch模型推向生产环境的成熟方案。核心流程：

导出：torch.onnx.export() 把PyTorch模型转ONNX
优化：graph optimization + 量化
部署：ONNX Runtime推理服务
验证：精度对比 + 性能基准测试

一句话建议：如果你的PyTorch模型需要跨平台、高性能部署，ONNX是目前最成熟、最实用的方案。

你用过ONNX做模型部署吗？有什么经验或教训？欢迎分享！

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2026-07-24，如有侵权请联系 cloudcommunity@tencent.com 删除

pytorch