python模型转换为ONNX格式

原创

于 2021-01-20 17:18:57 发布 · 1.3k 阅读

1 ·

CC 4.0 BY-SA版权

本文介绍了如何将训练好的Python模型转换为ONNX格式，包括三个步骤：训练模型、模型转换和使用ONNX Runtime运行模型。ONNX是一种通用的模型描述格式，支持多种框架间的模型互操作。

ONNX(Open Neural Network Exchange) 可以优化机器学习模型的推理，
许多框架（包括 TensorFlow、PyTorch、SciKit-Learn、Keras、Chainer、MXNet、MATLAB 和 SparkML）中的模型都可以导出或转换为标准 ONNX 格式。

下面将以sklearn模块中的LogisticRegression模型为例，做一个简单介绍：

示例1

Step 1: 训练模型

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression

# 加载数据
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X