EBM model serialization to ONNX
Project description
Ebm2onnx converts EBM models to ONNX. It allows to run an EBM model on any ONNX compliant runtime.
Features
Binary classification
Regression
Continuous, nominal, and ordinal variables
N-way interactions
Multi-class classification (support is still experimental in EBM)
Expose predictions probabilities
Expose local explanations
Export a model as part of a scikit-learn pipeline (experimental)
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', 'nominal'],
)
model.fit(x_train, y_train)
Then you can convert it to ONNX in a single function call:
import onnx
import ebm2onnx
onnx_model = ebm2onnx.to_onnx(
model,
ebm2onnx.get_dtype_from_pandas(x_train),
)
onnx.save_model(onnx_model, 'ebm_model.onnx')
If your dataset is not a pandas dataframe, you can provide the features’ types directly:
import ebm2onnx
onnx_model = ebm2onnx.to_onnx(
model,
dtype={
'Age': 'double',
'Fare': 'double',
'Pclass': 'int',
'Embarked': 'str',
}
)
onnx.save_model(onnx_model, 'ebm_model.onnx')
Try it live
You can live test the model conversion.
You can live test local explanations.
You can live test the export of a scikit-learn pipeline.
Supporting organizations
The following organizations are supporting Ebm2onnx:
SoftAtHome: Main supporter of Ebm2onnx development.
InterpretML: Ebm2onnx is hosted under the umbrella of the InterpretML organization.