当前位置： 首页 > news >正文

[模型部署] 1. 模型导出

news 2025/5/17 7:40:13

👋 你好！这里有实用干货与深度分享✨✨ 若有帮助，欢迎：​
👍 点赞 | ⭐ 收藏 | 💬 评论 | ➕ 关注 ，解锁更多精彩！​
📁 收藏专栏即可第一时间获取最新推送🔔。​
📖后续我将持续带来更多优质内容，期待与你一同探索知识，携手前行，共同进步🚀。​

​人工智能

模型导出

本文介绍如何将深度学习模型导出为不同的部署格式，包括ONNX、TorchScript等，并对各种格式的优缺点和最佳实践进行总结，帮助你高效完成模型部署准备。

1. 导出格式对比

格式优点缺点适用场景
ONNX- 跨平台跨框架
- 生态丰富
- 标准统一
- 广泛支持		- 可能存在算子兼容问题
- 部分高级特性支持有限		- 跨平台部署
- 使用标准推理引擎
- 需要广泛兼容性
TorchScript- 与PyTorch无缝集成
- 支持动态图结构
- 调试方便
- 性能优化		- 仅限PyTorch生态
- 文件体积较大		- PyTorch生产环境
- 需要动态特性
- 性能要求高
TensorRT- 极致优化性能
- 支持GPU加速
- 低延迟推理		- 仅支持NVIDIA GPU
- 配置复杂		- 高性能推理场景
- 实时应用
- 边缘计算
TensorFlow SavedModel- TensorFlow生态完整支持
- 部署便捷		- 跨框架兼容性差- TensorFlow生产环境

2. ONNX格式导出

2.1 基本导出

ONNX格式适用于跨平台部署，支持多种推理引擎（如ONNXRuntime、TensorRT、OpenVINO等）。

import torch
import torch.onnxdef export_to_onnx(model, input_shape, save_path):# 设置模型为评估模式model.eval()# 创建示例输入dummy_input = torch.randn(input_shape)# 导出模型torch.onnx.export(model,               # 要导出的模型dummy_input,        # 模型输入save_path,          # 保存路径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'}})print(f"Model exported to {save_path}")# 使用示例
model = YourModel()
model.load_state_dict(torch.load('model.pth'))
export_to_onnx(model, (1, 3, 224, 224), 'model.onnx')

2.2 验证导出模型

导出后必须进行全面验证，包括结构检查和数值对比：

  1. 结构验证
import onnx
import onnxruntime
import numpy as npdef verify_onnx_structure(onnx_path):# 加载并检查模型结构onnx_model = onnx.load(onnx_path)onnx.checker.check_model(onnx_model)# 打印模型信息print("模型输入:")for input in onnx_model.graph.input:print(f"- {input.name}: {input.type.tensor_type.shape}")print("\n模型输出:")for output in onnx_model.graph.output:print(f"- {output.name}: {output.type.tensor_type.shape}")
  1. 数值精度对比
def compare_outputs(model, onnx_path, input_data):# PyTorch结果model.eval()with torch.no_grad():torch_output = model(torch.from_numpy(input_data))# ONNX结果ort_output = verify_onnx_model(onnx_path, input_data)# 比较差异diff = np.abs(torch_output.numpy() - ort_output).max()print(f"最大误差: {diff}")return diff < 1e-5
  1. 验证 ONNX 模型
import onnx
import onnxruntime
import numpy as npdef verify_onnx_model(onnx_path, input_data):# 加载ONNX模型onnx_model = onnx.load(onnx_path)onnx.checker.check_model(onnx_model)# 创建推理会话ort_session = onnxruntime.InferenceSession(onnx_path)# 准备输入数据ort_inputs = {ort_session.get_inputs()[0].name: input_data}# 运行推理ort_outputs = ort_session.run(None, ort_inputs)return ort_outputs[0]input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
output = verify_onnx_model('model.onnx', input_data)

2.3 ONNX模型优化

使用ONNX Runtime提供的优化工具进一步提升性能：

import onnxruntime as ort
from onnxruntime.transformers import optimizerdef optimize_onnx_model(onnx_path, optimized_path):# 创建优化器配置opt_options = optimizer.OptimizationConfig(optimization_level=99,  # 最高优化级别enable_gelu_approximation=True,enable_layer_norm_optimization=True,enable_attention_fusion=True)# 优化模型optimized_model = optimizer.optimize_model(onnx_path, 'cpu',  # 或 'gpu'opt_options)# 保存优化后的模型optimized_model.save_model_to_file(optimized_path)print(f"优化后的模型已保存至 {optimized_path}")
  • optimizer.optimize_model() 第二个参数是优化目标设备，支持 ‘cpu’ 或 ‘gpu’。
    • 优化目标设备：指定模型优化时的目标硬件平台。例如：
      • ‘cpu’：针对 CPU 进行优化（如调整算子、量化参数等）。
      • ‘gpu’：针对 GPU 进行优化（如使用 CUDA 内核、张量核心等）。
        *运行时设备：优化后的模型可以在其他设备上运行，但性能可能受影响。例如：
      • 针对 CPU 优化的模型可以在 GPU 上运行，但可能无法充分利用 GPU 特性。
      • 针对 GPU 优化的模型在 CPU 上运行可能会报错或性能下降。
        建议保持优化目标与运行设备一致以获得最佳性能。

3. TorchScript格式导出

3.1 trace导出

适用于前向计算图结构固定的模型。

import torchdef export_torchscript_trace(model, input_shape, save_path):model.eval()example_input = torch.randn(input_shape)# 使用跟踪法导出traced_model = torch.jit.trace(model, example_input)traced_model.save(save_path)print(f"Traced model exported to {save_path}")return traced_model

3.2 script导出

适用于包含条件分支、循环等动态结构的模型。

import torch
import torch.nn as nn@torch.jit.script
class ScriptableModel(nn.Module):def __init__(self):super().__init__()self.conv = nn.Conv2d(3, 64, 3)self.relu = nn.ReLU()def forward(self, x):x = self.conv(x)x = self.relu(x)return xdef export_torchscript_script(model, save_path):scripted_model = torch.jit.script(model)scripted_model.save(save_path)print(f"Scripted model exported to {save_path}")return scripted_model

3.3 TorchScript模型验证

验证TorchScript模型的正确性：

def verify_torchscript_model(original_model, ts_model_path, input_data):# 原始模型输出original_model.eval()with torch.no_grad():original_output = original_model(input_data)# 加载TorchScript模型ts_model = torch.jit.load(ts_model_path)ts_model.eval()# TorchScript模型输出with torch.no_grad():ts_output = ts_model(input_data)# 比较差异diff = torch.abs(original_output - ts_output).max().item()print(f"最大误差: {diff}")return diff < 1e-5

4. 自定义算子处理

4.1 ONNX自定义算子

如需导出自定义算子，可通过ONNX扩展机制实现。

from onnx import helperdef create_custom_op():# 定义自定义算子custom_op = helper.make_node('CustomOp',           # 算子名称inputs=['input'],     # 输入outputs=['output'],   # 输出domain='custom.domain')return custom_opdef register_custom_op():# 注册自定义算子from onnxruntime.capi import _pybind_state as CC.register_custom_op('CustomOp', 'custom.domain')

4.2 TorchScript自定义算子

可通过C++扩展自定义TorchScript算子。

from torch.utils.cpp_extension import load# 编译自定义C++算子
custom_op = load(name="custom_op",sources=["custom_op.cpp"],verbose=True
)# 在模型中使用自定义算子
class ModelWithCustomOp(nn.Module):def forward(self, x):return custom_op.forward(x)

4.3 自定义算子示例

下面是一个完整的自定义算子实现示例：

// custom_op.cpp
#include <torch/extension.h>torch::Tensor custom_forward(torch::Tensor input) {return input.sigmoid().mul(2.0);
}PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {m.def("forward", &custom_forward, "Custom forward function");
}

# 在Python中使用
import torch
from torch.utils.cpp_extension import load# 编译自定义算子
custom_op = load(name="custom_op",sources=["custom_op.cpp"],verbose=True
)# 测试自定义算子
input_tensor = torch.randn(2, 3)
output = custom_op.forward(input_tensor)
print(output)

5. 模型部署示例

5.1 ONNXRuntime部署

import onnxruntime as ort
import numpy as np
from PIL import Image
import torchvision.transforms as transformsdef preprocess_image(image_path, input_shape):# 图像预处理transform = transforms.Compose([transforms.Resize((input_shape[2], input_shape[3])),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])image = Image.open(image_path).convert('RGB')image_tensor = transform(image).unsqueeze(0).numpy()return image_tensordef onnx_inference(onnx_path, image_path, input_shape=(1, 3, 224, 224)):# 加载ONNX模型session = ort.InferenceSession(onnx_path)# 预处理图像input_data = preprocess_image(image_path, input_shape)# 获取输入输出名称input_name = session.get_inputs()[0].nameoutput_name = session.get_outputs()[0].name# 执行推理result = session.run([output_name], {input_name: input_data})return result[0]

5.2 TorchScript部署

import torch
from PIL import Image
import torchvision.transforms as transformsdef torchscript_inference(model_path, image_path):# 加载TorchScript模型model = torch.jit.load(model_path)model.eval()# 图像预处理transform = transforms.Compose([transforms.Resize((224, 224)),transforms.ToTensor(),transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])# 加载并处理图像image = Image.open(image_path).convert('RGB')input_tensor = transform(image).unsqueeze(0)# 执行推理with torch.no_grad():output = model(input_tensor)return output

6. 常见问题与解决方案

6.1 ONNX导出失败

问题: 导出ONNX时出现算子不支持错误

解决方案:

# 尝试使用更高版本的opset
torch.onnx.export(model, dummy_input, "model.onnx", opset_version=13)# 或替换不支持的操作
class ModelWrapper(nn.Module):def __init__(self, model):super().__init__()self.model = modeldef forward(self, x):# 替换不支持的操作为等效操作return self.model(x)

6.2 TorchScript跟踪失败

问题: 动态控制流导致trace失败

解决方案:

# 使用script而非trace
scripted_model = torch.jit.script(model)# 或修改模型结构避免动态控制流
class TraceFriendlyModel(nn.Module):def __init__(self, original_model):super().__init__()self.model = original_modeldef forward(self, x):# 移除动态控制流return self.model.forward_fixed(x)

6.3 推理性能问题

问题: 导出模型推理速度慢

解决方案:

# 1. 使用量化
from torch.quantization import quantize_dynamic
quantized_model = quantize_dynamic(model, {nn.Linear, nn.Conv2d}, dtype=torch.qint8)# 2. 使用TensorRT优化ONNX
import tensorrt as trt
# TensorRT优化代码...# 3. 使用ONNX Runtime优化
import onnxruntime as ort
session = ort.InferenceSession("model.onnx", providers=['CUDAExecutionProvider'])

7. 最佳实践

  1. 选择合适的导出格式

    • ONNX：适合跨平台、跨框架部署，兼容性强
    • TorchScript：适合PyTorch生态内部署，支持灵活性高
    • 根据目标平台和性能需求选择

  2. 优化导出模型

    • 使用合适的opset版本（建议11及以上）
    • 启用常量折叠等优化选项
    • 导出后务必验证模型正确性
    • 考虑使用量化和剪枝优化模型大小

  3. 处理动态输入

    • 设置动态维度（如batch_size）
    • 测试不同输入大小，确保模型鲁棒性
    • 记录支持的输入范围和约束

  4. 文档和版本控制

    • 记录导出配置和依赖版本
    • 保存模型元数据（如输入输出规格）
    • 对模型文件进行版本化管理
    • 维护模型卡片(Model Card)记录关键信息

  5. 调试技巧

    • 使用ONNX Graph Viewer等可视化工具分析模型结构
    • 使用Netron查看计算图和参数分布
    • 比较原始与导出模型输出，检查数值精度差异
    • 遇到兼容性问题时查阅官方文档和社区经验

8. 参考资源

  • ONNX官方文档
  • PyTorch TorchScript教程
  • ONNX Runtime文档
  • TensorRT开发者指南
  • Netron模型可视化工具

​
​

📌 感谢阅读！若文章对你有用，别吝啬互动～​
👍 点个赞 | ⭐ 收藏备用 | 💬 留下你的想法 ，关注我，更多干货持续更新！

http://www.xdnf.cn/news/485677.html

相关文章：

  • 使用DDR4控制器实现多通道数据读写(十二)
  • 文章记单词 | 第88篇（六级）
  • ReentrantLock类详解
  • Spark，SparkSQL操作Mysql， 创建数据库和表
  • 【美团】Java后端一面复盘｜网络+线程+MySQL+Redis+设计模式+手撕算法
  • Neo4j 图书馆借阅系统知识图谱设计
  • Servlet 深度解析：生命周期、请求响应与状态管理
  • 表的设计、聚合函数
  • Linux信号的保存
  • 深入解析Spring Boot与微服务架构：从入门到实践
  • postgres的docker版本安装
  • 企业内容中台敏捷构建三步法
  • 火语言RPA--EmpireV7下载发布
  • Elasticsearch/OpenSearch 中doc_values的作用
  • 工商总局可视化模版 – 基于ECharts的大数据可视化HTML源码
  • mathematics-2024《Graph Convolutional Network for Image Restoration: A Survey》
  • 力扣HOT100之二叉树：102. 二叉树的层序遍历
  • Python循环控制
  • HarmonyOS开发组件基础
  • C#中UI线程的切换与后台线程的使用
  • SkyWalking的工作原理和搭建过程
  • 【Ansible基础】Ansible执行流程详解：从Playbook到实际任务
  • fpga系列 HDL : Microchip FPGA开发软件 Libero Soc 项目仿真示例
  • Femap许可优化策略
  • 如何选择靠谱的外卖霸王餐系统进行对接？
  • 编译opencv4.11gstreamer 参考
  • 4.3/Q1，Charls最新文章解读
  • 车道线检测----CLRERNet
  • 如何备考GRE？
  • python中列表的操作