ebm2onnx 1.0.2

EBM model serialization to ONNX

Ebm2onnx is an EBM model serialization to ONNX. It allows to run an EBM model on any ONNX compliant runtime.

Features

• Binary classification

• Regression

• Continuous variables

• Categorical variables

• Interactions

• Multi-class classification (support is still experimental in EBM)

The export of the models is tested against ONNX Runtime.

Get Started

Train an EBM model:

# prepare dataset
df = pd.read_csv('titanic_train.csv')
df = df.dropna()

feature_columns = ['Age', 'Fare', 'Pclass', 'Embarked']
label_column = "Survived"
y = df[[label_column]]
le = LabelEncoder()
y_enc = le.fit_transform(y)
x = df[feature_columns]
x_train, x_test, y_train, y_test = train_test_split(x, y_enc)

# train an EBM model
model = ExplainableBoostingClassifier(
    feature_types=['continuous', 'continuous', 'continuous','categorical'],
)
model.fit(x_train, y_train)

Then you can convert it to ONNX in a single function call:

import ebm2onnx

onnx_model = ebm2onnx.to_onnx(
    model,
    dtype={
        'Age': 'double',
        'Fare': 'double',
        'Pclass': 'int',
    }
)
onnx.save_model(onnx_model, 'ebm_model.onnx')