ONNX导出全解析：跨平台部署Python示例代码-深圳市維司達科技有限公司

ONNX导出全解析：跨平台部署Python示例代码

1. ONNX模型导出的核心价值

在AI模型从训练走向实际应用的过程中，跨平台部署能力是决定其能否落地的关键。ONNX（Open Neural Network Exchange）作为一种开放的神经网络交换格式，正在成为连接不同框架与硬件平台的桥梁。

以cv_resnet18_ocr-detection这一OCR文字检测模型为例，它基于PyTorch构建，具备高精度的文字区域定位能力。但若想将该模型部署到Windows、Linux、嵌入式设备或移动端，直接使用.pt权重文件会面临环境依赖复杂、推理引擎不兼容等问题。而通过ONNX导出，我们可以实现：

统一模型接口：一次导出，多端运行
硬件加速支持：适配TensorRT、OpenVINO、Core ML等优化引擎
轻量化部署：去除训练相关操作，减小模型体积
提升推理效率：结合后端优化策略，显著降低延迟

本文将围绕该OCR模型的ONNX导出流程，深入讲解关键步骤、参数设置及跨平台调用方法，并提供可直接运行的Python示例代码。

2. 模型结构与ONNX导出准备

2.1 OCR检测模型架构分析

cv_resnet18_ocr-detection采用经典的两阶段OCR架构：

主干网络（Backbone）：ResNet-18用于提取图像特征
特征融合模块（FPN）：增强多尺度特征表达能力
检测头（DBHead）：生成概率图和阈值图，通过可微分二值化（DB）输出文本框

这种设计兼顾了速度与精度，特别适合文档、截图等场景下的文字检测任务。

2.2 导出前的关键检查项

在执行ONNX导出之前，必须确保以下几点：

模型已切换至评估模式（model.eval()）
所有权重已正确加载且可在CPU上运行
输入输出节点名称明确，便于后续调用
动态轴设置合理，支持变尺寸输入

此外，还需确认PyTorch版本支持目标ONNX Opset版本（建议Opset ≥ 11），避免出现算子不兼容问题。

3. ONNX导出全流程详解

3.1 构建模型实例并加载权重

首先需要重建模型结构并载入预训练权重。以下是核心代码实现：

import torch from models.model import Model # 假设模型定义在此模块中 # 定义模型配置 model_config = { 'backbone': {'type': 'resnet18', 'pretrained': False, "in_channels": 3}, 'neck': {'type': 'FPN', 'inner_channels': 256}, 'head': {'type': 'DBHead', 'out_channels': 2, 'k': 50}, } # 实例化模型 model = Model(model_config=model_config) # 加载训练好的权重 weights_path = "/root/cv_resnet18_ocr-detection/workdirs/best_model.pth" state_dict = torch.load(weights_path, map_location=torch.device('cpu')) model.load_state_dict(state_dict) # 切换为评估模式 model.eval()

注意：务必使用map_location='cpu'确保模型可在无GPU环境下加载，这对后续跨平台部署至关重要。

3.2 构造虚拟输入张量

ONNX导出需要一个符合模型输入要求的示例张量。对于图像模型，通常为[B, C, H, W]格式：

# 构造输入张量（假设输入尺寸为800x800） dummy_input = torch.randn(1, 3, 800, 800) # B=1, C=3(RGB), H=800, W=800

你也可以根据实际需求调整输入尺寸，如640×640或1024×1024，具体选择见下文建议。

3.3 执行ONNX导出操作

调用torch.onnx.export完成模型转换：

import onnx onnx_path = "resnet18_ocr_detection_800x800.onnx" torch.onnx.export( model, dummy_input, onnx_path, export_params=True, # 保存训练好的参数 opset_version=11, # 使用ONNX Opset 11 do_constant_folding=True, # 合并常量以优化计算图 input_names=['input'], # 输入节点命名 output_names=['output'], # 输出节点命名 dynamic_axes={ 'input': {0: 'batch_size', 2: 'height', 3: 'width'}, # 支持动态批大小和图像尺寸 'output': {0: 'batch_size', 2: 'out_height', 3: 'out_width'} } ) # 验证导出结果 onnx_model = onnx.load(onnx_path) onnx.checker.check_model(onnx_model) print(f" ONNX模型已成功导出至: {onnx_path}")

参数说明：

参数	作用
`export_params`	是否包含学习到的权重
`opset_version`	算子集版本，影响兼容性
`do_constant_folding`	优化：合并常量运算
`dynamic_axes`	允许输入/输出具有动态维度

4. 输入尺寸选择与性能权衡

在WebUI界面中提供了三种常见输入尺寸选项，每种都有其适用场景：

输入尺寸	推理速度	内存占用	适用场景
640×640	快	低	移动端、实时检测
800×800	中等	中等	平衡型通用部署
1024×1024	慢	高	高精度文档识别

推荐实践：优先选择800×800作为默认尺寸，在保证精度的同时保持良好性能。

你可以通过修改dummy_input的形状来生成对应尺寸的ONNX模型，例如：

# 生成640x640版本 dummy_input = torch.randn(1, 3, 640, 640) onnx_path = "resnet18_ocr_detection_640x640.onnx" # 再次调用torch.onnx.export...

5. Python环境下的ONNX推理实战

5.1 安装依赖库

pip install onnxruntime opencv-python numpy

推荐使用onnxruntime-gpu以获得更高推理速度（需CUDA支持）。

5.2 图像预处理函数

import cv2 import numpy as np def preprocess_image(image_path, target_size=(800, 800)): """ 图像预处理：读取、缩放、归一化 """ image = cv2.imread(image_path) resized = cv2.resize(image, target_size) # 转换为 NCHW 格式并归一化 input_blob = resized.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32) / 255.0 return input_blob, image.shape[:2] # 返回原始尺寸用于坐标还原

5.3 ONNX推理执行

import onnxruntime as ort # 加载ONNX模型 session = ort.InferenceSession("resnet18_ocr_detection_800x800.onnx", providers=['CPUExecutionProvider']) # 可替换为'GPUExecutionProvider' # 准备输入数据 input_data, original_shape = preprocess_image("test.jpg") # 执行推理 outputs = session.run(None, {"input": input_data}) probability_map = outputs[0][0, 0] # 获取概率图 (H, W)

5.4 后处理：提取文本框坐标

import cv2 def postprocess(prob_map, original_shape, threshold=0.2): """ 将概率图转换为文本框坐标 """ h, w = original_shape prob_resized = cv2.resize(prob_map, (w, h)) _, binary = cv2.threshold(prob_resized, threshold, 255, cv2.THRESH_BINARY) binary = binary.astype(np.uint8) contours, _ = cv2.findContours(binary, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) boxes = [] for cnt in contours: if cv2.contourArea(cnt) < 50: # 过滤小区域 continue rect = cv2.minAreaRect(cnt) box = cv2.boxPoints(rect).astype(int) boxes.append(box.tolist()) return boxes # 调用后处理 detected_boxes = postprocess(probability_map, original_shape, threshold=0.2) print("检测到的文本框数量:", len(detected_boxes))

6. WebUI中的ONNX导出功能使用指南

6.1 访问ONNX导出页面

启动服务：bash start_app.sh
浏览器访问：http://<服务器IP>:7860
切换至“ONNX 导出”Tab页

6.2 设置输入尺寸并导出

在“输入高度”和“输入宽度”中填写目标尺寸（如800）
点击“导出 ONNX”按钮
等待提示“导出成功！”后点击“下载 ONNX 模型”

导出的模型将保存在容器内指定路径，可通过Web界面一键下载。

7. 常见问题与解决方案

7.1 导出失败：算子不支持

现象：抛出Unsupported ONNX opset version或特定算子错误
解决方法：

降低opset_version至10或9
查看PyTorch版本是否过新导致兼容性问题
使用torch.jit.trace先转为TorchScript再导出

7.2 推理结果为空

可能原因：

输入图像未正确归一化（应除以255.0）
模型输入尺寸与导出时不一致
阈值设置过高导致无法触发检测

调试建议：

打印probability_map.max()查看最大响应值
可视化概率图确认是否有激活区域

7.3 性能不佳

优化方向：

使用ONNX Runtime的GPU或TensorRT后端
开启sess_options.graph_optimization_level
对输入图像进行适当降采样

8. 总结

ONNX为AI模型的跨平台部署提供了标准化路径。通过对cv_resnet18_ocr-detection模型的完整导出与推理实践，我们验证了以下关键点：

正确的模型初始化与权重加载是导出前提
合理设置dynamic_axes可提升部署灵活性
输入尺寸需根据应用场景权衡精度与速度
ONNX Runtime提供了简洁高效的推理接口

掌握这些技能后，你不仅可以将此OCR模型部署到各类边缘设备，还能将其集成进Java、C++、JavaScript等非Python系统中，真正实现“一次训练，处处运行”。

获取更多AI镜像
想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。

ONNX导出全解析：跨平台部署Python示例代码