Ebm2onnx converts EBM models to ONNX. It allows to run an EBM model on any ONNX compliant runtime.
- 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
The export of the models is tested against ONNX Runtime.
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')- You can live test the model conversion.
- You can live test local explanations.
The following organizations are supporting Ebm2onnx:
- SoftAtHome: Main supporter of Ebm2onnx development.
- InterpretML: Ebm2onnx is hosted under the umbrella of the InterpretML organization.
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent