Source code for onnxmltools.convert.h2o.convert
# SPDX-License-Identifier: Apache-2.0
from uuid import uuid4
import json
import tempfile
import h2o
import onnx
from ..common.onnx_ex import get_maximum_opset_supported
from ..common._topology import convert_topology
from ..common.data_types import FloatTensorType
from ._parse import parse_h2o
# Invoke the registration of all our converters and shape calculators
from . import operator_converters, shape_calculators # noqa
[docs]def convert(model, name=None, initial_types=None, doc_string='', target_opset=None,
targeted_onnx=onnx.__version__, custom_conversion_functions=None,
custom_shape_calculators=None):
'''
This function produces an equivalent ONNX model of the given H2O MOJO model.
Supported model types:
- GBM, with limitations:
- poisson, gamma, tweedie distributions not supported
- multinomial distribution supported with 3 or more classes (use binomial otherwise)
Ohter limitations:
- modes with categorical splits not supported
:param model: H2O MOJO model loaded into memory (see below for example)
:param name: The name of the graph (type: GraphProto) in the produced ONNX model (type: ModelProto)
:param initial_types: a python list. Each element is a tuple of a variable name and a type defined in data_types.py
:param doc_string: A string attached onto the produced ONNX model
:param target_opset: number, for example, 7 for ONNX 1.2, and 8 for ONNX 1.3.
:param targeted_onnx: A string (for example, '1.1.2' and '1.2') used to specify the targeted ONNX version of the
produced model. If ONNXMLTools cannot find a compatible ONNX python package, an error may be thrown.
:param custom_conversion_functions: a dictionary for specifying the user customized conversion function
:param custom_shape_calculators: a dictionary for specifying the user customized shape calculator
:return: An ONNX model (type: ModelProto) which is equivalent to the input xgboost model
:examples:
>>> from onnxmltools.convert import convert_h2o
>>> file = open("/path/to/h2o_mojo.zip", "rb")
>>> mojo_content = file.read()
>>> file.close()
>>> h2o_onnx_model = convert_h2o(mojo_content)
'''
if name is None:
name = str(uuid4().hex)
if initial_types is None:
initial_types = [('input', FloatTensorType(shape=['None', 'None']))]
if isinstance(model, str):
model_path = model
else:
_, model_path = tempfile.mkstemp()
f = open(model_path, "wb")
f.write(model)
f.close()
mojo_str = h2o.print_mojo(model_path, format="json")
mojo_model = json.loads(mojo_str)
if mojo_model["params"]["algo"] != "gbm":
raise ValueError(
"Model type not supported (algo=%s). Only GBM Mojo supported for now." % mojo_model["params"]["algo"])
target_opset = target_opset if target_opset else get_maximum_opset_supported()
topology = parse_h2o(mojo_model, initial_types, target_opset, custom_conversion_functions, custom_shape_calculators)
topology.compile()
onnx_model = convert_topology(topology, name, doc_string, target_opset, targeted_onnx)
return onnx_model