📢 转载信息
原文链接：https://machinelearningmastery.com/export-your-ml-model-in-onnx-format/
原文作者：Jason Brownlee
ONNX（开放神经网络交换）是一种用于表示机器学习模型的开放格式。它允许模型在不同的深度学习框架之间移植，例如，你可以训练一个模型，然后将其转换为ONNX格式，并在另一个支持ONNX的框架或运行时中进行推理。
本指南将介绍如何将模型导出为ONNX格式，重点介绍如何使用TensorFlow/Keras和PyTorch框架进行导出。
为什么使用ONNX？
ONNX的主要目的是在不同框架之间提供模型的可移植性。在将模型部署到生产环境时，这非常有用。
例如，你可能希望使用PyTorch训练模型，但最终在C++推理引擎中运行它。或者，你可能想在基于TensorFlow的生产环境中部署一个使用Keras训练的模型。ONNX充当了框架之间的桥梁。
将模型转换为ONNX格式的常见场景包括：
- 跨框架部署： 在一个框架中训练，在另一个框架中推理。
- 性能优化： 使用支持ONNX的运行时（如ONNX Runtime）进行高效推理。
- 硬件加速： 许多硬件加速器（如NVIDIA TensorRT）支持导入ONNX模型。
如何导出模型为ONNX格式
导出过程取决于你最初训练模型所使用的框架。最流行的两个框架是TensorFlow（及其高级API Keras）和PyTorch。
使用TensorFlow/Keras导出模型到ONNX
对于TensorFlow模型，我们通常使用tf2onnx库。这是一个开源工具，可以将TensorFlow SavedModel或Keras模型转换为ONNX格式。
首先，你需要安装tf2onnx库。
pip install tf2onnx其次，你需要一个已训练的TensorFlow/Keras模型。假设你有一个保存好的模型，其路径为model/。
你可以使用python -m tf2onnx.convert命令将模型转换为ONNX格式。你需要指定输入和输出节点的名称。
确定输入和输出节点名称
确定模型的输入和输出节点名称是至关重要的一步，因为tf2onnx需要它们才能正确构建ONNX计算图。如果你不知道这些名称，可以在模型加载后检查。
import tensorflow as tf
# 加载模型
model = tf.keras.models.load_model('model/')
# 检查输入和输出签名
print(model.signatures)
通常，对于一个标准的Keras模型，签名键是'serving_default'，你可以从这个签名中获取输入和输出名称。
执行转换
一旦确定了节点名称，就可以执行转换。
python -m tf2onnx.convert --saved-model model/ --output model.onnx --signature_def serving_default如果需要指定特定的输入和输出名称（例如，当模型有多个输入/输出时），可以使用--inputs和--outputs标志。
使用PyTorch导出模型到ONNX
PyTorch提供了内置的支持，通过torch.onnx.export()函数可以直接将模型导出为ONNX格式。这使得PyTorch模型转换为ONNX变得非常直接。
准备工作
你需要一个已训练的PyTorch模型实例和一个示例输入张量。示例输入张量用于跟踪模型在执行过程中接收到的输入形状和数据类型。
import torch
# 假设你有一个已训练的模型实例
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval() # 确保模型处于评估模式
# 创建一个示例输入（形状和数据类型必须匹配你的模型预期）
example_input = torch.randn(1, 3, 224, 224) # 示例：批量大小为1，3通道，224x224图像
执行导出
使用torch.onnx.export()函数。
output_path = "model.onnx"
input_names = ["input"]
output_names = ["output"]
# 执行导出
torch.onnx.export(model,
example_input,
output_path,
verbose=True,
input_names=input_names,
output_names=output_names,
opset_version=13) # 建议指定opset版本
关键参数说明：
model：要导出的PyTorch模型。example_input：用于跟踪计算图的示例输入。output_path：ONNX文件的保存路径。input_names/output_names：为输入和输出节点指定名称（可选，但推荐）。opset_version：指定ONNX操作集的版本。不同的版本支持不同的操作。
验证ONNX模型
将模型导出为ONNX格式后，最好验证它是否正确转换，并且能够在ONNX Runtime中正确运行。
你可以使用onnxruntime库来加载和运行ONNX模型。
import onnxruntime as rt
import numpy as np
# 加载ONNX模型
ort_session = rt.InferenceSession("model.onnx")
# 准备输入数据（必须与导出时示例输入的形状和类型一致）
input_name = ort_session.get_inputs()[0].name
data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 运行推理
ort_outputs = ort_session.run(None, {input_name: data})
print("ONNX Runtime Inference successful.")
通过比较原始框架（如PyTorch或TensorFlow）的输出和ONNX Runtime的输出，你可以确保转换是准确的。
总结
ONNX为机器学习模型的互操作性和部署提供了一个强大的解决方案。无论是使用tf2onnx转换TensorFlow/Keras模型，还是使用torch.onnx.export转换PyTorch模型，掌握这一过程都是将AI模型投入生产的关键步骤。确保在转换过程中正确指定输入/输出节点以及ONNX操作集版本，是获得成功部署的关键。
🚀 想要体验更好更全面的AI调用？
欢迎使用青云聚合API，约为官网价格的十分之一，支持300+全球最新模型，以及全球各种生图生视频模型，无需翻墙高速稳定，文档丰富，小白也可以简单操作。
评论区