ONNX初学者指南

TOOL Nov 27, 2024

跨平台推理或缓慢、未优化的模型让你苦苦挣扎？让我们快速轻松地深入了解 ONNX — 它的工作原理、亮点所在以及如何像专业人士一样排除模型故障！

1、ONNX ？

ONNX（开放神经网络交换）就像是 AI 模型的通用转换器 — 一种允许你跨平台保存和共享机器学习或深度学习模型的格式，就像 h5 或 pth 文件一样。它是开源的并且与框架无关，这意味着你可以在 PyTorch 或 TensorFlow 中训练模型并将其部署到任何地方，从 ONNX Runtime 到自定义边缘设备，而无需担心兼容性问题。

1.1 为什么要使用 ONNX ？

在构建机器学习或深度学习模型时，训练通常会遵循特定的框架 — 但生产需要灵活性。这就是 ONNX 介入的地方，它充当了一座通用桥梁。它使模型可移植性更容易，加快部署速度，并确保跨环境的可扩展性，无论是云还是边缘设备。

1.2 ONNX 的一些常见场景

使用不同框架的团队之间共享模型。
降低管理多个框架和库的复杂性。
在资源有限的边缘设备上运行轻量级模型。

1.3 ONNX 的主要特点

ONNX 带来了很多好处：

互操作性：在 PyTorch 和 TensorFlow 等框架之间无缝传输模型。
跨平台：轻松在 CPU、GPU 或专用硬件上运行模型。
推理优化：使用图形修剪和量化等技术提高性能。
框架支持：适用于流行的 ML 框架，如 PyTorch、TensorFlow 和 scikit-learn。
标准化：一种简化模型交换和部署的通用文件格式。
开源：完全免费，拥有支持社区和定期更新。
更小的模型尺寸：减小模型尺寸，以便在边缘设备上更快、更高效地部署。

2、将模型导出到 ONNX

设置环境：

pip install onnx onnxruntime

从 PyTorch 导出：

import torch

# Define or load your model
model = MyModel() # Your trained model
model.eval()  # Set the model to evaluation mode

# Define a dummy input with the same shape as real input,to let the model know about the input size as a input.
dummy_input = torch.randn(1, 3, 224, 224)

# Export the model to ONNX
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=11,  # Specify the ONNX opset version
    input_names=['input'],
    output_names=['output']
)

从 Tensorflow 导出，使用 tf2onnx 转换器：

python -m tf2onnx.convert --saved-model ./saved_model --output model.onnx

从 Scikit-learn 导出：

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# Train a model
model = LogisticRegression().fit(X_train, y_train)

# Define input type
initial_type = [('float_input', FloatTensorType([None, X_train.shape[1]]))]

# Convert and save as ONNX
onnx_model = convert_sklearn(model, initial_types=initial_type)
with open("model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

3、部署 ONNX 模型

将模型转换为 ONNX 后，下一步是使用 ONNX Runtime — 这是一种多功能工具，支持多种编程语言（如 Python、C++ 和 Java），并在各种硬件平台上顺利运行。它是你高效且可扩展的模型部署的门户！

import onnxruntime as ort
import numpy as np

# Load the ONNX model
session = ort.InferenceSession("model.onnx")

# Prepare input
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)

# Run inference
result = session.run([output_name], {input_name: input_data})
print("Inference result:", result)