def convert_model(model, name, input_types):
"""
Runs the appropriate conversion method.
:param model: model
:return: *onnx* model
"""
from sklearn.base import BaseEstimator
if model.__class__.__name__.startswith("LGBM"):
from onnxmltools.convert import convert_lightgbm
model, prefix = convert_lightgbm(model, name, input_types), "LightGbm"
elif model.__class__.__name__.startswith("XGB"):
from onnxmltools.convert import convert_xgboost
model, prefix = convert_xgboost(model, name, input_types), "XGB"
elif model.__class__.__name__ == 'Booster':
import lightgbm
if isinstance(model, lightgbm.Booster):
from onnxmltools.convert import convert_lightgbm
model, prefix = convert_lightgbm(model, name,
input_types), "LightGbm"
else:
raise RuntimeError("Unable to convert model of type '{0}'.".format(
type(model)))
elif isinstance(model, BaseEstimator):
from onnxmltools.convert import convert_sklearn
model, prefix = convert_sklearn(model, name, input_types), "Sklearn"
else:
from onnxmltools.convert import convert_coreml
model, prefix = convert_coreml(model, name, input_types), "Cml"
if model is None:
raise RuntimeError("Unable to convert model of type '{0}'.".format(
type(model)))
return model, prefix
def convert_model(model, name, input_types):
"""
Runs the appropriate conversion method.
:param model: model, *scikit-learn*, *keras*, or *coremltools* object
:return: *onnx* model
"""
from sklearn.base import BaseEstimator
if model.__class__.__name__.startswith("LGBM"):
from onnxmltools.convert import convert_lightgbm
model, prefix = convert_lightgbm(model, name, input_types), "LightGbm"
elif isinstance(model, BaseEstimator):
from onnxmltools.convert import convert_sklearn
model, prefix = convert_sklearn(model, name, input_types), "Sklearn"
else:
from keras.models import Model
if isinstance(model, Model):
from onnxmltools.convert import convert_keras
model, prefix = convert_keras(model, name, input_types), "Keras"
else:
from onnxmltools.convert import convert_coreml
model, prefix = convert_coreml(model, name, input_types), "Cml"
if model is None:
raise RuntimeError("Unable to convert model of type '{0}'.".format(type(model)))
return model, prefix
def coreml_converter(args):
# When imported, CoreML tools checks for the current version of Keras and TF and prints warnings if they are
# outside its expected range. We don't want it to import these packages (since they are big and take seconds to
# load) and we don't want to clutter the console with unrelated Keras warnings when converting from CoreML.
import sys
sys.modules['keras'] = None
import coremltools
from onnxmltools.convert import convert_coreml
source_model = coremltools.utils.load_spec(args.source)
onnx_model = convert_coreml(source_model,
name=args.name,
target_opset=get_opset(args.ONNXVersion))
return onnx_model
def get_diff(self, input_file, ref_file):
this = os.path.dirname(__file__)
coreml_file = os.path.join(this, "models", input_file)
cml = coremltools.utils.load_spec(coreml_file)
onnx_model = convert_coreml(cml)
output_dir = os.path.join(this, "outmodels")
output_file = os.path.join(this, "outmodels", ref_file)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
save_text(onnx_model, output_file)
reference_model = os.path.join(this, "models", ref_file)
with open(reference_model, 'r') as ref_file:
with open(output_file, 'r') as output_file:
diff = set(ref_file).difference(output_file)
return diff