将 onnx 节点附加到转换后的模型

此示例展示了如何将一些 onnx 节点附加到转换后的模型以生成所需的输出。在此案例中,它移除了输出概率的第二列。

准确地说,大部分代码是使用 LLM 生成并修改以适应最新变化。

from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
import onnx

iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y)
clr = LogisticRegression(max_iter=500)
clr.fit(X_train, y_train)
LogisticRegression(max_iter=500)
在 Jupyter 环境中,请重新运行此单元格以显示 HTML 表示或信任此 notebook。
在 GitHub 上,HTML 表示无法渲染,请尝试使用 nbviewer.org 加载此页面。


model_to_convert 指要转换的 scikit-learn 分类器。

model_to_convert = clr  # model to convert
X_test = X_test[:1]  # data used to test or train, one row is enough

设置修改后的 ONNX 模型的输出文件名

output_filename = "output_file.onnx"  # Replace with your desired output filename

步骤 1:将模型转换为 ONNX 格式,禁用标签输出。定义 ONNX 模型的输入类型。输入类型是一个 float 张量,形状为 [None, X_test.shape[1]],其中 None 表示输入样本数量可以灵活,X_test.shape[1] 是每个输入样本的特征数量。“张量”本质上是一个多维数组,常用于机器学习中表示数据。“float 张量”特指包含浮点数的张量,浮点数是带有小数的数字。

initial_type = [("float_input", FloatTensorType([None, X_test.shape[1]]))]

将模型转换为 ONNX 格式。

  • target_opset=18 指定要使用的 ONNX 操作符版本。

  • options={…} 设置转换参数: - “zipmap”: False 确保输出是原始的概率数组而非字典。 - “output_class_labels”: False 确保输出

    仅包含概率,不包含类别标签。

ONNX(Open Neural Network Exchange,开放神经网络交换)是一种表示机器学习模型的开放格式。它允许不同机器学习框架之间的互操作性,从而能够在各种平台中使用模型。

onx = convert_sklearn(
    model_to_convert,
    initial_types=initial_type,
    target_opset={"": 18, "ai.onnx.ml": 3},
    options={
        id(model_to_convert): {"zipmap": False, "output_class_labels": False}
    },  # Ensures the output is only probabilities, not labels
)

步骤 2:加载 ONNX 模型,以便在需要时进一步修改从序列化的字符串表示中加载 ONNX 模型。ONNX 文件本质上是机器学习模型的序列化表示,可以在不同系统之间共享和使用。

onnx_model = onnx.load_model_from_string(onx.SerializeToString())

假设此模型中的第一个输出应该是概率张量提取表示概率的输出张量的名称。如果存在多个输出,选择第二个,否则选择第一个。

prob_output_name = (
    onnx_model.graph.output[1].name
    if len(onnx_model.graph.output) > 1
    else onnx_model.graph.output[0].name
)

添加 Gather 节点,仅提取正类别的概率 (索引 1)创建一个张量来指定要聚集的索引(索引 1),它代表正类别。

indices = onnx.helper.make_tensor(
    "indices", onnx.TensorProto.INT64, (1,), [1]
)  # Index 1 to gather positive class

在 ONNX 图中创建一个“Gather”节点,以提取正类别的概率。 - 输入: [prob_output_name, “indices”] 指定输入

到此节点(概率张量和索引张量)。

  • 输出: [“positive_class_prob”] 指定此节点的输出名称。

  • axis=1 表示沿着概率张量的列(特征)进行聚集。

“Gather”节点用于从张量中提取特定元素。在此处,它提取正类别的概率。

gather_node = onnx.helper.make_node(
    "Gather",
    inputs=[prob_output_name, "indices"],
    outputs=["positive_class_prob"],
    axis=1,  # Gather along columns (axis 1)
)

将 Gather 节点添加到 ONNX 图

onnx_model.graph.node.append(gather_node)
input: "probabilities"
input: "indices"
output: "positive_class_prob"
op_type: "Gather"
attribute {
  name: "axis"
  i: 1
  type: INT
}

为索引添加张量初始化器(Gather 节点需要)ONNX 中的初始化器用于定义计算中使用的常量张量。

onnx_model.graph.initializer.append(indices)
dims: 1
data_type: 7
int64_data: 1
name: "indices"

移除现有输出,仅添加正类别概率的新输出清除现有输出定义,以替换为新输出。

del onnx_model.graph.output[:]

定义正类别概率的新输出创建一个名为“positive_class_prob”的新输出张量规范。

positive_class_output = onnx.helper.make_tensor_value_info(
    "positive_class_prob", onnx.TensorProto.FLOAT, [None, 1]
)
onnx_model.graph.output.append(positive_class_output)
name: "positive_class_prob"
type {
  tensor_type {
    elem_type: 1
    shape {
      dim {
      }
      dim {
        dim_value: 1
      }
    }
  }
}

步骤 3:保存修改后的 ONNX 模型将修改后的 ONNX 模型保存到指定的输出文件名。生成的 ONNX 文件随后可以在支持 ONNX 的不同环境中使用,例如推理服务器或其他机器学习框架。

onnx.save(onnx_model, output_filename)

模型可以按如下方式打印。

print(onnx.printer.to_text(onnx_model))
<
   ir_version: 8,
   opset_import: ["ai.onnx.ml" : 1, "" : 18],
   producer_name: "skl2onnx",
   producer_version: "1.18.0",
   domain: "ai.onnx",
   model_version: 0,
   doc_string: ""
>
"2a0b430c8e794cac873666482db85939" (float[?,4] float_input) => (float[?,1] positive_class_prob)
   <int64[1] indices =  {1}>
{
   [LinearClassifier] label, probability_tensor = ai.onnx.ml.LinearClassifier <classlabels_ints: ints = [0, 1, 2], coefficients: floats = [-0.442649, 0.828556, -2.36352, -0.972382, 0.269401, -0.245, -0.185949, -0.761031, 0.173248, -0.583556, 2.54947, 1.73341], intercepts: floats = [9.78245, 3.05803, -12.8405], multi_class: int = 1, post_transform: string = "SOFTMAX"> (float_input)
   [Normalizer] probabilities = ai.onnx.ml.Normalizer <norm: string = "L1"> (probability_tensor)
   positive_class_prob = Gather <axis: int = 1> (probabilities, indices)
}

脚本总运行时间: (0 分 0.020 秒)

由 Sphinx-Gallery 生成的图集