ONNX模型部署指南

1. ONNX简介

ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式，旨在为AI模型提供一个统一的表示方式。它允许开发者在不同的深度学习框架之间转换模型，从而实现模型的可移植性和互操作性。

1.1 ONNX的主要特点

• 跨框架兼容性：支持从PyTorch、TensorFlow、Keras、MXNet等框架转换模型
• 硬件加速支持：可以在各种硬件平台上优化执行
• 广泛的操作符支持：包含丰富的神经网络操作符集
• 可扩展性：允许添加自定义操作符和属性

1.2 ONNX生态系统

ONNX生态系统包括多个组件：

• ONNX格式规范
• ONNX Runtime执行引擎
• 模型转换工具
• 模型优化工具
• 推理加速库

2. 模型转换为ONNX格式

2.1 从PyTorch转换

import torch
import torchvision.models as models

# 加载预训练模型
model = models.resnet50(pretrained=True)
model.eval()

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX模型
torch.onnx.export(model,               # 要转换的模型
                  dummy_input,         # 模型输入
                  "resnet50.onnx",     # 输出文件名
                  export_params=True,  # 存储训练好的参数权重
                  opset_version=11,    # ONNX算子集版本
                  do_constant_folding=True,  # 是否执行常量折叠优化
                  input_names=['input'],     # 输入节点名称
                  output_names=['output'],   # 输出节点名称
                  dynamic_axes={'input': {0: 'batch_size'},  # 动态轴
                                'output': {0: 'batch_size'}})

2.2 从TensorFlow/Keras转换

import tensorflow as tf
from tensorflow import keras
import tf2onnx
import onnx

# 加载Keras模型
model = keras.applications.MobileNetV2(weights='imagenet', include_top=True)

# 方法1：使用tf2onnx
model_proto, _ = tf2onnx.convert.from_keras(model, opset=11)
onnx.save(model_proto, "mobilenet_v2.onnx")

# 方法2：使用TensorFlow SavedModel中转
model_path = "saved_model"
tf.saved_model.save(model, model_path)
# 然后使用命令行转换：
# python -m tf2onnx.convert --saved-model saved_model --output mobilenet_v2.onnx

2.3 验证ONNX模型

import onnx

# 加载ONNX模型
model = onnx.load("model.onnx")

# 检查模型结构
onnx.checker.check_model(model)

# 打印模型信息
print(onnx.helper.printable_graph(model.graph))

3. ONNX Runtime部署

ONNX Runtime是微软开发的高性能推理引擎，支持在多种平台上部署ONNX模型。

3.1 Python环境部署

import onnxruntime as ort
import numpy as np
from PIL import Image
import torchvision.transforms as transforms

# 加载ONNX模型
session = ort.InferenceSession("resnet50.onnx")

# 准备输入数据
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 图像预处理
img = Image.open("sample.jpg")
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(img)
input_batch = input_tensor.unsqueeze(0).numpy()  # 创建一个批次

# 执行推理
result = session.run([output_name], {input_name: input_batch})
output = result[0]

# 处理输出结果
probabilities = np.exp(output) / np.sum(np.exp(output), axis=1, keepdims=True)
print(f"Top prediction: {np.argmax(probabilities)}")

3.2 C++环境部署

#include <onnxruntime_cxx_api.h>
#include <vector>
#include <iostream>

int main() {
    // 创建ONNX Runtime环境
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "onnx_model_example");
    Ort::SessionOptions session_options;
    session_options.SetIntraOpNumThreads(1);
    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

    // 加载ONNX模型
    const char* model_path = "model.onnx";
    Ort::Session session(env, model_path, session_options);

    // 获取模型输入输出信息
    Ort::AllocatorWithDefaultOptions allocator;
    const char* input_name = session.GetInputName(0, allocator);
    const char* output_name = session.GetOutputName(0, allocator);

    // 准备输入数据
    std::vector<float> input_data(1 * 3 * 224 * 224);  // 根据模型输入形状调整
    // 填充input_data...

    // 设置输入形状
    std::vector<int64_t> input_shape = {1, 3, 224, 224};  // 批次大小, 通道数, 高度, 宽度

    // 创建输入tensor
    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, input_data.data(), input_data.size(),
        input_shape.data(), input_shape.size());

    // 执行推理
    std::vector<Ort::Value> outputs = session.Run(
        Ort::RunOptions{nullptr},
        &input_name, &input_tensor, 1,
        &output_name, 1);

    // 处理输出结果
    float* output_data = outputs[0].GetTensorMutableData<float>();
    size_t output_size = outputs[0].GetTensorTypeAndShapeInfo().GetElementCount();

    // 找出最大概率的类别
    int max_class_id = 0;
    float max_prob = output_data[0];
    for (size_t i = 1; i < output_size; i++) {
        if (output_data[i] > max_prob) {
            max_prob = output_data[i];
            max_class_id = i;
        }
    }

    std::cout << "Predicted class ID: " << max_class_id << std::endl;

    return 0;
}

3.3 移动端部署

Android (Java)

import ai.onnxruntime.*;

public class ONNXModelRunner {
    private OrtEnvironment env;
    private OrtSession session;

    public ONNXModelRunner(String modelPath) throws OrtException {
        // 初始化ONNX Runtime环境
        env = OrtEnvironment.getEnvironment();
        
        // 配置会话选项
        OrtSession.SessionOptions sessionOptions = new OrtSession.SessionOptions();
        sessionOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT);
        
        // 加载模型
        session = env.createSession(modelPath, sessionOptions);
    }

    public float[] runInference(float[] inputData, long[] inputShape) throws OrtException {
        // 创建输入Tensor
        OnnxTensor inputTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), inputShape);
        
        // 准备输入映射
        Map<String, OnnxTensor> inputs = new HashMap<>();
        inputs.put(session.getInputNames().iterator().next(), inputTensor);
        
        // 执行推理
        OrtSession.Result results = session.run(inputs);
        
        // 获取输出
        float[] outputData = ((OnnxTensor) results.get(0)).getFloatBuffer().array();
        
        // 释放资源
        inputTensor.close();
        results.close();
        
        return outputData;
    }

    public void close() throws OrtException {
        session.close();
        env.close();
    }
}

iOS (Swift)

import CoreML
import Vision

class ONNXModelHandler {
    private var model: VNCoreMLModel
    
    init(modelURL: URL) throws {
        // 加载CoreML模型（可以从ONNX转换为CoreML格式）
        let coreMLModel = try MLModel(contentsOf: modelURL)
        model = try VNCoreMLModel(for: coreMLModel)
    }
    
    func predict(image: CGImage, completion: @escaping (VNClassificationObservation?) -> Void) {
        // 创建图像分析请求
        let request = VNCoreMLRequest(model: model) { request, error in
            guard let results = request.results as? [VNClassificationObservation],
                  let topResult = results.first else {
                completion(nil)
                return
            }
            completion(topResult)
        }
        
        // 配置请求
        request.imageCropAndScaleOption = .centerCrop
        
        // 执行请求
        let handler = VNImageRequestHandler(cgImage: image)
        try? handler.perform([request])
    }
}

4. 模型优化技术

4.1 量化

量化是将模型的浮点权重转换为低精度表示（如int8）的过程，可以显著减小模型大小并提高推理速度。

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# 动态量化
model_fp32 = "model_fp32.onnx"
model_quant = "model_quant.onnx"
quantize_dynamic(model_fp32, model_quant, weight_type=QuantType.QInt8)

4.2 图优化

import onnx
from onnxruntime.transformers import optimizer

# 加载模型
model = onnx.load("model.onnx")

# 优化模型
optimized_model = optimizer.optimize_model(
    "model.onnx",
    model_type="bert",  # 模型类型
    num_heads=12,       # 注意力头数量
    hidden_size=768     # 隐藏层大小
)

# 保存优化后的模型
optimized_model.save_model_to_file("optimized_model.onnx")

4.3 模型剪枝

模型剪枝是移除模型中不重要的权重或神经元，以减小模型大小并提高推理速度。

# 使用ONNX简化工具
import onnx
from onnxsim import simplify

# 加载ONNX模型
model = onnx.load("model.onnx")

# 简化模型
simplified_model, check = simplify(model)
assert check, "简化模型检查失败"

# 保存简化后的模型
onnx.save(simplified_model, "simplified_model.onnx")

5. 部署案例分析

5.1 Web应用部署

使用ONNX.js在浏览器中部署模型：

<!DOCTYPE html>
<html>
<head>
    <title>ONNX.js Demo</title>
    <script src="https://cdn.jsdelivr.net/npm/onnxjs/dist/onnx.min.js"></script>
</head>
<body>
    <input type="file" id="imageUpload" accept="image/*">
    <div id="result"></div>

    <script>
        async function runONNXModel() {
            // 创建ONNX会话
            const session = new onnx.InferenceSession();
            // 加载模型
            await session.loadModel("model.onnx");
            
            // 获取上传的图像
            const imageUpload = document.getElementById('imageUpload');
            imageUpload.addEventListener('change', async (e) => {
                const file = e.target.files[0];
                const img = new Image();
                img.onload = async () => {
                    // 预处理图像
                    const tensor = preprocessImage(img);
                    
                    // 运行推理
                    const outputMap = await session.run([tensor]);
                    const outputTensor = outputMap.values().next().value;
                    const predictions = outputTensor.data;
                    
                    // 显示结果
                    const resultDiv = document.getElementById('result');
                    resultDiv.innerHTML = `预测结果: ${predictions}`;
                };
                img.src = URL.createObjectURL(file);
            });
        }

        function preprocessImage(img) {
            // 图像预处理逻辑
            // ...
            return tensor;
        }

        // 初始化
        runONNXModel();
    </script>
</body>
</html>

5.2 边缘设备部署

在Raspberry Pi上部署ONNX模型：

# 安装ONNX Runtime
pip install onnxruntime

# 对于ARM设备，可能需要从源码编译
git clone --recursive https://github.com/microsoft/onnxruntime.git
cd onnxruntime
./build.sh --config Release --arm --build_shared_lib

Python部署代码：

import onnxruntime as ort
import numpy as np
import cv2
import time

# 加载模型
session = ort.InferenceSession("optimized_model.onnx")

# 获取摄像头
cap = cv2.VideoCapture(0)

while True:
    # 捕获帧
    ret, frame = cap.read()
    if not ret:
        break
        
    # 预处理
    input_img = cv2.resize(frame, (224, 224))
    input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)
    input_img = input_img.astype(np.float32) / 255.0
    input_img = input_img.transpose(2, 0, 1)  # HWC -> CHW
    input_img = np.expand_dims(input_img, axis=0)  # 添加批次维度
    
    # 获取输入名称
    input_name = session.get_inputs()[0].name
    
    # 推理
    start_time = time.time()
    outputs = session.run(None, {input_name: input_img})
    inference_time = time.time() - start_time
    
    # 处理输出
    output = outputs[0]
    class_id = np.argmax(output)
    
    # 显示结果
    cv2.putText(frame, f"Class: {class_id}, Time: {inference_time:.2f}s", 
                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
    cv2.imshow('ONNX Inference', frame)
    
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

5.3 云服务部署

使用Docker和ONNX Runtime部署REST API服务：

FROM python:3.8-slim

WORKDIR /app

# 安装依赖
RUN pip install onnxruntime numpy pillow flask gunicorn

# 复制应用代码和模型
COPY app.py /app/
COPY model.onnx /app/
COPY requirements.txt /app/

RUN pip install -r requirements.txt

# 暴露端口
EXPOSE 5000

# 启动服务
CMD ["gunicorn", "--bind", "0.0.0.0:5000", "app:app"]

Flask应用代码 (app.py)：

from flask import Flask, request, jsonify
import onnxruntime as ort
import numpy as np
from PIL import Image
import io
import time

app = Flask(__name__)

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

@app.route('/predict', methods=['POST'])
def predict():
    if 'image' not in request.files:
        return jsonify({'error': 'No image provided'}), 400
        
    # 读取图像
    file = request.files['image']
    img = Image.open(io.BytesIO(file.read()))
    
    # 预处理
    img = img.resize((224, 224))
    img = img.convert('RGB')
    input_data = np.array(img).transpose(2, 0, 1)  # HWC -> CHW
    input_data = input_data.astype(np.float32) / 255.0
    input_data = np.expand_dims(input_data, axis=0)  # 添加批次维度
    
    # 推理
    start_time = time.time()
    outputs = session.run([output_name], {input_name: input_data})
    inference_time = time.time() - start_time
    
    # 处理结果
    output = outputs[0]
    probabilities = np.exp(output) / np.sum(np.exp(output), axis=1, keepdims=True)
    class_id = int(np.argmax(probabilities))
    confidence = float(probabilities[0][class_id])
    
    return jsonify({
        'class_id': class_id,
        'confidence': confidence,
        'inference_time': inference_time
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

6. 性能调优与监控

6.1 性能分析

import onnxruntime as ort
import numpy as np
import time

# 启用性能分析
session_options = ort.SessionOptions()
session_options.enable_profiling = True
session = ort.InferenceSession("model.onnx", session_options)

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

# 预热
for _ in range(5):
    session.run(None, {input_name: input_data})

# 测量性能
iterations = 100
start_time = time.time()
for _ in range(iterations):
    session.run(None, {input_name: input_data})
end_time = time.time()

# 计算平均推理时间
avg_time = (end_time - start_time) / iterations
print(f"Average inference time: {avg_time * 1000:.2f} ms")

# 获取性能分析文件
profile_file = session.end_profiling()
print(f"Profiling file saved to: {profile_file}")

6.2 内存优化

import onnxruntime as ort
import numpy as np

# 配置会话选项
session_options = ort.SessionOptions()

# 内存优化设置
session_options.enable_mem_pattern = True
session_options.enable_mem_reuse = True

# 设置内存限制（以MB为单位）
session_options.set_memory_limit(512)  # 限制为512MB

# 创建会话
session = ort.InferenceSession("model.onnx", session_options)

# 执行推理
input_name = session.get_inputs()[0].name
input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {input_name: input_data})

6.3 多线程推理

import onnxruntime as ort
import numpy as np
import threading
import queue
import time

class InferenceWorker:
    def __init__(self, model_path, num_threads=4):
        self.num_threads = num_threads
        self.model_path = model_path
        self.request_queue = queue.Queue()
        self.result_queue = queue.Queue()
        self.workers = []
        self.running = True
        
        # 启动工作线程
        for _ in range(num_threads):
            thread = threading.Thread(target=self._worker_loop)
            thread.daemon = True
            thread.start()
            self.workers.append(thread)
    
    def _worker_loop(self):
        # 每个线程创建自己的会话
        session_options = ort.SessionOptions()
        session_options.inter_op_num_threads = 1
        session_options.intra_op_num_threads = 1
        session = ort.InferenceSession(self.model_path, session_options)
        input_name = session.get_inputs()[0].name
        
        while self.running:
            try:
                # 获取请求
                request_id, input_data = self.request_queue.get(timeout=1)
                
                # 执行推理
                outputs = session.run(None, {input_name: input_data})
                
                # 返回结果
                self.result_queue.put((request_id, outputs))
                
                # 标记任务完成
                self.request_queue.task_done()
            except queue.Empty:
                continue
    
    def infer_async(self, request_id, input_data):
        self.request_queue.put((request_id, input_data))
    
    def get_result(self, timeout=None):
        try:
            return self.result_queue.get(timeout=timeout)
        except queue.Empty:
            return None
    
    def shutdown(self):
        self.running = False
        for worker in self.workers:
            worker.join()

# 使用示例
if __name__ == "__main__":
    # 创建推理工作器
    worker = InferenceWorker("model.onnx", num_threads=4)
    
    # 提交多个推理请求
    for i in range(10):
        input_data = np.random.rand(1, 3, 224, 224).astype(np.float32)
        worker.infer_async(i, input_data)
    
    # 获取结果
    results = []
    for _ in range(10):
        result = worker.get_result()
        if result:
            request_id, outputs = result
            results.append((request_id, outputs))
    
    # 关闭工作器
    worker.shutdown()

7. 常见问题与解决方案

7.1 模型转换问题

7.2 部署问题

7.3 调试技巧

# 启用详细日志
import logging
logging.basicConfig(level=logging.INFO)
onnxruntime_logger = logging.getLogger("onnxruntime")
onnxruntime_logger.setLevel(logging.DEBUG)

# 检查中间结果
session_options = ort.SessionOptions()
session_options.log_severity_level = 0  # 详细日志
session = ort.InferenceSession("model.onnx", session_options)

# 比较原始框架和ONNX输出
def compare_outputs(original_output, onnx_output, rtol=1e-5, atol=1e-5):
    import numpy as np
    return np.allclose(original_output, onnx_output, rtol=rtol, atol=atol)

8. 未来发展趋势

• 更广泛的硬件支持：ONNX Runtime将支持更多专用硬件加速器
• 更高效的量化技术：混合精度量化和量化感知训练将提高性能
• 边缘AI部署：更多针对资源受限设备的优化
• 模型即服务：简化云端和边缘部署的工具和平台
• 联邦学习集成：支持分布式模型训练和部署

9. 总结

ONNX为深度学习模型部署提供了一个强大而灵活的解决方案，使模型能够在不同框架和平台之间无缝转换。通过ONNX Runtime，开发者可以在各种硬件上高效部署模型，从云服务器到边缘设备，从桌面应用到移动应用。

随着AI应用的普及，ONNX生态系统将继续发展，提供更多工具和优化技术，使模型部署更加简单和高效。掌握ONNX模型部署技术，将帮助开发者构建更强大、更灵活的AI应用。

10. 参考资源

• ONNX官方网站
• ONNX Runtime GitHub
• ONNX模型库
• PyTorch ONNX导出指南
• TensorFlow-ONNX转换工具