Skip to content
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
23 lines (18 sloc) 802 Bytes
from . import trees
from . import net
from . import bayes
def convert(estimator, kind=None, method='pymodule', dtype='float'):
"""Main entrypoint for converting a model"""
if kind is None:
kind = type(estimator).__name__
# Use name instead of instance to avoid hard dependency on the libraries
if kind in ['RandomForestClassifier', 'ExtraTreesClassifier', 'DecisionTreeClassifier']:
return trees.Wrapper(estimator, method, dtype=dtype)
elif kind == 'MLPClassifier':
return net.convert_sklearn_mlp(estimator, method)
elif kind == 'Sequential':
return net.convert_keras(estimator, method)
elif kind == 'GaussianNB':
return bayes.Wrapper(estimator, method)
raise ValueError("Unknown model type: '{}'".format(kind))
You can’t perform that action at this time.