ONNX export of an Extra Tree
Download Python samples A Zip archive containing all samples can be found here: Samples of ONNX export |
Extra Tree Regressor with Scikit-learn
from sklearn.datasets import make_regression
from sklearn.tree import ExtraTreeRegressor
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
# # Generate data for regression
X, y = make_regression(n_samples=300, n_features=10, n_informative=10, n_targets=1)
# # Construct Extra Tree-Model
model = ExtraTreeRegressor(criterion='squared_error', splitter='random', max_depth=None, min_samples_split=2, min_samples_leaf=1)
model.fit(X,y)
# # Convert model to ONNX
onnxfile = 'extratree-regressor.onnx'
initial_type = [('float_input', FloatTensorType([None, X.shape[1]]))]
onnx_model = convert_sklearn(model, initial_types=initial_type, target_opset=12)
# # Export to ONNX file
with open(onnxfile, "wb") as f:
f.write( onnx_model.SerializeToString())
f.close()
exit()
Extra Tree Classifier with Scikit-learn
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import ExtraTreeClassifier
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
# # Load dataset
X, y = load_iris(return_X_y = True)
# # Construct Decision Tree-Model
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.1)
model = ExtraTreeClassifier(criterion='gini', splitter='random', max_depth=None, min_samples_split=2, min_samples_leaf=1)
model.fit(X_train,y_train)
# # Convert model to ONNX
onnxfile = 'extratree-classifier.onnx'
initial_type = [('float_input', FloatTensorType([None, X.shape[1]]))]
# # Zipmap should be always turned off as it's not implemented in TF3800
onnx_model = convert_sklearn(model, initial_types=initial_type, options={type(model): {'zipmap':False}}, target_opset=12)
# # Export to ONNX file
with open(onnxfile, "wb") as f:
f.write( onnx_model.SerializeToString())
f.close()
exit()