lightGBM转onnx

一、数据读取

from sklearn.datasets import load_iris

from sklearn.model_selection import train_test_split

import lightgbm as lgb

from onnxconverter_common.data_types import FloatTensorType

import onnxmltools

iris = load_iris()

X, y = iris.data, iris.target

X_train, X_test, y_train, y_test = train_test_split(X, y)

train_data = lgb.Dataset(X_train, label=y_train)

二、模型训练

test_data = lgb.Dataset(X_test, label=y_test)

param = {'num_leaves': 3, 'objective': 'multiclass','num_class':3}

param['metric'] = ['multi_logloss','multi_error']

num_round = 10

bst = lgb.train(param, train_data, num_round, valid_sets=[test_data])

bst.save_model('model.txt')

三、转ONNX

with open('model.txt') as fd:

  for line in fd:

    if line.startswith("max_feature_idx"):

      max_feature_idx = int(line.split("=")[1])


 

lgb_regression_model = lgb.Booster(model_file='model.txt')

 

initial_type = [("float_input", FloatTensorType([None, max_feature_idx+1]))]

onnx_model = onnxmltools.convert_lightgbm(lgb_regression_model, initial_types = initial_type )

onnxmltools.utils.save_model(onnx_model, 'model.onnx')

四、ONNX模型预测

import onnxruntime as rt

import numpy

sess = rt.InferenceSession("model.onnx")

input_name = sess.get_inputs()[0].name

label_name = sess.get_outputs()[0].name

pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0]

print(pred_onx)