-
什么是ONNX:
ONNX是一个开放的生态系统,允许AI开发者选择合适的工具,不会被特定的框架或库所限制。 -
ONNX的用途:
- 跨框架模型转换:允许在不同的深度学习框架之间转换模型。
- 模型优化:提供了一套标准的模型优化工具。
- 部署灵活性:便于将模型部署到各种硬件平台。
- 为什么使用ONNX:
- 互操作性:解决了不同框架之间的兼容性问题。
- 性能优化:ONNX Runtime可以自动优化模型性能。
- 跨平台部署:简化了模型在不同环境中的部署过程。
- 如何使用ONNX:
以下是一个使用PyTorch将模型转换为ONNX格式的示例:
import torch
import torchvision
# 加载预训练的ResNet模型
model = torchvision.models.resnet18(pretrained=True)
model.eval()
# 创建一个示例输入
dummy_input = torch.randn(1, 3, 224, 224)
# 导出模型为ONNX格式
torch.onnx.export(model, # 要导出的模型
dummy_input, # 模型的输入
"resnet18.onnx", # 输出的ONNX文件名
export_params=True, # 存储训练好的参数权重
opset_version=10, # ONNX版本
do_constant_folding=True, # 是否执行常量折叠优化
input_names = ['input'], # 输入节点的名称
output_names = ['output'], # 输出节点的名称
dynamic_axes={'input' : {0 : 'batch_size'}, # 批处理维度动态
'output' : {0 : 'batch_size'}})
print("Model exported to resnet18.onnx")
- 使用ONNX Runtime进行推理:
import onnxruntime
import numpy as np
# 加载ONNX模型
session = onnxruntime.InferenceSession("resnet18.onnx")
# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 运行推理
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
result = session.run([output_name], {input_name: input_data})
print(result[0])