def convert(model, feature_names, target):
"""Convert a decision tree model to protobuf format.
Parameters
----------
decision_tree : DecisionTreeRegressor
A trained scikit-learn tree model.
feature_names: [str]
Name of the input columns.
target: str
Name of the output column.
Returns
-------
model_spec: An object of type Model_pb.
Protobuf representation of the model
"""
if not(_HAS_SKLEARN):
raise RuntimeError('scikit-learn not found. scikit-learn conversion API is disabled.')
_sklearn_util.check_expected_type(model, _tree.DecisionTreeRegressor)
_sklearn_util.check_fitted(model, lambda m: hasattr(m, 'tree_') and model.tree_ is not None)
return _MLModel(_convert_tree_ensemble(model, feature_names, target))
评论列表
文章目录