Module hummingbird.ml.operator_converters.onnx.onnxml_array_feature_extractor
Converter for ONNX-ML Array Feature Extractor.
Expand source code Browse git
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
"""
Converter for ONNX-ML Array Feature Extractor.
"""
from onnxconverter_common.registration import register_converter
from .. import constants
from .._array_feature_extractor_implementations import ArrayFeatureExtractor
def convert_onnx_array_feature_extractor(operator, device, extra_config):
"""
Converter for `ai.onnx.ml.ArrayFeatureExtractor`.
Args:
operator: An operator wrapping a `ai.onnx.ml.ArrayFeatureExtractor` model
device: String defining the type of device the converted operator should be run on
extra_config: Extra configuration used to select the best conversion strategy
Returns:
A PyTorch model
"""
# TODO, this will be tested as part of the ai.onnx.ml.OneHotEncoder tests
column_indices = []
initializers = extra_config[constants.ONNX_INITIALIZERS]
column_indices = initializers[operator.raw_operator.origin.input[1]].int64_data
return ArrayFeatureExtractor(column_indices, device)
register_converter("ONNXMLArrayFeatureExtractor", convert_onnx_array_feature_extractor)
Functions
def convert_onnx_array_feature_extractor(operator, device, extra_config)
-
Converter for
ai.onnx.ml.ArrayFeatureExtractor
.Args
operator
- An operator wrapping a
ai.onnx.ml.ArrayFeatureExtractor
model device
- String defining the type of device the converted operator should be run on
extra_config
- Extra configuration used to select the best conversion strategy
Returns
A PyTorch model
Expand source code Browse git
def convert_onnx_array_feature_extractor(operator, device, extra_config): """ Converter for `ai.onnx.ml.ArrayFeatureExtractor`. Args: operator: An operator wrapping a `ai.onnx.ml.ArrayFeatureExtractor` model device: String defining the type of device the converted operator should be run on extra_config: Extra configuration used to select the best conversion strategy Returns: A PyTorch model """ # TODO, this will be tested as part of the ai.onnx.ml.OneHotEncoder tests column_indices = [] initializers = extra_config[constants.ONNX_INITIALIZERS] column_indices = initializers[operator.raw_operator.origin.input[1]].int64_data return ArrayFeatureExtractor(column_indices, device)