From 8dcbd5da8afa197b3392ca3205e18732590114fd Mon Sep 17 00:00:00 2001 From: matthias_oh Date: Mon, 20 Mar 2023 15:22:34 +0800 Subject: [PATCH 1/5] Fix: buildInputDictSimple dataType Stringtype() --- onnxmltools/convert/sparkml/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxmltools/convert/sparkml/utils.py b/onnxmltools/convert/sparkml/utils.py index c64986bb..bfe64176 100644 --- a/onnxmltools/convert/sparkml/utils.py +++ b/onnxmltools/convert/sparkml/utils.py @@ -33,7 +33,7 @@ def buildInputDictSimple(dataframe): import numpy result = {} for field in dataframe.schema.fields: - if str(field.dataType) == 'StringType': + if str(field.dataType) == 'StringType' or str(field.dataType) == 'StringType()': result[field.name] = dataframe.select(field.name).toPandas().values else: result[field.name] = dataframe.select(field.name).toPandas().values.astype(numpy.float32) From 13f52ac73c3140e57892b45ae05eb6e28d8f8507 Mon Sep 17 00:00:00 2001 From: matthias_oh Date: Mon, 20 Mar 2023 15:23:40 +0800 Subject: [PATCH 2/5] Fix: getTensorTypeFromSpark for batch input --- onnxmltools/convert/sparkml/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxmltools/convert/sparkml/utils.py b/onnxmltools/convert/sparkml/utils.py index bfe64176..6a9c0aae 100644 --- a/onnxmltools/convert/sparkml/utils.py +++ b/onnxmltools/convert/sparkml/utils.py @@ -15,7 +15,7 @@ def buildInitialTypesSimple(dataframe): def getTensorTypeFromSpark(sparktype): if sparktype == 'StringType' or sparktype == 'StringType()': - return StringTensorType([1, 1]) + return StringTensorType([None, 1]) elif sparktype == 'DecimalType' or sparktype == 'DecimalType()' \ or sparktype == 'DoubleType' or sparktype == 'DoubleType()' \ or sparktype == 'FloatType' or sparktype == 'FloatType()' \ @@ -24,7 +24,7 @@ def getTensorTypeFromSpark(sparktype): or sparktype == 'ShortType' or sparktype == 'ShortType()' \ or sparktype == 'ByteType' or sparktype == 'ByteType()' \ or sparktype == 'BooleanType' or sparktype == 'BooleanType()': - return FloatTensorType([1, 1]) + return FloatTensorType([None, 1]) else: raise TypeError("Cannot map this type to Onnx types: " + sparktype) From 174e5d9197aa13408fd1a1605414f603833701a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Fri, 28 Jul 2023 16:52:24 +0200 Subject: [PATCH 3/5] Support for lightgbm >= 4.0 (#634) * Fix issues raised with lightgbm 4.0 Signed-off-by: Xavier Dupre * update CI Signed-off-by: Xavier Dupre * remove use of eval Signed-off-by: Xavier Dupre * fix requirements Signed-off-by: Xavier Dupre * update requirements Signed-off-by: Xavier Dupre * fix handle Signed-off-by: Xavier Dupre * fix handle Signed-off-by: Xavier Dupre * fix scipy version Signed-off-by: Xavier Dupre * svm Signed-off-by: Xavier Dupre * remove allow_failure Signed-off-by: Xavier Dupre --------- Signed-off-by: Xavier Dupre --- .azure-pipelines/linux-conda-CI.yml | 114 ++++++------------ .azure-pipelines/win32-conda-CI.yml | 74 ++---------- onnxmltools/convert/lightgbm/_parse.py | 15 ++- .../lightgbm/operator_converters/LightGbm.py | 34 ++++-- onnxmltools/utils/tests_helper.py | 21 +--- onnxmltools/utils/utils_backend.py | 13 -- requirements-dev.txt | 5 +- .../test_cml_DictVectorizerConverter.py | 3 +- .../coreml/test_cml_GLMClassifierConverter.py | 6 +- ...st_cml_SupportVectorClassifierConverter.py | 3 +- ...est_cml_TreeEnsembleClassifierConverter.py | 3 +- ...l_TreeEnsembleRegressorConverterXGBoost.py | 3 +- ...htGbmTreeEnsembleConverters_hummingbird.py | 6 +- .../test_LightGbmTreeEnsembleConverters.py | 21 ++-- tests/sparkml/sparkml_test_utils.py | 6 +- tests/svmlib/test_SVMConverters.py | 64 +++++++--- 16 files changed, 155 insertions(+), 236 deletions(-) diff --git a/.azure-pipelines/linux-conda-CI.yml b/.azure-pipelines/linux-conda-CI.yml index c7de5ab8..90cb056f 100644 --- a/.azure-pipelines/linux-conda-CI.yml +++ b/.azure-pipelines/linux-conda-CI.yml @@ -15,90 +15,46 @@ jobs: strategy: matrix: - Python310-1140-RT1150-xgb175: + Python311-1140-RT1151-xgb175: + python.version: '3.11' + ONNX_PATH: 'onnx==1.14.0' #'-i https://test.pypi.org/simple/ onnx==1.14.0rc3' + ONNXRT_PATH: 'onnxruntime==1.15.1' + COREML_PATH: NONE + lightgbm.version: '>=4.0' + xgboost.version: '>=1.7.5' + numpy.version: '' + scipy.version: '' + + Python310-1140-RT1151-xgb175: python.version: '3.10' ONNX_PATH: 'onnx==1.14.0' #'-i https://test.pypi.org/simple/ onnx==1.14.0rc3' - ONNXRT_PATH: '-i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/ ort_nightly==1.15.0.dev20230502003' + ONNXRT_PATH: 'onnxruntime==1.15.1' COREML_PATH: NONE + lightgbm.version: '<4.0' xgboost.version: '>=1.7.5' numpy.version: '' + scipy.version: '' Python310-1140-RT1140-xgb175: python.version: '3.10' ONNX_PATH: 'onnx==1.14.0' #'-i https://test.pypi.org/simple/ onnx==1.14.0rc3' ONNXRT_PATH: onnxruntime==1.14.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' COREML_PATH: NONE + lightgbm.version: '<4.0' xgboost.version: '>=1.7.5' numpy.version: '' - # Python310-1130-RT1140-xgb173: - # python.version: '3.10' - # ONNX_PATH: 'onnx==1.13.0' #'-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.14.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: NONE - # xgboost.version: '>=1.7.3' - # numpy.version: '' - # Python310-1130-RT1131-xgb173: - # python.version: '3.10' - # ONNX_PATH: 'onnx==1.13.0' #'-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.13.1 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: NONE - # xgboost.version: '>=1.7.3' - # numpy.version: '' - # Python310-1120-RT1121-xgb161: - # python.version: '3.10' - # ONNX_PATH: 'onnx==1.12.0' #'-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.12.1 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: NONE - # xgboost.version: '==1.6.1' - # numpy.version: '' - # Python39-1120-RT1110-xgb161: - # python.version: '3.9' - # ONNX_PATH: 'onnx==1.12.0' #'-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.11.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: coremltools==6.3 # git+https://github.com/apple/coremltools@3.1 - # xgboost.version: '>=1.6.1' - # numpy.version: '' - # Python39-1120-RT1110-xgb160: - # python.version: '3.9' - # ONNX_PATH: 'onnx==1.12.0' #'-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.11.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: git+https://github.com/apple/coremltools@3.1 - # xgboost.version: '==1.6.0' - # numpy.version: '' - # Python39-1110-RT1110: - # python.version: '3.9' - # ONNX_PATH: onnx==1.11.0 # '-i https://test.pypi.org/simple/ onnx==1.9.101' - # ONNXRT_PATH: onnxruntime==1.11.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: git+https://github.com/apple/coremltools@3.1 - # numpy.version: '' - # Python39-1110-RT1100-xgb120: - # python.version: '3.9' - # ONNX_PATH: onnx==1.11.0 # '-i https://test.pypi.org/simple/ onnx==1.9.101' - # ONNXRT_PATH: onnxruntime==1.10.0 - # COREML_PATH: git+https://github.com/apple/coremltools@3.1 - # xgboost.version: '>=1.2' - # numpy.version: '' - # Python39-1101-RT190-xgb120: - # python.version: '3.9' - # ONNX_PATH: onnx==1.10.1 - # ONNXRT_PATH: onnxruntime==1.9.0 - # COREML_PATH: coremltools==4.0 # git+https://github.com/apple/coremltools@3.1 - # xgboost.version: '>=1.2' - # numpy.version: '<=1.23.5' - # Python39-190-RT180-xgb120: - # python.version: '3.9' - # ONNX_PATH: onnx==1.9.0 - # ONNXRT_PATH: onnxruntime==1.8.0 - # COREML_PATH: coremltools==4.0 # git+https://github.com/apple/coremltools@3.1 - # xgboost.version: '>=1.2' - # numpy.version: '<=1.23.5' - # Python38-181-RT170-xgb120: - # python.version: '3.8' - # ONNX_PATH: onnx==1.8.1 - # ONNXRT_PATH: onnxruntime==1.7.0 - # COREML_PATH: coremltools==4.0 # git+https://github.com/apple/coremltools@3.1 - # xgboost.version: '>=1.2' - # numpy.version: '<=1.23.5' + scipy.version: '' + + Python39-1140-RT1151-xgb175-scipy180: + python.version: '3.9' + ONNX_PATH: 'onnx==1.14.0' #'-i https://test.pypi.org/simple/ onnx==1.14.0rc3' + ONNXRT_PATH: 'onnxruntime==1.15.1' + COREML_PATH: NONE + lightgbm.version: '>=4.0' + xgboost.version: '>=1.7.5' + numpy.version: '' + scipy.version: '==1.8.0' + maxParallel: 3 @@ -120,16 +76,18 @@ jobs: - script: | python -m pip install --upgrade pip - pip install xgboost$(xgboost.version) + pip install "xgboost$(xgboost.version)" + pip install "lightgbm$(lightgbm.version)" pip install $(ONNX_PATH) pip install $(ONNXRT_PATH) pip install "numpy$(numpy.version)" + pip install "scipy$(scipy.version)" displayName: 'Install xgboost, onnxruntime' - script: | python -m pip install coloredlogs flatbuffers packaging sympy numpy protobuf python -m pip install $(ONNXRT_PATH) - displayName: 'Install ort-nightly' + displayName: 'Install onnxruntime' - script: | pip install flake8 @@ -170,11 +128,6 @@ jobs: pytest tests/sparkml --durations=0 displayName: 'pytest - sparkml' - - script: | - export PYTHONPATH=. - pytest tests/svmlib --durations=0 - displayName: 'pytest - svmlib' - - script: | export PYTHONPATH=. pytest tests/utils --durations=0 @@ -191,6 +144,11 @@ jobs: pytest tests/h2o --durations=0 displayName: 'pytest - h2o' + - script: | + export PYTHONPATH=. + pytest tests/svmlib --durations=0 + displayName: 'pytest - svmlib' + - script: | pip install torch --extra-index-url https://download.pytorch.org/whl/cpu pip install hummingbird-ml --no-deps diff --git a/.azure-pipelines/win32-conda-CI.yml b/.azure-pipelines/win32-conda-CI.yml index 115d244b..68352693 100644 --- a/.azure-pipelines/win32-conda-CI.yml +++ b/.azure-pipelines/win32-conda-CI.yml @@ -15,10 +15,17 @@ jobs: strategy: matrix: - Python310-1150-RT1140: + Python311-1140-RT1151: + python.version: '3.11' + ONNX_PATH: 'onnx==1.14.0' # '-i https://test.pypi.org/simple/ onnx==1.14.0rc3' + ONNXRT_PATH: 'onnxruntime==1.15.1' + COREML_PATH: NONE + numpy.version: '' + + Python310-1140-RT1151: python.version: '3.10' ONNX_PATH: 'onnx==1.14.0' # '-i https://test.pypi.org/simple/ onnx==1.14.0rc3' - ONNXRT_PATH: '-i https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ORT-Nightly/pypi/simple/ ort_nightly==1.15.0.dev20230502003' + ONNXRT_PATH: 'onnxruntime==1.15.1' COREML_PATH: NONE numpy.version: '' @@ -29,69 +36,6 @@ jobs: COREML_PATH: NONE numpy.version: '' - # Python310-1130-RT1140: - # python.version: '3.10' - # ONNX_PATH: 'onnx==1.13.0' # '-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.14.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: NONE - # numpy.version: '' - - # Python310-1130-RT1131: - # python.version: '3.10' - # ONNX_PATH: 'onnx==1.13.0' # '-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.13.1 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: NONE - # numpy.version: '' - - # Python39-1120-RT1110: - # python.version: '3.9' - # ONNX_PATH: 'onnx==1.12.0' # '-i https://test.pypi.org/simple/ onnx==1.12.0rc4' - # ONNXRT_PATH: onnxruntime==1.11.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: coremltools==6.3 # git+https://github.com/apple/coremltools@3.1 - # numpy.version: '' - - # Python39-1110-RT1110: - # python.version: '3.9' - # ONNX_PATH: onnx==1.11.0 # '-i https://test.pypi.org/simple/ onnx==1.9.101' - # ONNXRT_PATH: onnxruntime==1.11.0 #'-i https://test.pypi.org/simple/ ort-nightly==1.11.0.dev20220311003' - # COREML_PATH: coremltools==6.2 # 6.0 doesn't work - # numpy.version: '' - - # Python39-1110-RT190: - # python.version: '3.9' - # ONNX_PATH: 'onnx==1.11.0' # '-i https://test.pypi.org/simple/ onnx==1.9.101' - # ONNXRT_PATH: onnxruntime==1.10.0 - # COREML_PATH: coremltools==5.2 # git+https://github.com/apple/coremltools@3.1 - # numpy.version: '' - - # Python39-1102-RT190: - # python.version: '3.9' - # ONNX_PATH: 'onnx==1.10.2' # '-i https://test.pypi.org/simple/ onnx==1.9.101' - # ONNXRT_PATH: onnxruntime==1.9.0 - # COREML_PATH: coremltools==5.2 # 5.0 doesn't work - # numpy.version: '<=1.23.5' - - # Python39-190-RT181: - # python.version: '3.9' - # ONNX_PATH: 'onnx==1.9.0' - # ONNXRT_PATH: onnxruntime==1.8.1 - # COREML_PATH: coremltools==4.0 # git+https://github.com/apple/coremltools@3.1 - # numpy.version: '<=1.23.5' - - # Python39-190-RT180: - # python.version: '3.9' - # ONNX_PATH: onnx==1.9.0 - # ONNXRT_PATH: onnxruntime==1.8.0 - # COREML_PATH: coremltools==4.0 # git+https://github.com/apple/coremltools@3.1 - # numpy.version: '<=1.23.5' - - # Python38-181-RT170: - # python.version: '3.8' - # ONNX_PATH: onnx==1.8.1 - # ONNXRT_PATH: onnxruntime==1.7.0 - # COREML_PATH: coremltools==4.0 # git+https://github.com/apple/coremltools@3.1 - # numpy.version: '<=1.23.5' - maxParallel: 3 steps: diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py index abace2b6..aee8869c 100644 --- a/onnxmltools/convert/lightgbm/_parse.py +++ b/onnxmltools/convert/lightgbm/_parse.py @@ -34,7 +34,10 @@ def __init__(self, booster): else: raise NotImplementedError( 'Unsupported LightGbm objective: %r.' % self.objective_) - average_output = self.booster_.attr('average_output') + try: + average_output = self.booster_.attr('average_output') + except AttributeError: + average_output = self.booster_.params.get("average_output", None) if average_output: self.boosting_type = 'rf' else: @@ -47,7 +50,10 @@ def _generate_classes(booster): if isinstance(booster, dict): num_class = booster['num_class'] else: - num_class = booster.attr('num_class') + try: + num_class = booster.attr('num_class') + except AttributeError: + num_class = booster.params.get('num_class', None) if num_class is None: dp = booster.dump_model(num_iteration=1) num_class = dp['num_class'] @@ -59,7 +65,10 @@ def get_objective(self): "Returns the objective." if hasattr(self, 'objective_') and self.objective_ is not None: return self.objective_ - objective = self.booster_.attr('objective') + try: + objective = self.booster_.attr('objective') + except AttributeError: + objective = self.booster_.params.get("objective", None) if objective is not None: return objective dp = self.booster_.dump_model(num_iteration=1) diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py index 7a3cbfe8..7ca59985 100644 --- a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py +++ b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py @@ -264,20 +264,34 @@ def dump_booster_model(self, num_iteration=None, start_iteration=0, """ if getattr(self, 'is_mock', False): return self.dump_model(), None - from lightgbm.basic import ( - _LIB, FEATURE_IMPORTANCE_TYPE_MAPPER, _safe_call, - json_default_with_numpy) + from lightgbm.basic import _LIB, _safe_call + try: + # lightgbm >= 4.0 + from lightgbm.basic import ( + _FEATURE_IMPORTANCE_TYPE_MAPPER as FITM, + _json_default_with_numpy as jdwn, + ) + except ImportError: + # lightgbm < 4.0 + from lightgbm.basic import ( + FEATURE_IMPORTANCE_TYPE_MAPPER as FITM, + json_default_with_numpy as jdwn, + ) if num_iteration is None: num_iteration = self.best_iteration - importance_type_int = FEATURE_IMPORTANCE_TYPE_MAPPER[importance_type] + importance_type_int = FITM[importance_type] buffer_len = 1 << 20 tmp_out_len = ctypes.c_int64(0) string_buffer = ctypes.create_string_buffer(buffer_len) ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)]) if verbose >= 2: print("[dump_booster_model] call CAPI: LGBM_BoosterDumpModel") + try: + handle = self._handle + except AttributeError: + handle = self.handle _safe_call(_LIB.LGBM_BoosterDumpModel( - self.handle, + handle, ctypes.c_int(start_iteration), ctypes.c_int(num_iteration), ctypes.c_int(importance_type_int), @@ -290,8 +304,14 @@ def dump_booster_model(self, num_iteration=None, start_iteration=0, string_buffer = ctypes.create_string_buffer(actual_len) ptr_string_buffer = ctypes.c_char_p( *[ctypes.addressof(string_buffer)]) + try: + # lightgbm >= 4.0 + handle = self._handle + except AttributeError: + # lightgbm < 4.0 + handle = self.handle _safe_call(_LIB.LGBM_BoosterDumpModel( - self.handle, + handle, ctypes.c_int(start_iteration), ctypes.c_int(num_iteration), ctypes.c_int(importance_type_int), @@ -359,7 +379,7 @@ def hook(self, obj): info=info, n_trees=self.num_trees(), verbose=verbose) ret['pandas_categorical'] = json.loads( json.dumps(self.pandas_categorical, - default=json_default_with_numpy)) + default=jdwn)) if verbose >= 2: print("[dump_booster_model] end.") return ret, info diff --git a/onnxmltools/utils/tests_helper.py b/onnxmltools/utils/tests_helper.py index dc2b4a4b..2c33ff7e 100644 --- a/onnxmltools/utils/tests_helper.py +++ b/onnxmltools/utils/tests_helper.py @@ -8,7 +8,7 @@ from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from ..convert.common.data_types import FloatTensorType -from .utils_backend import compare_backend, extract_options, evaluate_condition, is_backend_enabled +from .utils_backend import compare_backend, extract_options, is_backend_enabled TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) @@ -183,22 +183,9 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, if not is_backend_enabled(b): continue if isinstance(allow_failure, str): - allow = evaluate_condition(b, allow_failure) - else: - allow = allow_failure - if allow is None: - output = compare_backend(b, runtime_test, options=extract_options(basename), - context=context, verbose=verbose) - else: - try: - output = compare_backend(b, runtime_test, options=extract_options(basename), - context=context, verbose=verbose) - except AssertionError as e: - if isinstance(allow, bool) and allow: - warnings.warn("Issue with '{0}' due to {1}".format(basename, e)) - continue - else: - raise e + raise NotImplementedError("allow_failure is deprecated.") + output = compare_backend(b, runtime_test, options=extract_options(basename), + context=context, verbose=verbose) if output is not None: dest = os.path.join(folder, basename + ".backend.{0}.pkl".format(b)) names.append(dest) diff --git a/onnxmltools/utils/utils_backend.py b/onnxmltools/utils/utils_backend.py index 89f3974d..70ef0fd0 100644 --- a/onnxmltools/utils/utils_backend.py +++ b/onnxmltools/utils/utils_backend.py @@ -25,19 +25,6 @@ class OnnxRuntimeAssertionError(AssertionError): pass -def evaluate_condition(backend, condition): - """ - Evaluates a condition such as - ``pv.Version(onnxruntime.__version__) <= pv.Version('0.1.3')`` - """ - if backend == "onnxruntime": - import onnxruntime - import onnx - return eval(condition) - else: - raise NotImplementedError("Not implemented for backend '{0}'".format(backend)) - - def is_backend_enabled(backend): """ Tells if a backend is enabled. diff --git a/requirements-dev.txt b/requirements-dev.txt index 8e1b04c3..d4b6e564 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -9,13 +9,12 @@ mleap numpy openpyxl pandas -protobuf<4.0 psutil pyspark pytest pytest-cov pytest-spark -scikit-learn==1.1.0 -scipy==1.10.0 +scikit-learn>=1.2.0 +scipy wheel xgboost==1.7.5 diff --git a/tests/coreml/test_cml_DictVectorizerConverter.py b/tests/coreml/test_cml_DictVectorizerConverter.py index dd6130c6..8818f1ae 100644 --- a/tests/coreml/test_cml_DictVectorizerConverter.py +++ b/tests/coreml/test_cml_DictVectorizerConverter.py @@ -47,8 +47,7 @@ def test_dict_vectorizer(self): sklearn.__version__, sys.platform)) from e model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - dump_data_and_model(data, model, model_onnx, basename="CmlDictVectorizer-OneOff-SkipDim1", - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.3.0')") + dump_data_and_model(data, model, model_onnx, basename="CmlDictVectorizer-OneOff-SkipDim1") if __name__ == "__main__": diff --git a/tests/coreml/test_cml_GLMClassifierConverter.py b/tests/coreml/test_cml_GLMClassifierConverter.py index 26ae083e..66d654f9 100644 --- a/tests/coreml/test_cml_GLMClassifierConverter.py +++ b/tests/coreml/test_cml_GLMClassifierConverter.py @@ -53,8 +53,7 @@ def test_glm_classifier(self): lr_onnx = convert(lr_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(lr_onnx is not None) self.validate_zipmap(lr_onnx) - dump_data_and_model(X.astype(numpy.float32), lr, lr_onnx, basename="CmlbinLogitisticRegression", - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.3.0')") + dump_data_and_model(X.astype(numpy.float32), lr, lr_onnx, basename="CmlbinLogitisticRegression") # Ensure there is a probability output svm = LinearSVC() @@ -63,8 +62,7 @@ def test_glm_classifier(self): svm_onnx = convert(svm_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(svm_onnx is not None) self.validate_zipmap(svm_onnx) - dump_data_and_model(X.astype(numpy.float32), svm, svm_onnx, basename="CmlBinLinearSVC-NoProb", - allow_failure=True) + dump_data_and_model(X.astype(numpy.float32), svm, svm_onnx, basename="CmlBinLinearSVC-NoProb") if __name__ == "__main__": diff --git a/tests/coreml/test_cml_SupportVectorClassifierConverter.py b/tests/coreml/test_cml_SupportVectorClassifierConverter.py index e9f3dca3..a3f884b2 100644 --- a/tests/coreml/test_cml_SupportVectorClassifierConverter.py +++ b/tests/coreml/test_cml_SupportVectorClassifierConverter.py @@ -77,8 +77,7 @@ def test_support_vector_classifier_binary_no_prob(self): nodes = svm_onnx.graph.node self.assertEqual(len(nodes), 1) self._check_model_outputs(svm_onnx, ['classLabel']) - dump_data_and_model(X, svm, svm_onnx, basename="CmlBinSVC-Out0", - allow_failure=True) + dump_data_and_model(X, svm, svm_onnx, basename="CmlBinSVC-Out0") @unittest.skipIf( pv.Version(coremltools.__version__) > pv.Version("3.1"), diff --git a/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py b/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py index 2c61d11f..74a5cfdf 100644 --- a/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py +++ b/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py @@ -46,8 +46,7 @@ def test_tree_ensemble_classifier(self): model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) self.validate_zipmap(model_onnx) - dump_data_and_model(X, model, model_onnx, basename="CmlBinRandomForestClassifier", - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.3.0')") + dump_data_and_model(X, model, model_onnx, basename="CmlBinRandomForestClassifier") if __name__ == "__main__": diff --git a/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py b/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py index 01e14745..728e6361 100644 --- a/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py +++ b/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py @@ -48,8 +48,7 @@ def test_tree_ensemble_regressor_xgboost(self): if sys.version_info[0] >= 3: # python 2.7 returns TypeError: can't pickle instancemethod objects dump_data_and_model(X.astype(numpy.float32), model, model_onnx, - basename="CmlXGBoostRegressor-OneOff-Reshape", - allow_failure=True) + basename="CmlXGBoostRegressor-OneOff-Reshape") if __name__ == "__main__": diff --git a/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py b/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py index 1f46d014..7565711c 100644 --- a/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py +++ b/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py @@ -51,7 +51,6 @@ def test_lightgbm_booster_classifier(self): target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) dump_data_and_model(X, model, model_onnx, - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.3.0')", basename=prefix + "BoosterBin" + model.__class__.__name__) @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") @@ -76,7 +75,6 @@ def test_lightgbm_booster_classifier_zipmap(self): [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) dump_data_and_model(X, model, model_onnx, - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.3.0')", basename=prefix + "BoosterBin" + model.__class__.__name__) @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") @@ -92,9 +90,8 @@ def test_lightgbm_booster_multi_classifier(self): [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) dump_data_and_model(X, model, model_onnx, - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.3.0')", basename=prefix + "BoosterBin" + model.__class__.__name__) - sess = InferenceSession(model_onnx.SerializeToString()) + sess = InferenceSession(model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]) out = sess.get_outputs() names = [o.name for o in out] assert names == ['label', 'probabilities'] @@ -112,7 +109,6 @@ def test_lightgbm_booster_regressor(self): [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) dump_data_and_model(X, model, model_onnx, - allow_failure="pv.Version(onnx.__version__) < pv.Version('1.0.0')", basename=prefix + "BoosterBin" + model.__class__.__name__) # Base test implementation comparing ONNXML and ONNX models. diff --git a/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py b/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py index 67cf6565..241642fe 100644 --- a/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py +++ b/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py @@ -24,10 +24,13 @@ class TestLightGbmTreeEnsembleModels(unittest.TestCase): - def test_lightgbm_classifier(self): + def test_lightgbm_classifier_binary(self): model = LGBMClassifier(n_estimators=3, min_child_samples=1, num_thread=1) - dump_binary_classification(model, allow_failure=pv.Version(onnx.__version__) < pv.Version('1.3.0')) - dump_multiple_classification(model, allow_failure=pv.Version(onnx.__version__) < pv.Version('1.3.0')) + dump_binary_classification(model) + + def test_lightgbm_classifier_multiple(self): + model = LGBMClassifier(n_estimators=3, min_child_samples=1, num_thread=1) + dump_multiple_classification(model) def test_lightgbm_classifier_zipmap(self): X = [[0, 1], [1, 1], [2, 0], [1, 2]] @@ -52,7 +55,7 @@ def test_lightgbm_classifier_nozipmap(self): assert "zipmap" not in str(onx).lower() onxs = onx[0].SerializeToString() try: - sess = onnxruntime.InferenceSession(onxs) + sess = onnxruntime.InferenceSession(onxs, providers=["CPUExecutionProvider"]) except Exception as e: raise AssertionError( "Model cannot be loaded by onnxruntime due to %r\n%s." % ( @@ -74,11 +77,11 @@ def test_lightgbm_classifier_nozipmap2(self): assert "zipmap" not in str(onx).lower() onxs = onx.SerializeToString() try: - sess = onnxruntime.InferenceSession(onxs) + sess = onnxruntime.InferenceSession(onxs, providers=["CPUExecutionProvider"]) except Exception as e: raise AssertionError( "Model cannot be loaded by onnxruntime due to %r\n%s." % ( - e, onx[0])) + e, onx)) exp = model.predict(X), model.predict_proba(X) got = sess.run(None, {'X': X}) assert_almost_equal(exp[0], got[0]) @@ -108,7 +111,6 @@ def test_lightgbm_booster_classifier(self): [('input', FloatTensorType([None, 2]))], target_opset=TARGET_OPSET) dump_data_and_model(X, model, model_onnx, - allow_failure=pv.Version(onnx.__version__) < pv.Version('1.3.0'), basename=prefix + "BoosterBin" + model.__class__.__name__) def test_lightgbm_booster_classifier_nozipmap(self): @@ -124,7 +126,6 @@ def test_lightgbm_booster_classifier_nozipmap(self): zipmap=False, target_opset=TARGET_OPSET) assert "zipmap" not in str(model_onnx).lower() dump_data_and_model(X, model, model_onnx, - allow_failure=pv.Version(onnx.__version__) < pv.Version('1.3.0'), basename=prefix + "BoosterBin" + model.__class__.__name__) def test_lightgbm_booster_classifier_zipmap(self): @@ -140,7 +141,6 @@ def test_lightgbm_booster_classifier_zipmap(self): target_opset=TARGET_OPSET) assert "zipmap" in str(model_onnx).lower() dump_data_and_model(X, model, model_onnx, - allow_failure=pv.Version(onnx.__version__) < pv.Version('1.3.0'), basename=prefix + "BoosterBin" + model.__class__.__name__) def test_lightgbm_booster_multi_classifier(self): @@ -155,14 +155,13 @@ def test_lightgbm_booster_multi_classifier(self): [('input', FloatTensorType([None, 2]))], target_opset=TARGET_OPSET) dump_data_and_model(X, model, model_onnx, - allow_failure=pv.Version(onnx.__version__) < pv.Version('1.3.0'), basename=prefix + "BoosterBin" + model.__class__.__name__) try: from onnxruntime import InferenceSession except ImportError: # onnxruntime not installed (python 2.7) return - sess = InferenceSession(model_onnx.SerializeToString()) + sess = InferenceSession(model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]) out = sess.get_outputs() names = [o.name for o in out] assert names == ['label', 'probabilities'] diff --git a/tests/sparkml/sparkml_test_utils.py b/tests/sparkml/sparkml_test_utils.py index f71b2f2a..45db647f 100644 --- a/tests/sparkml/sparkml_test_utils.py +++ b/tests/sparkml/sparkml_test_utils.py @@ -13,7 +13,7 @@ from pyspark.ml.linalg import VectorUDT from pyspark.sql.types import ArrayType, FloatType, DoubleType from onnxmltools.utils.utils_backend import ( - compare_backend, extract_options, evaluate_condition, is_backend_enabled, + compare_backend, extract_options, is_backend_enabled, OnnxRuntimeAssertionError, compare_outputs, ExpectedAssertionError) from onnxmltools.utils.utils_backend_onnxruntime import _create_column @@ -272,9 +272,7 @@ def dump_data_and_sparkml_model(input, expected, model, onnx=None, basename="mod if not is_backend_enabled(b): continue if isinstance(allow_failure, str): - allow = evaluate_condition(b, allow_failure) - else: - allow = allow_failure + raise NotImplementedError("allow_failure is deprecated.") if allow is None: output = compare_backend(b, runtime_test, options=extract_options(basename), context=context, verbose=verbose) diff --git a/tests/svmlib/test_SVMConverters.py b/tests/svmlib/test_SVMConverters.py index 5fbfc276..49339529 100644 --- a/tests/svmlib/test_SVMConverters.py +++ b/tests/svmlib/test_SVMConverters.py @@ -5,15 +5,20 @@ """ import tempfile import numpy +import scipy import os try: from libsvm.svm import C_SVC as SVC, EPSILON_SVR as SVR, NU_SVC as NuSVC, NU_SVR as NuSVR import libsvm.svm as svm import libsvm.svmutil as svmutil + DISABLED = False except ImportError: - import svm - from svm import C_SVC as SVC, EPSILON_SVR as SVR, NU_SVC as NuSVC, NU_SVR as NuSVR - import svmutil + try: + import svm + from svm import C_SVC as SVC, EPSILON_SVR as SVR, NU_SVC as NuSVC, NU_SVR as NuSVR + import svmutil + except ImportError: + DISABLED = True import onnxruntime import numpy as np @@ -22,7 +27,6 @@ from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from sklearn.datasets import load_iris -from onnxmltools.convert.libsvm import convert from onnxmltools.convert.common.data_types import FloatTensorType from onnxmltools.utils import dump_data_and_model @@ -33,6 +37,15 @@ # This was recently added. noprint = None +if not DISABLED: + try: + from scipy import ctypeslib + except ImportError: + DISABLED = True + +if not DISABLED: + from onnxmltools.convert.libsvm import convert + TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) @@ -114,6 +127,8 @@ def decision_function(self, X): class TestSvmLibSVM(unittest.TestCase): + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmc_linear(self): iris = load_iris() @@ -136,9 +151,10 @@ def test_convert_svmc_linear(self): node = convert(libsvm_model, "LibSvmSvmcLinear", [('input', FloatTensorType())], target_opset=TARGET_OPSET) self.assertTrue(node is not None) dump_data_and_model(X[:5].astype(numpy.float32), SkAPIClProba2(libsvm_model), node, - basename="LibSvmSvmcLinear-Dec2", - allow_failure=pv.Version(onnxruntime.__version__) < pv.Version('0.5.0')) + basename="LibSvmSvmcLinear-Dec2") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmc(self): iris = load_iris() @@ -163,6 +179,8 @@ def test_convert_svmc(self): dump_data_and_model(X[:5].astype(numpy.float32), SkAPIClProba2(libsvm_model), node, basename="LibSvmSvmc-Dec2") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmr_linear(self): iris = load_iris() @@ -184,6 +202,8 @@ def test_convert_svmr_linear(self): dump_data_and_model(X[:5].astype(numpy.float32), SkAPIReg(libsvm_model), node, basename="LibSvmSvmrLinear-Dec3") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmr(self): iris = load_iris() @@ -206,6 +226,8 @@ def test_convert_svmr(self): dump_data_and_model(X[:5].astype(numpy.float32), SkAPIReg(libsvm_model), node, basename="LibSvmSvmr") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_nusvmr(self): iris = load_iris() @@ -228,6 +250,8 @@ def test_convert_nusvmr(self): dump_data_and_model(X[:5].astype(numpy.float32), SkAPIReg(libsvm_model), node, basename="LibSvmNuSvmr") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_nusvmc(self): iris = load_iris() @@ -251,9 +275,10 @@ def test_convert_nusvmc(self): target_opset=TARGET_OPSET) self.assertTrue(node is not None) dump_data_and_model(X[:5].astype(numpy.float32), SkAPIClProba2(libsvm_model), node, - basename="LibSvmNuSvmc-Dec2", - allow_failure=pv.Version(onnxruntime.__version__) <= pv.Version('0.1.3')) + basename="LibSvmNuSvmc-Dec2") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmc_linear_raw(self): iris = load_iris() @@ -278,9 +303,10 @@ def test_convert_svmc_linear_raw(self): self.assertTrue(node is not None) # known svm runtime dimension error in ONNX Runtime dump_data_and_model(X[:5].astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmSvmcLinearRaw-Dec3", verbose=False, - allow_failure=pv.Version(onnxruntime.__version__) < pv.Version('0.5.0')) + basename="LibSvmSvmcLinearRaw-Dec3", verbose=False) + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmc_raw(self): iris = load_iris() @@ -305,10 +331,11 @@ def test_convert_svmc_raw(self): target_opset=TARGET_OPSET) self.assertTrue(node is not None) dump_data_and_model(X[:5].astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmSvmcRaw", - allow_failure=pv.Version(onnxruntime.__version__) < pv.Version('0.5.0')) + basename="LibSvmSvmcRaw") @unittest.skip(reason="libsvm crashes.") + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_nusvmc_linear_raw(self): iris = load_iris() @@ -333,9 +360,10 @@ def test_convert_nusvmc_linear_raw(self): self.assertTrue(node is not None) X2 = numpy.vstack([X[:5], X[60:65]]) # 5x0, 5x1 dump_data_and_model(X2.astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmNuSvmcRaw", verbose=False, - allow_failure=pv.Version(onnxruntime.__version__) <= pv.Version('0.1.3')) + basename="LibSvmNuSvmcRaw", verbose=False) + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmc_rbf_raw_multi(self): iris = load_iris() @@ -360,9 +388,10 @@ def test_convert_svmc_rbf_raw_multi(self): self.assertTrue(node is not None) X2 = numpy.vstack([X[:2], X[60:62], X[110:112], X[147:149]]) # 5x0, 5x1 dump_data_and_model(X2.astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmNuSvmcRaw", verbose=False, - allow_failure=pv.Version(onnxruntime.__version__) <= pv.Version('0.1.3')) + basename="LibSvmNuSvmcRaw", verbose=False) + @unittest.skipIf(DISABLED, + reason="svmlib not really maintained") def test_convert_svmc_linear_raw_multi(self): iris = load_iris() @@ -387,8 +416,7 @@ def test_convert_svmc_linear_raw_multi(self): self.assertTrue(node is not None) X2 = numpy.vstack([X[:2], X[60:62], X[110:112], X[147:149]]) # 5x0, 5x1 dump_data_and_model(X2.astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmSvmcRaw-Dec3", verbose=False, - allow_failure=pv.Version(onnxruntime.__version__) <= pv.Version('0.1.3')) + basename="LibSvmSvmcRaw-Dec3", verbose=False) if __name__ == "__main__": From 2b709760625838bd412a6e3c6e6d6532e7690971 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Xavier=20Dupr=C3=A9?= Date: Tue, 1 Aug 2023 09:59:49 +0200 Subject: [PATCH 4/5] Apply black and ruff. (#635) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Apply black and ruff. Signed-off-by: Xavier Dupre * lint Signed-off-by: Xavier Dupre * update ci Signed-off-by: Xavier Dupre * fix req Signed-off-by: Xavier Dupre * restore __init__.py Signed-off-by: Xavier Dupre * ci Signed-off-by: Xavier Dupre * requirements Signed-off-by: Xavier Dupre * fix handle Signed-off-by: Xavier Dupre * disable coremltools test Signed-off-by: Xavier Dupre * fix missing import Signed-off-by: Xavier Dupre * fix missing import Signed-off-by: Xavier Dupre * import Signed-off-by: Xavier Dupre --------- Signed-off-by: Xavier Dupre Signed-off-by: Xavier Dupré Signed-off-by: Xavier Dupre --- .azure-pipelines/linux-CI-nightly.yml | 66 -- .azure-pipelines/linux-conda-CI.yml | 5 - .azure-pipelines/win32-CI-nightly.yml | 64 -- .azure-pipelines/win32-conda-CI.yml | 5 - .flake8 | 10 - .github/workflows/black-ruff.yml | 16 + docs/conf.py | 43 +- docs/examples/plot_convert_h2o.py | 27 +- docs/examples/plot_convert_keras.py | 34 +- docs/examples/plot_convert_libsvm.py | 26 +- docs/examples/plot_convert_lightgbm.py | 33 +- docs/examples/plot_convert_sklearn.py | 29 +- docs/examples/plot_convert_sparkml.py | 51 +- docs/examples/plot_convert_xgboost.py | 33 +- onnxmltools/convert/common/_container.py | 13 +- onnxmltools/convert/common/interface.py | 3 +- onnxmltools/convert/common/onnx_ex.py | 6 +- onnxmltools/convert/common/tree_ensemble.py | 4 +- onnxmltools/convert/common/utils.py | 2 + onnxmltools/convert/coreml/_parse.py | 582 ++++++++----- onnxmltools/convert/coreml/convert.py | 72 +- .../ArrayFeatureExtractor.py | 17 +- .../operator_converters/DictVectorizer.py | 21 +- .../operator_converters/FeatureVectorizer.py | 33 +- .../operator_converters/GLMClassifier.py | 118 ++- .../operator_converters/GLMRegressor.py | 32 +- .../coreml/operator_converters/Identity.py | 9 +- .../coreml/operator_converters/Imputer.py | 30 +- .../coreml/operator_converters/Normalizer.py | 20 +- .../operator_converters/OneHotEncoder.py | 25 +- .../convert/coreml/operator_converters/SVC.py | 133 +-- .../convert/coreml/operator_converters/SVR.py | 63 +- .../coreml/operator_converters/Scaler.py | 18 +- .../operator_converters/TensorToLabel.py | 83 +- .../TensorToProbabilityMap.py | 93 ++- .../operator_converters/TreeEnsemble.py | 240 ++++-- .../neural_network/Activation.py | 140 +++- .../operator_converters/neural_network/Add.py | 58 +- .../neural_network/Average.py | 11 +- .../neural_network/BatchNorm.py | 100 ++- .../neural_network/Bias.py | 31 +- .../neural_network/BidirectionalLSTM.py | 363 +++++--- .../neural_network/Concat.py | 12 +- .../neural_network/Convolution.py | 60 +- .../neural_network/Crop.py | 22 +- .../operator_converters/neural_network/Dot.py | 61 +- .../neural_network/Embed.py | 91 +- .../neural_network/Flatten.py | 24 +- .../operator_converters/neural_network/GRU.py | 172 ++-- .../neural_network/ImageScaler.py | 43 +- .../neural_network/InnerProduct.py | 85 +- .../neural_network/L2Normalize.py | 18 +- .../operator_converters/neural_network/LRN.py | 18 +- .../neural_network/LSTM.py | 247 ++++-- .../neural_network/LoadConstant.py | 23 +- .../neural_network/LoadConstantND.py | 23 +- .../operator_converters/neural_network/Max.py | 11 +- .../neural_network/MeanImage.py | 28 +- .../neural_network/MeanVarianceNorm.py | 14 +- .../operator_converters/neural_network/Min.py | 10 +- .../neural_network/Multiply.py | 58 +- .../operator_converters/neural_network/Pad.py | 34 +- .../neural_network/Permute.py | 12 +- .../neural_network/Pool.py | 324 ++++--- .../neural_network/Reduce.py | 42 +- .../neural_network/ReorganizeData.py | 15 +- .../neural_network/Reshape.py | 39 +- .../neural_network/ReshapeStatic.py | 29 +- .../neural_network/Scale.py | 93 ++- .../neural_network/SequenceRepeat.py | 16 +- .../neural_network/SimpleRNN.py | 208 +++-- .../neural_network/Slice.py | 59 +- .../neural_network/Softmax.py | 6 +- .../neural_network/Split.py | 18 +- .../neural_network/UnaryFunction.py | 106 ++- .../neural_network/Upsample.py | 20 +- .../ArrayFeatureExtractor.py | 16 +- .../coreml/shape_calculators/Classifier.py | 104 ++- .../shape_calculators/DictVectorizer.py | 27 +- .../shape_calculators/FeatureVectorizer.py | 54 +- .../coreml/shape_calculators/Identity.py | 9 +- .../coreml/shape_calculators/OneHotEncoder.py | 20 +- .../coreml/shape_calculators/Regressor.py | 43 +- .../coreml/shape_calculators/TensorToLabel.py | 27 +- .../TensorToProbabilityMap.py | 56 +- .../neural_network/BatchNorm.py | 13 +- .../neural_network/BidirectionalLSTM.py | 43 +- .../neural_network/Concat.py | 34 +- .../neural_network/Convolution.py | 98 ++- .../shape_calculators/neural_network/Crop.py | 17 +- .../shape_calculators/neural_network/Dot.py | 13 +- .../shape_calculators/neural_network/Embed.py | 23 +- .../neural_network/Flatten.py | 17 +- .../shape_calculators/neural_network/GRU.py | 36 +- .../neural_network/IdentityFloat.py | 31 +- .../neural_network/InnerProduct.py | 22 +- .../shape_calculators/neural_network/LSTM.py | 25 +- .../neural_network/LoadConstant.py | 9 +- .../neural_network/LoadConstantND.py | 9 +- .../shape_calculators/neural_network/Merge.py | 33 +- .../shape_calculators/neural_network/Pad.py | 11 +- .../neural_network/Permute.py | 18 +- .../shape_calculators/neural_network/Pool.py | 44 +- .../neural_network/Reduce.py | 12 +- .../neural_network/ReorganizeData.py | 50 +- .../neural_network/Reshape.py | 12 +- .../neural_network/ReshapeStatic.py | 12 +- .../neural_network/SequenceRepeat.py | 22 +- .../shape_calculators/neural_network/Slice.py | 24 +- .../shape_calculators/neural_network/Split.py | 19 +- .../neural_network/Upsample.py | 15 +- .../neural_network/__init__.py | 3 +- onnxmltools/convert/h2o/_parse.py | 35 +- onnxmltools/convert/h2o/convert.py | 62 +- .../convert/h2o/operator_converters/h2o.py | 140 ++-- .../h2o/shape_calculators/h2otreemojo.py | 8 +- onnxmltools/convert/libsvm/__init__.py | 1 + onnxmltools/convert/libsvm/_parse.py | 45 +- onnxmltools/convert/libsvm/convert.py | 53 +- .../operator_converters/SVMConverter.py | 161 ++-- .../libsvm/shape_calculators/Classifier.py | 18 +- .../libsvm/shape_calculators/Regressor.py | 11 +- onnxmltools/convert/lightgbm/_parse.py | 146 ++-- onnxmltools/convert/lightgbm/convert.py | 84 +- .../lightgbm/operator_converters/LightGbm.py | 788 +++++++++++------- .../lightgbm/shape_calculators/Classifier.py | 27 +- .../lightgbm/shape_calculators/Regressor.py | 2 +- onnxmltools/convert/main.py | 395 ++++++--- onnxmltools/convert/sparkml/_parse.py | 92 +- onnxmltools/convert/sparkml/convert.py | 75 +- .../aft_survival_regression.py | 41 +- .../sparkml/operator_converters/binarizer.py | 18 +- .../bucketed_random_projection_lsh.py | 53 +- .../sparkml/operator_converters/bucketizer.py | 112 ++- .../operator_converters/chi_sq_selector.py | 31 +- .../sparkml/operator_converters/common.py | 54 +- .../operator_converters/count_vectorizer.py | 37 +- .../sparkml/operator_converters/dct.py | 45 +- .../decision_tree_classifier.py | 52 +- .../decision_tree_regressor.py | 43 +- .../element_wise_product.py | 30 +- .../operator_converters/gbt_classifier.py | 125 ++- .../sparkml/operator_converters/imputer.py | 86 +- .../operator_converters/index_to_string.py | 36 +- .../sparkml/operator_converters/k_means.py | 131 +-- .../operator_converters/linear_classifier.py | 88 +- .../operator_converters/linear_regressor.py | 39 +- .../operator_converters/min_hash_lsh.py | 72 +- .../operator_converters/mlp_classifier.py | 31 +- .../operator_converters/naive_bayes.py | 242 ++++-- .../sparkml/operator_converters/normalizer.py | 29 +- .../operator_converters/one_vs_rest.py | 58 +- .../operator_converters/onehot_encoder.py | 49 +- .../sparkml/operator_converters/pca.py | 27 +- .../polynomial_expansion.py | 115 ++- .../random_forest_classifier.py | 52 +- .../random_forest_regressor.py | 44 +- .../sparkml/operator_converters/scaler.py | 50 +- .../operator_converters/stop_words_remover.py | 30 +- .../operator_converters/string_indexer.py | 22 +- .../sparkml/operator_converters/tokenizer.py | 49 +- .../tree_ensemble_common.py | 175 ++-- .../operator_converters/tree_helper.py | 49 +- .../operator_converters/vector_assembler.py | 22 +- .../operator_converters/vector_indexer.py | 89 +- .../operator_converters/vector_slicer.py | 31 +- .../sparkml/operator_converters/word2vec.py | 81 +- .../convert/sparkml/ops_input_output.py | 25 +- onnxmltools/convert/sparkml/ops_names.py | 112 ++- onnxmltools/convert/sparkml/utils.py | 39 +- onnxmltools/convert/xgboost/__init__.py | 1 + onnxmltools/convert/xgboost/_parse.py | 136 +-- onnxmltools/convert/xgboost/common.py | 8 +- onnxmltools/convert/xgboost/convert.py | 58 +- .../xgboost/operator_converters/XGBoost.py | 339 +++++--- .../xgboost/shape_calculators/Classifier.py | 22 +- .../xgboost/shape_calculators/Regressor.py | 2 +- onnxmltools/proto/__init__.py | 28 +- onnxmltools/utils/__init__.py | 6 +- onnxmltools/utils/main.py | 8 +- onnxmltools/utils/tests_dl_helper.py | 37 +- onnxmltools/utils/tests_helper.py | 204 +++-- onnxmltools/utils/utils_backend.py | 91 +- .../utils/utils_backend_onnxruntime.py | 200 +++-- onnxmltools/utils/visualize.py | 55 +- pyproject.toml | 24 + requirements-dev.txt | 5 +- setup.py | 52 +- tests/baseline/test_convert_baseline.py | 35 +- tests/catboost/test_CatBoost_converter.py | 93 ++- .../test_cml_AllNeuralNetworkConverters.py | 662 +++++++++++---- .../test_cml_DictVectorizerConverter.py | 28 +- .../coreml/test_cml_GLMClassifierConverter.py | 25 +- .../coreml/test_cml_GLMRegressorConverter.py | 19 +- tests/coreml/test_cml_ImputerConverter.py | 26 +- .../coreml/test_cml_OneHotEncoderConverter.py | 37 +- tests/coreml/test_cml_ScalerConverter.py | 19 +- ...st_cml_SupportVectorClassifierConverter.py | 39 +- ...est_cml_SupportVectorRegressorConverter.py | 17 +- ...est_cml_TreeEnsembleClassifierConverter.py | 19 +- ...test_cml_TreeEnsembleRegressorConverter.py | 19 +- ...l_TreeEnsembleRegressorConverterXGBoost.py | 20 +- tests/h2o/test_h2o_converters.py | 82 +- ...htGbmTreeEnsembleConverters_hummingbird.py | 246 ++++-- .../test_LightGbmTreeEnsembleConverters.py | 217 +++-- .../test_LightGbmTreeEnsembleConvertersPkl.py | 52 +- ...st_LightGbmTreeEnsembleConverters_split.py | 68 +- .../lightgbm/test_lightgbm_missing_values.py | 42 +- .../lightgbm/test_lightgbm_tree_structure.py | 422 +++++----- tests/lightgbm/test_objective_functions.py | 68 +- tests/sparkml/__init__.py | 14 +- tests/sparkml/profile_pipeline.py | 83 +- tests/sparkml/r_pipeline.py | 44 +- tests/sparkml/sparkml_test_base.py | 15 +- tests/sparkml/sparkml_test_utils.py | 134 +-- tests/sparkml/test_PCA.py | 50 +- tests/sparkml/test_aft_survival_regression.py | 45 +- tests/sparkml/test_binarizer.py | 32 +- .../test_bucketed_random_projection_lsh.py | 77 +- tests/sparkml/test_bucketizer.py | 35 +- tests/sparkml/test_chi_sql_selector.py | 51 +- tests/sparkml/test_count_vectorizer.py | 113 ++- tests/sparkml/test_dct.py | 42 +- .../sparkml/test_decision_tree_classifier.py | 217 +++-- .../test_decision_tree_classifier_category.py | 85 +- tests/sparkml/test_decision_tree_regressor.py | 118 ++- tests/sparkml/test_decision_tree_rules.py | 353 ++++---- tests/sparkml/test_element_wise_product.py | 48 +- tests/sparkml/test_gbt_classifier.py | 54 +- tests/sparkml/test_gbt_regressor.py | 44 +- tests/sparkml/test_imputer.py | 99 ++- tests/sparkml/test_index_to_string.py | 56 +- tests/sparkml/test_k_means.py | 120 ++- tests/sparkml/test_linear_classifier.py | 93 ++- tests/sparkml/test_linear_regressor.py | 121 ++- tests/sparkml/test_min_hash_lsh.py | 64 +- tests/sparkml/test_mlp_classifier.py | 51 +- tests/sparkml/test_naive_bayes.py | 101 ++- tests/sparkml/test_normalizer.py | 95 ++- tests/sparkml/test_one_vs_rest.py | 46 +- tests/sparkml/test_onehot_encoder.py | 199 ++++- tests/sparkml/test_pipeline.py | 256 ++++-- tests/sparkml/test_polynomial_expansion.py | 49 +- .../sparkml/test_random_forest_classifier.py | 85 +- .../test_random_forest_classifier_tree.py | 47 +- tests/sparkml/test_random_forest_regressor.py | 84 +- tests/sparkml/test_scaler.py | 174 +++- tests/sparkml/test_stop_words_remover.py | 29 +- tests/sparkml/test_string_indexer.py | 72 +- tests/sparkml/test_tokenizer.py | 31 +- tests/sparkml/test_tree_helper.py | 671 ++++++++++++--- tests/sparkml/test_vector_assembler.py | 47 +- tests/sparkml/test_vector_indexer.py | 112 ++- tests/sparkml/test_vector_slicer.py | 49 +- tests/sparkml/test_word2vec.py | 61 +- tests/svmlib/test_SVMConverters.py | 263 ++++-- tests/utils/test_utils.py | 40 +- tests/xgboost/test_xgboost_13.py | 40 +- tests/xgboost/test_xgboost_converters.py | 606 +++++++++----- tests/xgboost/test_xgboost_pipeline.py | 99 ++- tests/xgboost/test_xgboost_unpickle_06.py | 18 +- 261 files changed, 13067 insertions(+), 6477 deletions(-) delete mode 100644 .azure-pipelines/linux-CI-nightly.yml delete mode 100644 .azure-pipelines/win32-CI-nightly.yml delete mode 100644 .flake8 create mode 100644 .github/workflows/black-ruff.yml create mode 100644 pyproject.toml diff --git a/.azure-pipelines/linux-CI-nightly.yml b/.azure-pipelines/linux-CI-nightly.yml deleted file mode 100644 index 0135a2fd..00000000 --- a/.azure-pipelines/linux-CI-nightly.yml +++ /dev/null @@ -1,66 +0,0 @@ -# Python package -# Create and test a Python package on multiple Python versions. -# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more: -# https://docs.microsoft.com/azure/devops/pipelines/languages/python - -trigger: -- master - -jobs: - -- job: 'Test' - pool: - vmImage: 'Ubuntu-16.04' - strategy: - matrix: - Python39-nightly: - python.version: '3.9' - ORT_PATH: -i https://test.pypi.org/simple/ ort-nightly - COREML_PATH: git+https://github.com/apple/coremltools@3.1 - Python38-nightly: - python.version: '3.8' - ORT_PATH: -i https://test.pypi.org/simple/ ort-nightly - COREML_PATH: git+https://github.com/apple/coremltools@3.1 - Python37-nightly: - python.version: '3.7' - ORT_PATH: -i https://test.pypi.org/simple/ ort-nightly - COREML_PATH: git+https://github.com/apple/coremltools@3.1 - maxParallel: 3 - - steps: - - script: sudo install -d -m 0777 /home/vsts/.conda/envs - displayName: Fix Conda permissions - - - task: CondaEnvironment@1 - inputs: - createCustomEnvironment: true - environmentName: 'py$(python.version)' - packageSpecs: 'python=$(python.version)' - - - script: | - python -m pip install --upgrade pip - conda config --set always_yes yes --set changeps1 no - conda install -c conda-forge protobuf - conda install -c conda-forge numpy - conda install -c conda-forge cmake - python -m pip install $(COREML_PATH) - python -m pip install $(ONNX_PATH) - python -m pip install hummingbird-ml --no-deps - python -m pip install -r requirements.txt - python -m pip install -r requirements-dev.txt - python -m pip install $(ORT_PATH) - python -m pip install pytest - displayName: 'Install dependencies' - - - script: | - pip install -e . - python -c "import onnxconverter_common;print(onnxconverter_common.__version__)" - python -c "import onnxruntime;print(onnxruntime.__version__)" - pytest tests --ignore=tests/sparkml --doctest-modules --junitxml=junit/test-results.xml - displayName: 'pytest - onnxmltools' - - - task: PublishTestResults@2 - inputs: - testResultsFiles: '**/test-results.xml' - testRunTitle: 'Python $(python.version)' - condition: succeededOrFailed() diff --git a/.azure-pipelines/linux-conda-CI.yml b/.azure-pipelines/linux-conda-CI.yml index 90cb056f..a08470cc 100644 --- a/.azure-pipelines/linux-conda-CI.yml +++ b/.azure-pipelines/linux-conda-CI.yml @@ -89,11 +89,6 @@ jobs: python -m pip install $(ONNXRT_PATH) displayName: 'Install onnxruntime' - - script: | - pip install flake8 - python -m flake8 ./onnxmltools - displayName: 'run flake8 check' - - script: | pip install -e . displayName: 'local installation' diff --git a/.azure-pipelines/win32-CI-nightly.yml b/.azure-pipelines/win32-CI-nightly.yml deleted file mode 100644 index 3aad5d61..00000000 --- a/.azure-pipelines/win32-CI-nightly.yml +++ /dev/null @@ -1,64 +0,0 @@ -# Python package -# Create and test a Python package on multiple Python versions. -# Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more: -# https://docs.microsoft.com/azure/devops/pipelines/languages/python - -trigger: -- master - -jobs: - -- job: 'Test' - pool: - vmImage: 'windows-latest' - strategy: - matrix: - Python39-nightly: - python.version: '3.9' - ONNXRT_PATH: -i https://test.pypi.org/simple/ ort-nightly - COREML_PATH: git+https://github.com/apple/coremltools@3.1 - Python38-nightly: - python.version: '3.8' - ONNXRT_PATH: -i https://test.pypi.org/simple/ ort-nightly - COREML_PATH: git+https://github.com/apple/coremltools@3.1 - Python37-nightly: - python.version: '3.7' - ONNXRT_PATH: -i https://test.pypi.org/simple/ ort-nightly - COREML_PATH: git+https://github.com/apple/coremltools@3.1 - maxParallel: 3 - - steps: - - task: UsePythonVersion@0 - inputs: - versionSpec: '$(python.version)' - architecture: 'x64' - - - powershell: Write-Host "##vso[task.prependpath]$env:CONDA\Scripts" - displayName: Add conda to PATH - - - script: conda create --yes --quiet --name py$(python.version) -c conda-forge python=$(python.version) numpy protobuf - displayName: Create Anaconda environment - - - script: | - call activate py$(python.version) - python -m pip install --upgrade pip numpy - pip install %COREML_PATH% %ONNX_PATH% - pip install humming-bird-ml --no-deps - pip install -r requirements.txt - pip install -r requirements-dev.txt - pip install %ONNXRT_PATH% - displayName: 'Install dependencies' - - - script: | - call activate py$(python.version) - pip install -e . - python -c "import onnxconverter_common;print(onnxconverter_common.__version__)" - python -c "import onnxruntime;print(onnxruntime.__version__)" - python -m pytest tests --ignore=tests/sparkml --doctest-modules --junitxml=junit/test-results.xml - displayName: 'pytest - onnxmltools' - - - task: PublishTestResults@2 - inputs: - testResultsFiles: '**/test-results.xml' - testRunTitle: 'Python $(python.version)' - condition: succeededOrFailed() diff --git a/.azure-pipelines/win32-conda-CI.yml b/.azure-pipelines/win32-conda-CI.yml index 68352693..ebb3108e 100644 --- a/.azure-pipelines/win32-conda-CI.yml +++ b/.azure-pipelines/win32-conda-CI.yml @@ -75,11 +75,6 @@ jobs: python -m pip install $(ONNXRT_PATH) displayName: 'Install ort-nightly' - - script: | - call activate py$(python.version) - python -m flake8 ./onnxmltools - displayName: 'run flake8 check' - - script: | call activate py$(python.version) python -m pip install -e . diff --git a/.flake8 b/.flake8 deleted file mode 100644 index 98f34b13..00000000 --- a/.flake8 +++ /dev/null @@ -1,10 +0,0 @@ -[flake8] -max-line-length = 120 -per-file-ignores = - __init__.py:F401 -exclude = - sparkml - coreml - libsvm - xgboost - utils diff --git a/.github/workflows/black-ruff.yml b/.github/workflows/black-ruff.yml new file mode 100644 index 00000000..09da3fc3 --- /dev/null +++ b/.github/workflows/black-ruff.yml @@ -0,0 +1,16 @@ +name: Black Format Checker +on: [push, pull_request] +jobs: + black-format-check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - uses: psf/black@193ee766ca496871f93621d6b58d57a6564ff81b # stable 23.7.0 + with: + options: "--diff --check" + src: "." + ruff-format-check: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: chartboost/ruff-action@v1.0.0 diff --git a/docs/conf.py b/docs/conf.py index d2384094..bc955198 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -4,49 +4,42 @@ # Configuration file for the Sphinx documentation builder. -import os -import sys -import shutil import onnxmltools import sphinx_readable_theme -import tabulate -import sphinx_gallery.gen_gallery - - # -- Project information ----------------------------------------------------- -project = 'onnxmltools' -copyright = '2018-2020, Microsoft' -author = 'Microsoft' +project = "onnxmltools" +copyright = "2018-2020, Microsoft" +author = "Microsoft" version = onnxmltools.__version__ release = version # -- General configuration --------------------------------------------------- extensions = [ - 'sphinx.ext.intersphinx', - 'sphinx.ext.imgmath', - 'sphinx.ext.ifconfig', - 'sphinx.ext.viewcode', + "sphinx.ext.intersphinx", + "sphinx.ext.imgmath", + "sphinx.ext.ifconfig", + "sphinx.ext.viewcode", "sphinx.ext.autodoc", - 'sphinx.ext.githubpages', + "sphinx.ext.githubpages", "sphinx_gallery.gen_gallery", - 'sphinx.ext.autodoc', + "sphinx.ext.autodoc", ] -templates_path = ['_templates'] -source_suffix = ['.rst'] +templates_path = ["_templates"] +source_suffix = [".rst"] -master_doc = 'index' +master_doc = "index" language = "en" exclude_patterns = [] -pygments_style = 'default' +pygments_style = "default" # -- Options for HTML output ------------------------------------------------- -html_static_path = ['_static'] +html_static_path = ["_static"] html_theme = "readable" html_theme_path = [sphinx_readable_theme.get_html_theme_path()] html_logo = "ONNXMLTools_logo_main.png" @@ -58,19 +51,19 @@ # -- Options for intersphinx extension --------------------------------------- # Example configuration for intersphinx: refer to the Python standard library. -intersphinx_mapping = {'https://docs.python.org/': None} +intersphinx_mapping = {"https://docs.python.org/": None} # -- Options for Sphinx Gallery ---------------------------------------------- sphinx_gallery_conf = { - 'examples_dirs': 'examples', - 'gallery_dirs': 'auto_examples', + "examples_dirs": "examples", + "gallery_dirs": "auto_examples", } # -- Setup actions ----------------------------------------------------------- + def setup(app): # Placeholder to initialize the folder before # generating the documentation. return app - diff --git a/docs/examples/plot_convert_h2o.py b/docs/examples/plot_convert_h2o.py index 2078baee..556b3492 100644 --- a/docs/examples/plot_convert_h2o.py +++ b/docs/examples/plot_convert_h2o.py @@ -21,18 +21,19 @@ import os import numpy import onnx +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer +import matplotlib.pyplot as plt import sklearn -from sklearn.linear_model import LogisticRegression from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import onnxruntime as rt import h2o from h2o.estimators.gbm import H2OGradientBoostingEstimator -import skl2onnx -import onnxmltools from onnxconverter_common.data_types import FloatTensorType +import onnxmltools from onnxmltools.convert import convert_h2o + iris = load_iris() X, y = iris.data, iris.target X_train, X_test, y_train, y_test = train_test_split(X, y) @@ -56,7 +57,7 @@ # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_h2o(pth, initial_types=initial_type) h2o.cluster().shutdown() @@ -68,8 +69,7 @@ sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) ################################## @@ -77,22 +77,23 @@ # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/docs/examples/plot_convert_keras.py b/docs/examples/plot_convert_keras.py index accb1ccd..a6fd4ded 100644 --- a/docs/examples/plot_convert_keras.py +++ b/docs/examples/plot_convert_keras.py @@ -19,10 +19,12 @@ """ +import os +import matplotlib.pyplot as plt +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer import numpy import onnx import sklearn -from sklearn.linear_model import LogisticRegression from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import keras @@ -30,11 +32,11 @@ from keras.layers import Dense import onnxruntime as rt -import skl2onnx import onnxmltools from onnxconverter_common.data_types import FloatTensorType from onnxmltools.convert import convert_keras + iris = load_iris() X, y = iris.data, iris.target y_multi = numpy.zeros((y.shape[0], 3), dtype=numpy.int64) @@ -42,11 +44,9 @@ X_train, X_test, y_train, y_test = train_test_split(X, y_multi) model = Sequential() -model.add(Dense(units=10, activation='relu', input_dim=4)) -model.add(Dense(units=3, activation='softmax')) -model.compile(loss='categorical_crossentropy', - optimizer='sgd', - metrics=['accuracy']) +model.add(Dense(units=10, activation="relu", input_dim=4)) +model.add(Dense(units=3, activation="softmax")) +model.compile(loss="categorical_crossentropy", optimizer="sgd", metrics=["accuracy"]) model.fit(X_train, y_train, epochs=5, batch_size=16) print("keras prediction") print(model.predict(X_test.astype(numpy.float32))) @@ -55,7 +55,7 @@ # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_keras(model, initial_types=initial_type) ################################### @@ -65,8 +65,7 @@ sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name output_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [output_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([output_name], {input_name: X_test.astype(numpy.float32)})[0] print("ONNX prediction") print(pred_onx) @@ -75,22 +74,23 @@ # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/docs/examples/plot_convert_libsvm.py b/docs/examples/plot_convert_libsvm.py index afe9ebf5..cb6c25dc 100644 --- a/docs/examples/plot_convert_libsvm.py +++ b/docs/examples/plot_convert_libsvm.py @@ -18,18 +18,18 @@ +++++++++++++ """ - +import os +import matplotlib.pyplot as plt +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer import numpy import onnx import sklearn -from sklearn.linear_model import LogisticRegression from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split from svm import C_SVC as SVC import svmutil import onnxruntime as rt -import skl2onnx import onnxmltools from onnxconverter_common.data_types import FloatTensorType from onnxmltools.convert import convert_libsvm @@ -54,7 +54,7 @@ # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_libsvm(clr, initial_types=initial_type) ################################### @@ -64,8 +64,7 @@ sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) ################################## @@ -73,22 +72,23 @@ # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/docs/examples/plot_convert_lightgbm.py b/docs/examples/plot_convert_lightgbm.py index 8c68d8b3..3f9ce102 100644 --- a/docs/examples/plot_convert_lightgbm.py +++ b/docs/examples/plot_convert_lightgbm.py @@ -18,17 +18,17 @@ +++++++++++++ """ - +import os +import matplotlib.pyplot as plt +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer import numpy import onnx import sklearn -from sklearn.linear_model import LogisticRegression from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import lightgbm from lightgbm import LGBMClassifier, Dataset, train as train_lgbm import onnxruntime as rt -import skl2onnx import onnxmltools from onnxconverter_common.data_types import FloatTensorType from onnxmltools.convert import convert_lightgbm @@ -44,7 +44,7 @@ # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_lightgbm(clr, initial_types=initial_type) ################################### @@ -54,8 +54,7 @@ sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) ############################################### @@ -68,17 +67,16 @@ dtrain = Dataset(X_train, label=y_train) -param = {'objective': 'multiclass', 'num_class': 3} +param = {"objective": "multiclass", "num_class": 3} bst = train_lgbm(param, dtrain, 10) -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_lightgbm(bst, initial_types=initial_type) sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) @@ -87,22 +85,23 @@ # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/docs/examples/plot_convert_sklearn.py b/docs/examples/plot_convert_sklearn.py index 676478ae..847e3dde 100644 --- a/docs/examples/plot_convert_sklearn.py +++ b/docs/examples/plot_convert_sklearn.py @@ -30,7 +30,9 @@ A very basic example using random forest and the iris dataset. """ - +import os +import matplotlib.pyplot as plt +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer import numpy import onnx import sklearn @@ -56,7 +58,7 @@ # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_sklearn(clr, initial_types=initial_type) with open("rf_iris.onnx", "wb") as f: @@ -68,8 +70,7 @@ sess = rt.InferenceSession("rf_iris.onnx") input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) ####################################### @@ -77,7 +78,7 @@ clr = LogisticRegression() clr.fit(X_train, y_train) -initial_type = [('float_input', FloatTensorType([None, X_train.shape[1]]))] +initial_type = [("float_input", FloatTensorType([None, X_train.shape[1]]))] onx = convert_sklearn(clr, initial_types=initial_type) with open("logreg_iris.onnx", "wb") as f: f.write(onx.SerializeToString()) @@ -85,8 +86,7 @@ sess = rt.InferenceSession("logreg_iris.onnx") input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run([label_name], - {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) @@ -95,22 +95,23 @@ # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/docs/examples/plot_convert_sparkml.py b/docs/examples/plot_convert_sparkml.py index 4a5d96f6..a573b2a1 100644 --- a/docs/examples/plot_convert_sparkml.py +++ b/docs/examples/plot_convert_sparkml.py @@ -18,8 +18,11 @@ +++++++++++++ """ + import os import numpy +import matplotlib.pyplot as plt +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer from pandas import DataFrame import onnx import sklearn @@ -27,12 +30,8 @@ from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import onnxruntime as rt -import skl2onnx import pyspark from pyspark.sql import SparkSession -from pyspark.ml.classification import LogisticRegression, LinearSVC -from pyspark.ml.linalg import VectorUDT, SparseVector -from pyspark.ml.classification import LogisticRegression from pyspark.ml.feature import VectorAssembler, StringIndexer import onnxmltools from onnxconverter_common.data_types import FloatTensorType @@ -43,6 +42,7 @@ def start_spark(options=None): import os import sys import pyspark + executable = sys.executable os.environ["SPARK_HOME"] = pyspark.__path__[0] os.environ["PYSPARK_PYTHON"] = executable @@ -50,7 +50,7 @@ def start_spark(options=None): builder = SparkSession.builder.appName("pyspark-unittesting").master("local[1]") if options: - for k,v in options.items(): + for k, v in options.items(): builder.config(k, v) spark = builder.getOrCreate() @@ -60,17 +60,18 @@ def start_spark(options=None): def stop_spark(spark): spark.sparkContext.stop() + iris = load_iris() X, y = iris.data, iris.target X_train, X_test, y_train, y_test = train_test_split(X, y) df = DataFrame(X_train, columns="x1 x2 x3 x4".split()) -df['class'] = y_train +df["class"] = y_train # df.to_csv("data_train.csv", index=False, header=False) -this_script_dir = os.path.abspath('.') -if os.name == 'nt' and os.environ.get('HADOOP_HOME') is None: - print('setting HADOOP_HOME to: ', this_script_dir) - os.environ['HADOOP_HOME'] = this_script_dir +this_script_dir = os.path.abspath(".") +if os.name == "nt" and os.environ.get("HADOOP_HOME") is None: + print("setting HADOOP_HOME to: ", this_script_dir) + os.environ["HADOOP_HOME"] = this_script_dir spark_session = start_spark() @@ -78,12 +79,12 @@ def stop_spark(spark): data = spark_session.createDataFrame(df) feature_cols = data.columns[:-1] -assembler = VectorAssembler(inputCols=feature_cols, outputCol='features') +assembler = VectorAssembler(inputCols=feature_cols, outputCol="features") train_data = assembler.transform(data) -train_data = train_data.select(['features', 'class']) -label_indexer = StringIndexer(inputCol='class', outputCol='label').fit(train_data) +train_data = train_data.select(["features", "class"]) +label_indexer = StringIndexer(inputCol="class", outputCol="label").fit(train_data) train_data = label_indexer.transform(train_data) -train_data = train_data.select(['features', 'label']) +train_data = train_data.select(["features", "label"]) train_data.show(10) lr = LogisticRegression(maxIter=100, tol=0.0001) @@ -94,8 +95,8 @@ def stop_spark(spark): # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_types = [('features', FloatTensorType([None, 4]))] -onx = convert_sparkml(model, 'sparkml logistic regression', initial_types) +initial_types = [("features", FloatTensorType([None, 4]))] +onx = convert_sparkml(model, "sparkml logistic regression", initial_types) stop_spark(spark_session) @@ -106,8 +107,7 @@ def stop_spark(spark): sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) ################################## @@ -115,22 +115,23 @@ def stop_spark(spark): # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/docs/examples/plot_convert_xgboost.py b/docs/examples/plot_convert_xgboost.py index 96838b1c..444b7ff3 100644 --- a/docs/examples/plot_convert_xgboost.py +++ b/docs/examples/plot_convert_xgboost.py @@ -18,18 +18,18 @@ +++++++++++++ """ - +import os import numpy +import matplotlib.pyplot as plt +from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer import onnx import sklearn -from sklearn.linear_model import LogisticRegression from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split import xgboost from xgboost import XGBClassifier, DMatrix, train as train_xgb import onnxruntime as rt -import skl2onnx import onnxmltools from onnxconverter_common.data_types import FloatTensorType from onnxmltools.convert import convert_xgboost @@ -45,7 +45,7 @@ # Convert a model into ONNX # +++++++++++++++++++++++++ -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_xgboost(clr, initial_types=initial_type) ################################### @@ -55,8 +55,7 @@ sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) ############################################### @@ -69,17 +68,16 @@ dtrain = DMatrix(X_train, label=y_train) -param = {'objective': 'multi:softmax', 'num_class': 3} +param = {"objective": "multi:softmax", "num_class": 3} bst = train_xgb(param, dtrain, 10) -initial_type = [('float_input', FloatTensorType([None, 4]))] +initial_type = [("float_input", FloatTensorType([None, 4]))] onx = convert_xgboost(bst, initial_types=initial_type) sess = rt.InferenceSession(onx.SerializeToString()) input_name = sess.get_inputs()[0].name label_name = sess.get_outputs()[0].name -pred_onx = sess.run( - [label_name], {input_name: X_test.astype(numpy.float32)})[0] +pred_onx = sess.run([label_name], {input_name: X_test.astype(numpy.float32)})[0] print(pred_onx) @@ -88,22 +86,23 @@ # ++++++++++++++++++++++ # # Finally, let's see the graph converted with *onnxmltools*. -import os -import matplotlib.pyplot as plt -from onnx.tools.net_drawer import GetPydotGraph, GetOpNodeProducer pydot_graph = GetPydotGraph( - onx.graph, name=onx.graph.name, rankdir="TB", + onx.graph, + name=onx.graph.name, + rankdir="TB", node_producer=GetOpNodeProducer( - "docstring", color="yellow", fillcolor="yellow", style="filled")) + "docstring", color="yellow", fillcolor="yellow", style="filled" + ), +) pydot_graph.write_dot("model.dot") -os.system('dot -O -Gdpi=300 -Tpng model.dot') +os.system("dot -O -Gdpi=300 -Tpng model.dot") image = plt.imread("model.dot.png") fig, ax = plt.subplots(figsize=(40, 20)) ax.imshow(image) -ax.axis('off') +ax.axis("off") ################################# diff --git a/onnxmltools/convert/common/_container.py b/onnxmltools/convert/common/_container.py index bd76e611..5d098083 100644 --- a/onnxmltools/convert/common/_container.py +++ b/onnxmltools/convert/common/_container.py @@ -2,7 +2,7 @@ from onnxconverter_common.container import ( RawModelContainer, - CommonSklearnModelContainer + CommonSklearnModelContainer, ) @@ -19,10 +19,10 @@ class H2OModelContainer(CommonSklearnModelContainer): class SparkmlModelContainer(RawModelContainer): - def __init__(self, sparkml_model): super(SparkmlModelContainer, self).__init__(sparkml_model) - # Sparkml models have no input and output specified, so we create them and store them in this container. + # Sparkml models have no input and output specified, + # so we create them and store them in this container. self._inputs = [] self._outputs = [] @@ -35,18 +35,19 @@ def output_names(self): return [variable.raw_name for variable in self._outputs] def add_input(self, variable): - # The order of adding variables matters. The final model's input names are sequentially added as this list + # The order of adding variables matters. The final model's + # input names are sequentially added as this list if variable not in self._inputs: self._inputs.append(variable) def add_output(self, variable): - # The order of adding variables matters. The final model's output names are sequentially added as this list + # The order of adding variables matters. + # The final model's output names are sequentially added as this list if variable not in self._outputs: self._outputs.append(variable) class CoremlModelContainer(RawModelContainer): - def __init__(self, coreml_model): super(CoremlModelContainer, self).__init__(coreml_model) diff --git a/onnxmltools/convert/common/interface.py b/onnxmltools/convert/common/interface.py index 35195d51..af0721c3 100644 --- a/onnxmltools/convert/common/interface.py +++ b/onnxmltools/convert/common/interface.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # This file defines the interface of the converter internal object for callback, -# So the usage of the methods and properties list here will not be affected among the different versions. +# So the usage of the methods and properties list here +# will not be affected among the different versions. from onnxconverter_common.interface import * # noqa diff --git a/onnxmltools/convert/common/onnx_ex.py b/onnxmltools/convert/common/onnx_ex.py index 3a9fd043..289a440e 100644 --- a/onnxmltools/convert/common/onnx_ex.py +++ b/onnxmltools/convert/common/onnx_ex.py @@ -2,9 +2,9 @@ from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import ( - get_maximum_opset_supported as _get_maximum_opset_supported) + get_maximum_opset_supported as _get_maximum_opset_supported, +) def get_maximum_opset_supported(): - return min(onnx_opset_version(), - max(15, _get_maximum_opset_supported())) + return min(onnx_opset_version(), max(15, _get_maximum_opset_supported())) diff --git a/onnxmltools/convert/common/tree_ensemble.py b/onnxmltools/convert/common/tree_ensemble.py index 488f71e7..32b44b37 100644 --- a/onnxmltools/convert/common/tree_ensemble.py +++ b/onnxmltools/convert/common/tree_ensemble.py @@ -22,6 +22,8 @@ def _process_process_tree_attributes(attrs): continue wrong_types.append(f"Unexpected type {type(v)} for attribute {k!r}.") if len(wrong_types) > 0: - raise TypeError("Unexpected type for one or several attributes:\n" + "\n".join(wrong_types)) + raise TypeError( + "Unexpected type for one or several attributes:\n" + "\n".join(wrong_types) + ) if update: attrs.update(update) diff --git a/onnxmltools/convert/common/utils.py b/onnxmltools/convert/common/utils.py index 963c0ebc..5b59e5ec 100644 --- a/onnxmltools/convert/common/utils.py +++ b/onnxmltools/convert/common/utils.py @@ -3,6 +3,7 @@ try: from onnxconverter_common.utils import hummingbird_installed # noqa except ImportError: + def hummingbird_installed(): """ Checks that *Hummingbird* is available. @@ -21,6 +22,7 @@ def tf2onnx_installed(): """ try: import tf2onnx # noqa F401 + return True except ImportError: return False diff --git a/onnxmltools/convert/coreml/_parse.py b/onnxmltools/convert/coreml/_parse.py index 0c0f76ce..778d91da 100644 --- a/onnxmltools/convert/coreml/_parse.py +++ b/onnxmltools/convert/coreml/_parse.py @@ -1,64 +1,103 @@ # SPDX-License-Identifier: Apache-2.0 import warnings -import onnx from ..common._container import CoremlModelContainer from ..common._topology import Topology -from ..common.data_types import * +from ..common.data_types import ( + find_type_conversion, + FloatTensorType, + Int64TensorType, + StringTensorType, + DictionaryType, + Int64Type, + FloatType, + StringType, +) def _parse_coreml_feature(feature_info, target_opset, batch_size=1): - ''' - Encode type information from CoreML's FeatureType protobuf message in converter's type system. - - Scalar types such as Int64FeatureType, DoubleFeatureType, and StringFeatureType in CoreML are interpreted as - [batch_size, 1]-tensor. Tensor-like types such as ArrayFeature in CoreML is viewed as tensors with a prepend - batch_size; for example, we use [batch_size, C, H, W] to denote [C, H, W]-array in CoreML. - :param feature_info: CoreML FeatureDescription (https://apple.github.io/coremltools/coremlspecification/sections/DataStructuresAndFeatureTypes.html#featuretype) + """ + Encode type information from CoreML's FeatureType protobuf message in + converter's type system. + + Scalar types such as Int64FeatureType, DoubleFeatureType, + and StringFeatureType in CoreML are interpreted as + [batch_size, 1]-tensor. Tensor-like types such as + ArrayFeature in CoreML is viewed as tensors with a prepend + batch_size; for example, we use [batch_size, C, H, W] + to denote [C, H, W]-array in CoreML. + + :param feature_info: CoreML FeatureDescription + (https://apple.github.io/coremltools/coremlspecification/ + sections/DataStructuresAndFeatureTypes.html#featuretype) :param target_opset: the target ospet number in the converted model. - :param batch_size: default batch size prepend to scalars and tensors variables from CoreML - :return: one of our Int64Type, FloatType, StringType, Int64TensorType, FloatTensorType, or DictionaryType - ''' + :param batch_size: default batch size prepend to scalars + and tensors variables from CoreML + :return: one of our Int64Type, FloatType, StringType, + Int64TensorType, FloatTensorType, or DictionaryType + """ raw_type = feature_info.type doc_string = feature_info.shortDescription - type_name = raw_type.WhichOneof('Type') + type_name = raw_type.WhichOneof("Type") - if type_name == 'int64Type': + if type_name == "int64Type": return Int64Type(doc_string=doc_string) - elif type_name == 'doubleType': + elif type_name == "doubleType": return FloatType(doc_string=doc_string) - elif type_name == 'stringType': + elif type_name == "stringType": return StringType(doc_string=doc_string) - elif type_name == 'imageType': - # Produce [C, H, W]-tensor, where C is the number of color channels, H the height, and W the width. + elif type_name == "imageType": + # Produce [C, H, W]-tensor, where C is the number of color channels, + # H the height, and W the width. color_space = raw_type.imageType.colorSpace shape = [batch_size] if doc_string: - if doc_string[-1] not in ['.', '!', '?']: - doc_string += '. ' + if doc_string[-1] not in [".", "!", "?"]: + doc_string += ". " else: - doc_string += ' ' + doc_string += " " if color_space == 10: # gray scale shape.append(1) - doc_string += 'Image(s) in gray scale. If there are N images, it is a 4-D tensor with shape [N, 1, H, W]' + doc_string += ( + "Image(s) in gray scale. If there are N images, " + "it is a 4-D tensor with shape [N, 1, H, W]" + ) elif color_space == 20: # RGB (20) shape.append(3) - doc_string += 'Image(s) in RGB format. It is a [N, C, H, W]-tensor. The 1st/2nd/3rd slices along the ' \ - 'C-axis are red, green, and blue channels, respectively.' + doc_string += ( + "Image(s) in RGB format. It is a [N, C, H, W]-tensor. " + "The 1st/2nd/3rd slices along the " + "C-axis are red, green, and blue channels, respectively." + ) elif color_space == 30: # BGR (30) shape.append(3) - doc_string += 'Image(s) in BGR format. It is a [N, C, H, W]-tensor. The 1st/2nd/3rd slices along the ' \ - 'C-axis are blue, green, and red channels, respectively.' + doc_string += ( + "Image(s) in BGR format. It is a [N, C, H, W]-tensor. " + "The 1st/2nd/3rd slices along the " + "C-axis are blue, green, and red channels, respectively." + ) else: - raise ValueError('Unknown image format. Only gray-level, RGB, and BGR are supported') + raise ValueError( + "Unknown image format. Only gray-level, RGB, and BGR are supported" + ) shape.append(raw_type.imageType.height) shape.append(raw_type.imageType.width) - color_space_map = {10: 'Gray8', 20: 'Rgb8', 30: 'Bgr8'} - return FloatTensorType(shape, color_space_map[color_space], doc_string=doc_string, - denotation='IMAGE', channel_denotations=['DATA_BATCH', 'DATA_CHANNEL', 'DATA_FEATURE', 'DATA_FEATURE']) - elif type_name == 'multiArrayType': + color_space_map = {10: "Gray8", 20: "Rgb8", 30: "Bgr8"} + return FloatTensorType( + shape, + color_space_map[color_space], + doc_string=doc_string, + denotation="IMAGE", + channel_denotations=[ + "DATA_BATCH", + "DATA_CHANNEL", + "DATA_FEATURE", + "DATA_FEATURE", + ], + ) + elif type_name == "multiArrayType": element_type_id = raw_type.multiArrayType.dataType shape = [d for d in raw_type.multiArrayType.shape] if len(shape) == 1: @@ -68,7 +107,10 @@ def _parse_coreml_feature(feature_info, target_opset, batch_size=1): # [C, H, W] shape = [batch_size, shape[0], shape[1], shape[2]] else: - shape = [batch_size, 1] # Missing shape information. We will try inferring it. + shape = [ + batch_size, + 1, + ] # Missing shape information. We will try inferring it. if element_type_id in [65568, 65600]: # CoreML FLOAT32 & DOUBLE @@ -77,81 +119,110 @@ def _parse_coreml_feature(feature_info, target_opset, batch_size=1): # CoreML INT32 return Int64TensorType(shape, doc_string=doc_string) else: - raise ValueError('Invalid element type') - elif type_name == 'dictionaryType': - key_type = raw_type.dictionaryType.WhichOneof('KeyType') - if key_type == 'int64KeyType': + raise ValueError("Invalid element type") + elif type_name == "dictionaryType": + key_type = raw_type.dictionaryType.WhichOneof("KeyType") + if key_type == "int64KeyType": if target_opset < 7: - return DictionaryType(Int64TensorType([1]), FloatTensorType([1]), doc_string=doc_string) + return DictionaryType( + Int64TensorType([1]), FloatTensorType([1]), doc_string=doc_string + ) else: - return DictionaryType(Int64TensorType([]), FloatTensorType([]), doc_string=doc_string) - elif key_type == 'stringKeyType': + return DictionaryType( + Int64TensorType([]), FloatTensorType([]), doc_string=doc_string + ) + elif key_type == "stringKeyType": if target_opset < 7: - return DictionaryType(StringTensorType([1]), FloatTensorType([1]), doc_string=doc_string) + return DictionaryType( + StringTensorType([1]), FloatTensorType([1]), doc_string=doc_string + ) else: - return DictionaryType(StringTensorType([]), FloatTensorType([]), doc_string=doc_string) + return DictionaryType( + StringTensorType([]), FloatTensorType([]), doc_string=doc_string + ) else: - raise ValueError('Unsupported key type: {}'.format(key_type)) + raise ValueError("Unsupported key type: {}".format(key_type)) else: - raise ValueError('Unsupported feature type: {}'.format(type_name)) + raise ValueError("Unsupported feature type: {}".format(type_name)) def _parse_model(topology, scope, model, inputs=None, outputs=None): - ''' - This is a delegate function of all top-level parsing functions. It does nothing but call a proper function + """ + This is a delegate function of all top-level parsing + functions. It does nothing but call a proper function to parse the given model. - ''' + """ if inputs is None: inputs = list() if outputs is None: outputs = list() - model_type = model.WhichOneof('Type') - if model_type in ['pipeline', 'pipelineClassifier', 'pipelineRegressor']: + model_type = model.WhichOneof("Type") + if model_type in ["pipeline", "pipelineClassifier", "pipelineRegressor"]: _parse_pipeline_model(topology, scope, model, inputs, outputs) - elif model_type in ['neuralNetworkClassifier', 'neuralNetworkRegressor', 'neuralNetwork']: + elif model_type in [ + "neuralNetworkClassifier", + "neuralNetworkRegressor", + "neuralNetwork", + ]: _parse_neural_network_model(topology, scope, model, inputs, outputs) else: _parse_simple_model(topology, scope, model, inputs, outputs) def _parse_simple_model(topology, parent_scope, model, inputs, outputs): - ''' + """ Parse a model containing only one operator (aka simple model). Steps: 1. Create local scope for allocating local variables and operators - 2. Create operator and then feed the model's inputs and outputs to the operator + 2. Create operator and then feed the model's + inputs and outputs to the operator 3. Connect local variables and their corresponding parent variables Note: - 1. Notice that a CoreML operator can contain no input and output, so we directly use model's inputs (outputs). - 2. Input and output names can be identical in CoreML, but they must be different for ONNX. - ''' + 1. Notice that a CoreML operator can contain no input + and output, so we directly use model's inputs (outputs). + 2. Input and output names can be identical in CoreML, + but they must be different for ONNX. + """ # Create local scope for the considered model - scope = topology.declare_scope('single', [parent_scope] + parent_scope.parent_scopes) + scope = topology.declare_scope( + "single", [parent_scope] + parent_scope.parent_scopes + ) # Create operator for the considered model - this_operator = scope.declare_local_operator(model.WhichOneof('Type'), model) + this_operator = scope.declare_local_operator(model.WhichOneof("Type"), model) - # Allocate inputs for the operator and then connect them with inputs from outside + # Allocate inputs for the operator and then connect + # them with inputs from outside for var in model.description.input: - # We assume that no duplicated raw name exists. Note that we set prepend=True because model inputs should + # We assume that no duplicated raw name exists. + # Note that we set prepend=True because model inputs should # not hide any intermediate variables. variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size), - prepend=True) + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + prepend=True, + ) this_operator.inputs.append(variable) - # Connect local variables and variables passed into this scope. Our assumptions are described below. - # 1. Assume a variable with 'A' as its CoreML name is passed in. There must be at least one local variable gets a - # raw name 'A'. That is, for each parent variable, at least one local duplicate is available. - # 2. It's possible to find multiple local variables associated with the same raw name. For example, raw name 'A' can - # be associated with 'A' and 'A1' in ONNX. In this case, we connect the first one to parent input. + # Connect local variables and variables passed into this scope. + # Our assumptions are described below. + # 1. Assume a variable with 'A' as its CoreML name is passed in. + # There must be at least one local variable gets a + # raw name 'A'. That is, for each parent variable, at + # least one local duplicate is available. + # 2. It's possible to find multiple local variables associated + # with the same raw name. For example, raw name 'A' can + # be associated with 'A' and 'A1' in ONNX. In this case, + # we connect the first one to parent input. for parent_variable in inputs: raw_name = parent_variable.raw_name child_variable = scope.variables[scope.variable_name_mapping[raw_name][0]] - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(parent_variable) operator.outputs.append(child_variable) @@ -159,160 +230,219 @@ def _parse_simple_model(topology, parent_scope, model, inputs, outputs): for var in model.description.output: # We assume that no duplicated output raw name exists. variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size)) + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + ) this_operator.outputs.append(variable) - # Connect local variables and variables passed into this scope. Our assumptions are described below. - # 1. Assume a variable with 'A' as its CoreML name is passed in. There must be at least one local variable gets a - # raw name 'A'. That is, for each parent variable, at least one local duplicate is available. - # 2. It's possible to find multiple local variables associated with the same raw name. For example, raw name 'A' can - # be associated with 'A' and 'A1' in ONNX. In this case, we connect the last one to parent output. + # Connect local variables and variables passed into this scope. + # Our assumptions are described below. + # 1. Assume a variable with 'A' as its CoreML name is passed in. + # There must be at least one local variable gets a + # raw name 'A'. That is, for each parent variable, at + # least one local duplicate is available. + # 2. It's possible to find multiple local variables associated + # with the same raw name. For example, raw name 'A' can + # be associated with 'A' and 'A1' in ONNX. In this case, + # we connect the last one to parent output. for parent_variable in outputs: raw_name = parent_variable.raw_name child_variable = scope.variables[scope.variable_name_mapping[raw_name][-1]] - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(child_variable) operator.outputs.append(parent_variable) def _parse_pipeline_model(topology, parent_scope, model, inputs, outputs): - ''' + """ Parse a pipeline including multiple sub-models. Steps: 1. Create local scope for allocating local variables and operators - 2. Sequentially parse the sub-models and create their inputs and outputs variables - 3. Connect model's (not sub-model's) inputs and outputs with proper variables created when parsing sub-models - 4. Link local variables and the corresponding parent variables (only model's inputs and outputs are considered) + 2. Sequentially parse the sub-models and create their + inputs and outputs variables + 3. Connect model's (not sub-model's) inputs and outputs + with proper variables created when parsing sub-models + 4. Link local variables and the corresponding parent + variables (only model's inputs and outputs are considered) Note: 1. A CoreML sub-model can use the same variable for its input and output. - 2. Two CoreML variables may have the same name but different types. - ''' + 2. Two CoreML variables may have the same name but + different types. + """ # Create local scope - scope = topology.declare_scope('pipeline', [parent_scope] + parent_scope.parent_scopes) + scope = topology.declare_scope( + "pipeline", [parent_scope] + parent_scope.parent_scopes + ) # Use the same name to denote sub-models - pipeline_type = model.WhichOneof('Type') - if pipeline_type == 'pipelineClassifier': + pipeline_type = model.WhichOneof("Type") + if pipeline_type == "pipelineClassifier": sub_models = model.pipelineClassifier.pipeline.models - elif pipeline_type == 'pipelineRegressor': + elif pipeline_type == "pipelineRegressor": sub_models = model.pipelineRegressor.pipeline.models - elif pipeline_type == 'pipeline': + elif pipeline_type == "pipeline": sub_models = model.pipeline.models else: - raise ValueError('Unsupported CoreML pipeline type: {0}'.format(pipeline_type)) + raise ValueError("Unsupported CoreML pipeline type: {0}".format(pipeline_type)) # Sequentially parse the sub-models for sub_model in sub_models: - # Declare the sub-model's input and output in this scope. Those input and output variables will be passed into - # the sub-model's parsing function and connected with proper child variables. + # Declare the sub-model's input and output in this scope. + # Those input and output variables will be passed into + # the sub-model's parsing function and connected with + # proper child variables. sub_inputs = [] for var in sub_model.description.input: variable = scope.get_local_variable_or_declare_one( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size)) + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + ) sub_inputs.append(variable) sub_outputs = [] for var in sub_model.description.output: variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size)) + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + ) sub_outputs.append(variable) _parse_model(topology, scope, sub_model, sub_inputs, sub_outputs) - # Declare the model's (not sub-model's) inputs and then link them with sub-model's inputs + # Declare the model's (not sub-model's) inputs and then + # link them with sub-model's inputs for var in model.description.input: - # Find the first variable with the same raw name declared when parsing the sub-models + # Find the first variable with the same raw name + # declared when parsing the sub-models child_variable = scope.variables[scope.variable_name_mapping[var.name][0]] - # Create model's input variable. Note that we set prepend=True because model inputs should not hide any + # Create model's input variable. Note that we set + # prepend=True because model inputs should not hide any # intermediate variables. variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size), - prepend=True) - # Feed the input to the sub-model's input. It's possible to add type conversion here by using a casting operator - # rather than identity, but we haven't see the need of doing so in practices. - operator = scope.declare_local_operator('identity') + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + prepend=True, + ) + # Feed the input to the sub-model's input. It's possible + # to add type conversion here by using a casting operator + # rather than identity, but we haven't see the need + # of doing so in practices. + operator = scope.declare_local_operator("identity") operator.inputs.append(variable) operator.outputs.append(child_variable) for parent_variable in inputs: raw_name = parent_variable.raw_name child_variable = scope.variables[scope.variable_name_mapping[raw_name][0]] - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(parent_variable) operator.outputs.append(child_variable) - # Declare the model's (not sub-model's) inputs and then link them with sub-model's inputs + # Declare the model's (not sub-model's) inputs and then link + # them with sub-model's inputs for var in model.description.output: - # Find the latest variable with the same raw name declared when parsing the sub-models + # Find the latest variable with the same raw name + # declared when parsing the sub-models child_variable = scope.variables[scope.variable_name_mapping[var.name][-1]] # Create model's output variable variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size)) - # Connect the input and a sub-model's input. It's possible to add type conversion here by using a casting - # operator rather than identity, but we haven't see the need of doing so in practices. - operator = scope.declare_local_operator('identity') + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + ) + # Connect the input and a sub-model's input. + # It's possible to add type conversion here by using a casting + # operator rather than identity, but we haven't see the + # need of doing so in practices. + operator = scope.declare_local_operator("identity") operator.inputs.append(child_variable) operator.outputs.append(variable) for parent_variable in outputs: raw_name = parent_variable.raw_name child_variable = scope.variables[scope.variable_name_mapping[raw_name][-1]] - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(child_variable) operator.outputs.append(parent_variable) def _parse_neural_network_model(topology, parent_scope, model, inputs, outputs): - ''' + """ Parse a neural network model. Steps: 1. Create local scope for allocating local variables and operators 2. Sequentially parse the preprocessors and layers - 3. Connect model's (neither layers' nor preprocessors') inputs and outputs with proper variables created when + 3. Connect model's (neither layers' nor preprocessors') + inputs and outputs with proper variables created when parsing sub-models. - 4. Link local variables and the corresponding parent variables (only model's inputs and outputs are considered) + 4. Link local variables and the corresponding parent + variables (only model's inputs and outputs are considered) Note: - 1. A CoreML preprocessor/layer can use the same variable for its input and output. + 1. A CoreML preprocessor/layer can use the same variable + for its input and output. 2. Two CoreML variables may have the same name but different types. - 3. Preprocessor sometime may not include any information about its input - ''' + 3. Preprocessor sometime may not include any information + about its input + """ # Create local scope to which all subsequent variables and operators belongs - scope = topology.declare_scope('NeuralNetwork', [parent_scope] + parent_scope.parent_scopes) + scope = topology.declare_scope( + "NeuralNetwork", [parent_scope] + parent_scope.parent_scopes + ) network = None - network_type = model.WhichOneof('Type') - if network_type == 'neuralNetworkClassifier': + network_type = model.WhichOneof("Type") + if network_type == "neuralNetworkClassifier": network = model.neuralNetworkClassifier - elif network_type == 'neuralNetworkRegressor': + elif network_type == "neuralNetworkRegressor": network = model.neuralNetworkRegressor - elif network_type == 'neuralNetwork': + elif network_type == "neuralNetwork": network = model.neuralNetwork else: - raise ValueError('Unknown network type {}'.format(network_type)) + raise ValueError("Unknown network type {}".format(network_type)) for op in network.preprocessing: - operator = scope.declare_local_operator(op.WhichOneof('preprocessor') + 'Preprocessor', op) + operator = scope.declare_local_operator( + op.WhichOneof("preprocessor") + "Preprocessor", op + ) - # Infer the variable name to be processed if feature name is an empty string - name = op.featureName if op.featureName != '' else model.description.input[0].name + # Infer the variable name to be processed if + # feature name is an empty string + name = ( + op.featureName if op.featureName != "" else model.description.input[0].name + ) # Find out input variable original = scope.get_local_variable_or_declare_one(name) - original.type = FloatTensorType() # A newly-declared variable has no type, so we add it. + original.type = ( + FloatTensorType() + ) # A newly-declared variable has no type, so we add it. operator.inputs.append(original) # Declare a variable for storing the processed result processed = scope.declare_local_variable(name) - processed.type = FloatTensorType() # A newly-declared variable has no type, so we add it + processed.type = ( + FloatTensorType() + ) # A newly-declared variable has no type, so we add it operator.outputs.append(processed) for op in network.layers: - operator = scope.declare_local_operator(op.WhichOneof('layer'), op) + operator = scope.declare_local_operator(op.WhichOneof("layer"), op) # Find out input variable and connect them with the operator for name in op.input: variable = scope.get_local_variable_or_declare_one(name) - # Although most neural network operators only accepts floats, we still need to handle the only exception, - # embedding layer. In the furture, we should create a Cast operator right inside embedding's converter. - if operator.type == 'embedding': + # Although most neural network operators only accepts + # floats, we still need to handle the only exception, + # embedding layer. In the furture, we should create + # a Cast operator right inside embedding's converter. + if operator.type == "embedding": variable.type = Int64TensorType() else: variable.type = FloatTensorType() @@ -321,97 +451,145 @@ def _parse_neural_network_model(topology, parent_scope, model, inputs, outputs): # Declare variables for catching the operator's outputs for name in op.output: variable = scope.declare_local_variable(name) - variable.type = FloatTensorType() # A newly-declared variable has no type, so we add it + variable.type = ( + FloatTensorType() + ) # A newly-declared variable has no type, so we add it operator.outputs.append(variable) sink_variables = scope.find_sink_variables() - # Declare the model's inputs and outputs. Then, connect them with proper variables computed by the main network + # Declare the model's inputs and outputs. + # Then, connect them with proper variables computed by the main network for var in model.description.input: - # Search for the first variable (declared when parsing network layers) associated with the considered raw name + # Search for the first variable (declared when parsing + # network layers) associated with the considered raw name child_variable = scope.variables[scope.variable_name_mapping[var.name][0]] - # Declare model input. To prevent intermediate variables form being hidden by model inputs, prepend is True. + # Declare model input. To prevent intermediate variables + # form being hidden by model inputs, prepend is True. variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size), - prepend=True) - - # A heuristic which forces the input of embedding to be integer tensor rather than float tensor. - # Ideally this should be done by adding a cast operator, but ONNX doesn't have float-to-int casting. - # If this variable is produced by another component in a CoreML pipeline, a bug may occur especially + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + prepend=True, + ) + + # A heuristic which forces the input of embedding to + # be integer tensor rather than float tensor. + # Ideally this should be done by adding a cast operator, + # but ONNX doesn't have float-to-int casting. + # If this variable is produced by another component in a + # CoreML pipeline, a bug may occur especially # when the source component's output type is float tensor. if isinstance(child_variable.type, Int64TensorType): variable.type = Int64TensorType(variable.type.shape) # Feed model input to the associated model input - operator_type = find_type_conversion(source_type=variable.type, target_type=child_variable.type) + operator_type = find_type_conversion( + source_type=variable.type, target_type=child_variable.type + ) operator = scope.declare_local_operator(operator_type) operator.inputs.append(variable) operator.outputs.append(child_variable) - # Connect local input variables with proper variables from parent scope + # Connect local input variables with proper variables + # from parent scope for parent_variable in inputs: raw_name = parent_variable.raw_name child_variable = scope.variables[scope.variable_name_mapping[raw_name][0]] - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(parent_variable) operator.outputs.append(child_variable) for var in model.description.output: - # CoreML's predicted label is not connected with any operator, so we handle it later as a special case. - special_variable_names = [model.description.predictedFeatureName, model.description.predictedProbabilitiesName] - if model.WhichOneof('Type') == 'neuralNetworkClassifier' and var.name in special_variable_names: + # CoreML's predicted label is not connected with any operator, + # so we handle it later as a special case. + special_variable_names = [ + model.description.predictedFeatureName, + model.description.predictedProbabilitiesName, + ] + if ( + model.WhichOneof("Type") == "neuralNetworkClassifier" + and var.name in special_variable_names + ): continue - # Search for the latest variable (declared when parsing network layers) associated with the considered raw name + # Search for the latest variable (declared when parsing + # network layers) associated with the considered raw name child_variable = scope.variables[scope.variable_name_mapping[var.name][-1]] # Create model output variable variable = scope.declare_local_variable( - var.name, _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size)) + var.name, + _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ), + ) # Feed result calculated by the network to the output variable - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(child_variable) operator.outputs.append(variable) - # If predicted label exists, connect probability tensor and label by a special operator - if model.WhichOneof('Type') == 'neuralNetworkClassifier' and model.description.predictedFeatureName: - # Find out the description of predicted label and declare a label variable + # If predicted label exists, connect probability tensor and + # label by a special operator + if ( + model.WhichOneof("Type") == "neuralNetworkClassifier" + and model.description.predictedFeatureName + ): + # Find out the description of predicted label and declare + # a label variable label_variable = None for var in model.description.output: if var.name == model.description.predictedFeatureName: - label_type = _parse_coreml_feature(var, topology.target_opset, topology.default_batch_size) + label_type = _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ) label_variable = scope.declare_local_variable(var.name, label_type) break - operator = scope.declare_local_operator('tensorToLabel', model) + operator = scope.declare_local_operator("tensorToLabel", model) probability_name = model.description.predictedProbabilitiesName if probability_name in scope.variable_name_mapping: # Find the latest probability variable - operator.inputs.append(scope.variables[scope.variable_name_mapping[probability_name][-1]]) + operator.inputs.append( + scope.variables[scope.variable_name_mapping[probability_name][-1]] + ) else: - # If predicted probability tensor is missing in CoreML model, it defaults to the first sink of the network + # If predicted probability tensor is missing in CoreML model, + # it defaults to the first sink of the network operator.inputs.append(sink_variables[0]) operator.outputs.append(label_variable) - # Probability tensor is implicitly converted into a dictionary (i.e., map) in CoreML. We handle this case here. - if model.WhichOneof('Type') == 'neuralNetworkClassifier' and model.description.predictedProbabilitiesName: - operator = scope.declare_local_operator('tensorToProbabilityMap', model) + # Probability tensor is implicitly converted into a dictionary + # (i.e., map) in CoreML. We handle this case here. + if ( + model.WhichOneof("Type") == "neuralNetworkClassifier" + and model.description.predictedProbabilitiesName + ): + operator = scope.declare_local_operator("tensorToProbabilityMap", model) probability_name = model.description.predictedProbabilitiesName if probability_name in scope.variable_name_mapping: # Find the latest probability variable - operator.inputs.append(scope.variables[scope.variable_name_mapping[probability_name][-1]]) + operator.inputs.append( + scope.variables[scope.variable_name_mapping[probability_name][-1]] + ) else: - # If predicted probability tensor is missing in CoreML model, it defaults to the first sink of the network + # If predicted probability tensor is missing in CoreML model, + # it defaults to the first sink of the network operator.inputs.append(sink_variables[0]) - # Find out the description of predicted probabilities and declare a variable for probability map + # Find out the description of predicted probabilities + # and declare a variable for probability map for var in model.description.output: if var.name == model.description.predictedProbabilitiesName: - probability_type = _parse_coreml_feature(var, topology.target_opset, - topology.default_batch_size) - probability_variable = scope.declare_local_variable(var.name, probability_type) + probability_type = _parse_coreml_feature( + var, topology.target_opset, topology.default_batch_size + ) + probability_variable = scope.declare_local_variable( + var.name, probability_type + ) operator.outputs.append(probability_variable) break @@ -419,54 +597,80 @@ def _parse_neural_network_model(topology, parent_scope, model, inputs, outputs): for parent_variable in outputs: raw_name = parent_variable.raw_name child_variable = scope.variables[scope.variable_name_mapping[raw_name][-1]] - operator = scope.declare_local_operator('identity') + operator = scope.declare_local_operator("identity") operator.inputs.append(child_variable) operator.outputs.append(parent_variable) -def parse_coreml(model, initial_types=None, target_opset=None, custom_conversion_functions=None, custom_shape_calculators=None): - ''' +def parse_coreml( + model, + initial_types=None, + target_opset=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): + """ This is the root function of the whole parsing procedure. :param model: CoreML model - :param initial_types: A list providing some types for some root variables. Each element is a tuple of a variable + :param initial_types: A list providing some types for + some root variables. Each element is a tuple of a variable name and a type defined in data_types.py. - :param target_opset: number, for example, 7 for ONNX 1.2, and 8 for ONNX 1.3. - :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: a Topology object. It's a intermediate representation of the input CoreML model - ''' - - # Add model-level input and output names into a set. The set will be fed into our Topology so that all its elements + :param target_opset: number, for example, + 7 for ONNX 1.2, and 8 for ONNX 1.3. + :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: a Topology object. It's a intermediate + representation of the input CoreML model + """ + + # Add model-level input and output names into a set. + # The set will be fed into our Topology so that all its elements # will not be used to declare variables reserved_variable_names = set() for var in list(model.description.input) + list(model.description.output): reserved_variable_names.add(var.name) - # Determine the batch size for parsing CoreML model's input and output features. Note that batch size is always + # Determine the batch size for parsing CoreML model's input + # and output features. Note that batch size is always # missing in all CoreML models. - default_batch_size = 'None' + default_batch_size = "None" - # Topology is shared by both of CoreML and scikit-learn conversion frameworks, so we have a wrapper class, - # CoremlModelContainer, to make sure our topology-related functions can seamlessly handle both of CoreML and + # Topology is shared by both of CoreML and scikit-learn + # conversion frameworks, so we have a wrapper class, + # CoremlModelContainer, to make sure our topology-related + # functions can seamlessly handle both of CoreML and # scikit-learn. - topology = Topology(CoremlModelContainer(model), - default_batch_size, - initial_types, - reserved_variable_names, - target_opset=target_opset, - custom_conversion_functions=custom_conversion_functions, - custom_shape_calculators=custom_shape_calculators) - scope = topology.declare_scope('__root__') - - # Instead of using CoremlModelContainer, we directly pass the model in because _parse_model is CoreML-specific. + topology = Topology( + CoremlModelContainer(model), + default_batch_size, + initial_types, + reserved_variable_names, + target_opset=target_opset, + custom_conversion_functions=custom_conversion_functions, + custom_shape_calculators=custom_shape_calculators, + ) + scope = topology.declare_scope("__root__") + + # Instead of using CoremlModelContainer, we directly + # pass the model in because _parse_model is CoreML-specific. _parse_model(topology, scope, model) topology.compile() for variable in topology.find_root_and_sink_variables(): - color_space = getattr(variable.type, 'color_space', None) + color_space = getattr(variable.type, "color_space", None) if color_space: - if topology.metadata_props.setdefault('Image.BitmapPixelFormat', color_space) != color_space: - warnings.warn('Conflicting pixel formats found. In ONNX, all input/output images must use the same pixel format.') + if ( + topology.metadata_props.setdefault( + "Image.BitmapPixelFormat", color_space + ) + != color_space + ): + warnings.warn( + "Conflicting pixel formats found. In ONNX, " + "all input/output images must use the same pixel format." + ) # Use original CoreML names for model-level input(s)/output(s) if variable.raw_name not in reserved_variable_names: continue diff --git a/onnxmltools/convert/coreml/convert.py b/onnxmltools/convert/coreml/convert.py index 4c692258..575a4b39 100644 --- a/onnxmltools/convert/coreml/convert.py +++ b/onnxmltools/convert/coreml/convert.py @@ -9,30 +9,43 @@ from ._parse import parse_coreml # Import modules to invoke function registrations -from . import operator_converters -from . import shape_calculators -from .operator_converters import neural_network as nn_converters -from .shape_calculators import neural_network as nn_shape_calculators -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 converts the specified CoreML model into its ONNX counterpart. Some information such as the produced +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 converts the specified CoreML model into its + ONNX counterpart. Some information such as the produced ONNX model name can be specified. - :param model: A `CoreML model `_ or - a CoreML MLModel object - :param initial_types: A list providing some types for some root variables. Each element is a tuple of a variable + :param model: A `CoreML model + `_ + or a CoreML MLModel object + :param initial_types: A list providing some types + for some root variables. Each element is a tuple of a variable name and a type defined in *data_types.py*. - :param name: The name of the graph (type: GraphProto) in the produced ONNX model (type: ModelProto) + :param name: The name of the graph (type: GraphProto) + in the produced ONNX model (type: ModelProto) :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 CoreML model + :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 CoreML model Example of initial types: Assume that 'A' and 'B' are two root variable names used in the CoreML @@ -43,7 +56,7 @@ def convert(model, name=None, initial_types=None, doc_string='', target_opset=No from onnxmltools.convert.common.data_types import FloatTensorType initial_type = [('A', FloatTensorType([40, 12, 1, 1])), ('B', FloatTensorType([1, 32, 1, 1]))] - ''' + """ if isinstance(model, coremltools.models.MLModel): spec = model.get_spec() else: @@ -54,27 +67,38 @@ def convert(model, name=None, initial_types=None, doc_string='', target_opset=No target_opset = target_opset if target_opset else get_maximum_opset_supported() # Parse CoreML model as our internal data structure (i.e., Topology) - topology = parse_coreml(spec, initial_types, target_opset, custom_conversion_functions, custom_shape_calculators) + topology = parse_coreml( + spec, + initial_types, + target_opset, + custom_conversion_functions, + custom_shape_calculators, + ) - # Parse CoreML description, author, and license. Those information will be attached to the final ONNX model. + # Parse CoreML description, author, and license. + # Those information will be attached to the final ONNX model. metadata = spec.description.metadata metadata_props = [] if metadata: if not doc_string and metadata.shortDescription: - doc_string = metadata.shortDescription # If doc_string is not specified, we use description from CoreML + doc_string = ( + metadata.shortDescription + ) # If doc_string is not specified, we use description from CoreML if metadata.author: entry = onnx_proto.StringStringEntryProto() - entry.key = 'author' + entry.key = "author" entry.value = metadata.author metadata_props.append(entry) if metadata.license: entry = onnx_proto.StringStringEntryProto() - entry.key = 'license' + entry.key = "license" entry.value = metadata.license metadata_props.append(entry) # Convert our Topology object into ONNX. The outcome is an ONNX model. - onnx_model = convert_topology(topology, name, doc_string, target_opset, targeted_onnx) + onnx_model = convert_topology( + topology, name, doc_string, target_opset, targeted_onnx + ) # Edit ONNX model's attributes related to CoreML's meta information if len(metadata_props) > 0: diff --git a/onnxmltools/convert/coreml/operator_converters/ArrayFeatureExtractor.py b/onnxmltools/convert/coreml/operator_converters/ArrayFeatureExtractor.py index a2580b25..5dd27867 100644 --- a/onnxmltools/convert/coreml/operator_converters/ArrayFeatureExtractor.py +++ b/onnxmltools/convert/coreml/operator_converters/ArrayFeatureExtractor.py @@ -5,17 +5,22 @@ def convert_array_feature_extractor(scope, operator, container): - op_type = 'ArrayFeatureExtractor' - attrs = {'name': operator.full_name} + op_type = "ArrayFeatureExtractor" + attrs = {"name": operator.full_name} target_indexes = operator.raw_operator.arrayFeatureExtractor.extractIndex - index_buffer_name = scope.get_unique_variable_name('target_indexes') - container.add_initializer(index_buffer_name, onnx_proto.TensorProto.INT64, [len(target_indexes)], target_indexes) + index_buffer_name = scope.get_unique_variable_name("target_indexes") + container.add_initializer( + index_buffer_name, + onnx_proto.TensorProto.INT64, + [len(target_indexes)], + target_indexes, + ) inputs = [operator.inputs[0].full_name, index_buffer_name] outputs = [operator.outputs[0].full_name] - container.add_node(op_type, inputs, outputs, op_domain='ai.onnx.ml', **attrs) + container.add_node(op_type, inputs, outputs, op_domain="ai.onnx.ml", **attrs) -register_converter('arrayFeatureExtractor', convert_array_feature_extractor) +register_converter("arrayFeatureExtractor", convert_array_feature_extractor) diff --git a/onnxmltools/convert/coreml/operator_converters/DictVectorizer.py b/onnxmltools/convert/coreml/operator_converters/DictVectorizer.py index 91cbaa3c..6c85eb9f 100644 --- a/onnxmltools/convert/coreml/operator_converters/DictVectorizer.py +++ b/onnxmltools/convert/coreml/operator_converters/DictVectorizer.py @@ -4,16 +4,21 @@ def convert_dictionary_vectorizer(scope, operator, container): - op_type = 'DictVectorizer' - attrs = {'name': operator.full_name} + op_type = "DictVectorizer" + attrs = {"name": operator.full_name} raw_model = operator.raw_operator.dictVectorizer - if raw_model.HasField('stringToIndex'): - attrs['string_vocabulary'] = raw_model.stringToIndex.vector + if raw_model.HasField("stringToIndex"): + attrs["string_vocabulary"] = raw_model.stringToIndex.vector else: - attrs['int64_vocabulary'] = raw_model.int64ToIndex.vector + attrs["int64_vocabulary"] = raw_model.int64ToIndex.vector - container.add_node(op_type, [operator.inputs[0].full_name], [operator.outputs[0].full_name], - op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + [operator.inputs[0].full_name], + [operator.outputs[0].full_name], + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('dictVectorizer', convert_dictionary_vectorizer) +register_converter("dictVectorizer", convert_dictionary_vectorizer) diff --git a/onnxmltools/convert/coreml/operator_converters/FeatureVectorizer.py b/onnxmltools/convert/coreml/operator_converters/FeatureVectorizer.py index dab128fc..2f6c45f3 100644 --- a/onnxmltools/convert/coreml/operator_converters/FeatureVectorizer.py +++ b/onnxmltools/convert/coreml/operator_converters/FeatureVectorizer.py @@ -5,27 +5,40 @@ def convert_feature_vectorizer(scope, operator, container): - op_type = 'FeatureVectorizer' - attrs = {'name': operator.full_name} + op_type = "FeatureVectorizer" + attrs = {"name": operator.full_name} inputs = [] input_dims = [] for variable in operator.inputs: if type(variable.type) in [Int64TensorType, Int64Type]: - # We use scaler to convert integers into floats because output is a single tensor and all tensor elements + # We use scaler to convert integers into floats + # because output is a single tensor and all tensor elements # should be in the same type. - scaler_name = scope.get_unique_operator_name('Scaler') - scaled_name = scope.get_unique_variable_name(variable.full_name + '_scaled') - scaler_attrs = {'name': scaler_name, 'scale': [1.], 'offset': [0.]} - container.add_node('Scaler', [variable.full_name], [scaled_name], op_domain='ai.onnx.ml', **scaler_attrs) + scaler_name = scope.get_unique_operator_name("Scaler") + scaled_name = scope.get_unique_variable_name(variable.full_name + "_scaled") + scaler_attrs = {"name": scaler_name, "scale": [1.0], "offset": [0.0]} + container.add_node( + "Scaler", + [variable.full_name], + [scaled_name], + op_domain="ai.onnx.ml", + **scaler_attrs + ) inputs.append(scaled_name) else: inputs.append(variable.full_name) # We assume feature vectorizer always combines inputs with shapes [1, C] or [C] input_dims.append(variable.type.shape[1]) - attrs['inputdimensions'] = input_dims + attrs["inputdimensions"] = input_dims - container.add_node(op_type, inputs, [operator.outputs[0].full_name], op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + inputs, + [operator.outputs[0].full_name], + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('featureVectorizer', convert_feature_vectorizer) +register_converter("featureVectorizer", convert_feature_vectorizer) diff --git a/onnxmltools/convert/coreml/operator_converters/GLMClassifier.py b/onnxmltools/convert/coreml/operator_converters/GLMClassifier.py index 016fb0bc..285dd1f7 100644 --- a/onnxmltools/convert/coreml/operator_converters/GLMClassifier.py +++ b/onnxmltools/convert/coreml/operator_converters/GLMClassifier.py @@ -44,20 +44,26 @@ def convert_glm_classifier(scope, operator, container): # ends at T' and T' is not linked # with other operators. from coremltools.proto.GLMClassifier_pb2 import GLMClassifier - op_type = 'LinearClassifier' - attrs = {'name': operator.full_name} - zipmap_attrs = {'name': scope.get_unique_operator_name('ZipMap')} + + op_type = "LinearClassifier" + attrs = {"name": operator.full_name} + zipmap_attrs = {"name": scope.get_unique_operator_name("ZipMap")} glm = operator.raw_operator.glmClassifier - transform_table = {GLMClassifier.Logit: 'LOGISTIC', GLMClassifier.Probit: 'PROBIT'} + transform_table = {GLMClassifier.Logit: "LOGISTIC", GLMClassifier.Probit: "PROBIT"} if glm.postEvaluationTransform not in transform_table: - raise ValueError('Unsupported post-transformation: {}'.format(glm.postEvaluationTransform)) - attrs['post_transform'] = transform_table[glm.postEvaluationTransform] - - encoding_table = {GLMClassifier.ReferenceClass: True, GLMClassifier.OneVsRest: False} + raise ValueError( + "Unsupported post-transformation: {}".format(glm.postEvaluationTransform) + ) + attrs["post_transform"] = transform_table[glm.postEvaluationTransform] + + encoding_table = { + GLMClassifier.ReferenceClass: True, + GLMClassifier.OneVsRest: False, + } if glm.classEncoding not in encoding_table: - raise ValueError('Unsupported class encoding: {}'.format(glm.classEncoding)) - attrs['multi_class'] = encoding_table[glm.classEncoding] + raise ValueError("Unsupported class encoding: {}".format(glm.classEncoding)) + attrs["multi_class"] = encoding_table[glm.classEncoding] # Determine the dimensionality of the model weights. dim_target = len(glm.weights) @@ -67,16 +73,16 @@ def convert_glm_classifier(scope, operator, container): for i, w in enumerate(glm.weights): matrix_w[i, :] = w.value - if glm.WhichOneof('ClassLabels') == 'stringClassLabels': - class_labels = list(s.encode('utf-8') for s in glm.stringClassLabels.vector) - attrs['classlabels_strings'] = class_labels - zipmap_attrs['classlabels_strings'] = class_labels - elif glm.WhichOneof('ClassLabels') == 'int64ClassLabels': + if glm.WhichOneof("ClassLabels") == "stringClassLabels": + class_labels = list(s.encode("utf-8") for s in glm.stringClassLabels.vector) + attrs["classlabels_strings"] = class_labels + zipmap_attrs["classlabels_strings"] = class_labels + elif glm.WhichOneof("ClassLabels") == "int64ClassLabels": class_labels = list(int(i) for i in glm.int64ClassLabels.vector) - attrs['classlabels_ints'] = class_labels - zipmap_attrs['classlabels_int64s'] = class_labels + attrs["classlabels_ints"] = class_labels + zipmap_attrs["classlabels_int64s"] = class_labels else: - raise ValueError('Unknown class label type') + raise ValueError("Unknown class label type") coefficients = matrix_w.flatten().tolist() intercepts = list(float(i) for i in glm.offset) @@ -85,8 +91,8 @@ def convert_glm_classifier(scope, operator, container): coefficients = list(map(lambda x: -1 * x, coefficients)) + coefficients intercepts = list(map(lambda x: -1 * x, intercepts)) + intercepts - attrs['coefficients'] = coefficients - attrs['intercepts'] = intercepts + attrs["coefficients"] = coefficients + attrs["intercepts"] = intercepts # For classifiers, due to the different representation of classes' probabilities, we need to add some # operators for type conversion. It turns out that we have the following topology. @@ -101,36 +107,70 @@ def convert_glm_classifier(scope, operator, container): for variable in operator.outputs: if raw_model.description.predictedFeatureName == variable.raw_name: label_output_name = variable.full_name - if raw_model.description.predictedProbabilitiesName != '' and \ - raw_model.description.predictedProbabilitiesName == variable.raw_name: + if ( + raw_model.description.predictedProbabilitiesName != "" + and raw_model.description.predictedProbabilitiesName == variable.raw_name + ): proba_output_name = variable.full_name inputs = [variable.full_name for variable in operator.inputs] - proba_tensor_name = scope.get_unique_variable_name('ProbabilityTensor') + proba_tensor_name = scope.get_unique_variable_name("ProbabilityTensor") if proba_output_name is not None: # Add tree model ONNX node with probability output - container.add_node(op_type, inputs, [label_output_name, proba_tensor_name], op_domain='ai.onnx.ml', **attrs) - - # Add a normalizer to make sure that the sum of all classes' probabilities is 1. It doesn't affect binary - # classification. For multi-class clssifiers, if one applies sigmoid function independently to all raw scores, - # we have to add a normalization so that the sum of all probabilities remains 1. Of course, if softmax is used - # to convert raw scores into probabilities, this normalization doesn't change anything. + container.add_node( + op_type, + inputs, + [label_output_name, proba_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) + + # Add a normalizer to make sure that the sum of all + # classes' probabilities is 1. It doesn't affect binary + # classification. For multi-class clssifiers, + # if one applies sigmoid function independently to all raw scores, + # we have to add a normalization so that the sum of + # all probabilities remains 1. Of course, if softmax is used + # to convert raw scores into probabilities, + # this normalization doesn't change anything. if len(class_labels) > 2: - normalized_proba_tensor_name = scope.get_unique_variable_name(proba_tensor_name + '_normalized') - container.add_node('Normalizer', proba_tensor_name, normalized_proba_tensor_name, op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('Normalizer'), norm='L1') + normalized_proba_tensor_name = scope.get_unique_variable_name( + proba_tensor_name + "_normalized" + ) + container.add_node( + "Normalizer", + proba_tensor_name, + normalized_proba_tensor_name, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("Normalizer"), + norm="L1", + ) else: - # If we don't need a normalization, we just pass the original probability tensor to the following ZipMap + # If we don't need a normalization, we just pass the + # original probability tensor to the following ZipMap normalized_proba_tensor_name = proba_tensor_name - # Add ZipMap to convert normalized probability tensor into probability map - container.add_node('ZipMap', [normalized_proba_tensor_name], [proba_output_name], - op_domain='ai.onnx.ml', **zipmap_attrs) + # Add ZipMap to convert normalized probability tensor + # into probability map + container.add_node( + "ZipMap", + [normalized_proba_tensor_name], + [proba_output_name], + op_domain="ai.onnx.ml", + **zipmap_attrs + ) else: - # Add linear classifier with isolated probability output, which means that the probability + # Add linear classifier with isolated probability + # output, which means that the probability # tensor won't be accessed by any others. - container.add_node(op_type, inputs, [label_output_name, proba_tensor_name], op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + inputs, + [label_output_name, proba_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('glmClassifier', convert_glm_classifier) +register_converter("glmClassifier", convert_glm_classifier) diff --git a/onnxmltools/convert/coreml/operator_converters/GLMRegressor.py b/onnxmltools/convert/coreml/operator_converters/GLMRegressor.py index 87282c3e..813c1d85 100644 --- a/onnxmltools/convert/coreml/operator_converters/GLMRegressor.py +++ b/onnxmltools/convert/coreml/operator_converters/GLMRegressor.py @@ -7,15 +7,21 @@ def convert_glm_regressor(scope, operator, container): from coremltools.proto.GLMRegressor_pb2 import GLMRegressor - op_type = 'LinearRegressor' + op_type = "LinearRegressor" glm = operator.raw_operator.glmRegressor - attrs = {'name': operator.full_name} + attrs = {"name": operator.full_name} - transform_table = {GLMRegressor.NoTransform: 'NONE', GLMRegressor.Logit: 'LOGISTIC', GLMRegressor.Probit: 'PROBIT'} + transform_table = { + GLMRegressor.NoTransform: "NONE", + GLMRegressor.Logit: "LOGISTIC", + GLMRegressor.Probit: "PROBIT", + } if glm.postEvaluationTransform in transform_table: - attrs['post_transform'] = transform_table[glm.postEvaluationTransform] + attrs["post_transform"] = transform_table[glm.postEvaluationTransform] else: - raise ValueError('Unsupported post-transformation: {}'.format(glm.postEvaluationTransform)) + raise ValueError( + "Unsupported post-transformation: {}".format(glm.postEvaluationTransform) + ) # Determine the dimensionality of the model weights. Conceptually, # the shape of the weight matrix in CoreML is E-by-F, where E and F @@ -28,11 +34,17 @@ def convert_glm_regressor(scope, operator, container): for i, w in enumerate(glm.weights): matrix_w[:, i] = w.value - attrs['targets'] = dim_target - attrs['coefficients'] = matrix_w.flatten() - attrs['intercepts'] = glm.offset + attrs["targets"] = dim_target + attrs["coefficients"] = matrix_w.flatten() + attrs["intercepts"] = glm.offset - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('glmRegressor', convert_glm_regressor) +register_converter("glmRegressor", convert_glm_regressor) diff --git a/onnxmltools/convert/coreml/operator_converters/Identity.py b/onnxmltools/convert/coreml/operator_converters/Identity.py index 3b1efcb0..47766e7a 100644 --- a/onnxmltools/convert/coreml/operator_converters/Identity.py +++ b/onnxmltools/convert/coreml/operator_converters/Identity.py @@ -4,7 +4,12 @@ def convert_identity(scope, operator, container): - container.add_node('Identity', operator.input_full_names, operator.output_full_names, name=operator.full_name) + container.add_node( + "Identity", + operator.input_full_names, + operator.output_full_names, + name=operator.full_name, + ) -register_converter('identity', convert_identity) +register_converter("identity", convert_identity) diff --git a/onnxmltools/convert/coreml/operator_converters/Imputer.py b/onnxmltools/convert/coreml/operator_converters/Imputer.py index 3f04a0df..987ef511 100644 --- a/onnxmltools/convert/coreml/operator_converters/Imputer.py +++ b/onnxmltools/convert/coreml/operator_converters/Imputer.py @@ -4,18 +4,24 @@ def convert_imputer(scope, operator, container): - op_type = 'Imputer' - attrs = {'name': operator.full_name} + op_type = "Imputer" + attrs = {"name": operator.full_name} imputer = operator.raw_operator.imputer - if imputer.HasField('replaceDoubleValue'): - attrs['replaced_value_float'] = imputer.replaceDoubleValue - elif imputer.HasField('replaceInt64Value'): - attrs['replaced_value_int64'] = imputer.replaceInt64Value - if imputer.HasField('imputedDoubleArray'): - attrs['imputed_value_floats'] = imputer.imputedDoubleArray.vector - elif imputer.HasField('imputedInt64Array'): - attrs['imputed_value_int64s'] = imputer.imputedInt64Array.vector - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + if imputer.HasField("replaceDoubleValue"): + attrs["replaced_value_float"] = imputer.replaceDoubleValue + elif imputer.HasField("replaceInt64Value"): + attrs["replaced_value_int64"] = imputer.replaceInt64Value + if imputer.HasField("imputedDoubleArray"): + attrs["imputed_value_floats"] = imputer.imputedDoubleArray.vector + elif imputer.HasField("imputedInt64Array"): + attrs["imputed_value_int64s"] = imputer.imputedInt64Array.vector + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('imputer', convert_imputer) +register_converter("imputer", convert_imputer) diff --git a/onnxmltools/convert/coreml/operator_converters/Normalizer.py b/onnxmltools/convert/coreml/operator_converters/Normalizer.py index b101b7b8..e6fa0d70 100644 --- a/onnxmltools/convert/coreml/operator_converters/Normalizer.py +++ b/onnxmltools/convert/coreml/operator_converters/Normalizer.py @@ -4,16 +4,22 @@ def convert_normalizer(scope, operator, container): - op_type = 'Normalizer' - attrs = {'name': operator.full_name} - norms = ['MAX', 'L1', 'L2'] + op_type = "Normalizer" + attrs = {"name": operator.full_name} + norms = ["MAX", "L1", "L2"] norm_type = operator.raw_operator.normalizer.normType if norm_type in range(3): - attrs['norm'] = norms[norm_type] + attrs["norm"] = norms[norm_type] else: - raise ValueError('Invalid norm type: ' + norm_type) + raise ValueError("Invalid norm type: " + norm_type) - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('normalizer', convert_normalizer) +register_converter("normalizer", convert_normalizer) diff --git a/onnxmltools/convert/coreml/operator_converters/OneHotEncoder.py b/onnxmltools/convert/coreml/operator_converters/OneHotEncoder.py index 2af86e5b..98014d10 100644 --- a/onnxmltools/convert/coreml/operator_converters/OneHotEncoder.py +++ b/onnxmltools/convert/coreml/operator_converters/OneHotEncoder.py @@ -4,17 +4,24 @@ def convert_one_hot_encoder(scope, operator, container): - op_type = 'OneHotEncoder' - attrs = {'name': operator.full_name} + op_type = "OneHotEncoder" + attrs = {"name": operator.full_name} raw_model = operator.raw_operator.oneHotEncoder - if raw_model.HasField('int64Categories'): - attrs['cats_int64s'] = list(int(i) for i in raw_model.int64Categories.vector) - if raw_model.HasField('stringCategories'): - attrs['cats_strings'] = list(str(s).encode('utf-8') for s in raw_model.stringCategories.vector) + if raw_model.HasField("int64Categories"): + attrs["cats_int64s"] = list(int(i) for i in raw_model.int64Categories.vector) + if raw_model.HasField("stringCategories"): + attrs["cats_strings"] = list( + str(s).encode("utf-8") for s in raw_model.stringCategories.vector + ) - container.add_node(op_type, [operator.inputs[0].full_name], [operator.outputs[0].full_name], - op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + [operator.inputs[0].full_name], + [operator.outputs[0].full_name], + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('oneHotEncoder', convert_one_hot_encoder) +register_converter("oneHotEncoder", convert_one_hot_encoder) diff --git a/onnxmltools/convert/coreml/operator_converters/SVC.py b/onnxmltools/convert/coreml/operator_converters/SVC.py index 8c33b645..a8e06d9f 100644 --- a/onnxmltools/convert/coreml/operator_converters/SVC.py +++ b/onnxmltools/convert/coreml/operator_converters/SVC.py @@ -6,11 +6,11 @@ def extract_support_vectors_as_dense_tensor(svm_model): - support_type = svm_model.WhichOneof('supportVectors') - if support_type == 'denseSupportVectors': + support_type = svm_model.WhichOneof("supportVectors") + if support_type == "denseSupportVectors": vectors = svm_model.denseSupportVectors.vectors support_vectors = np.array([v.values for v in vectors]).flatten() - elif support_type == 'sparseSupportVectors': + elif support_type == "sparseSupportVectors": # Since current ONNX doesn't support sparse representations, we have to save them as dense tensors. It may # dramatically prolong prediction time and increase memory usage. vectors = svm_model.sparseSupportVectors.vectors @@ -27,69 +27,87 @@ def extract_support_vectors_as_dense_tensor(svm_model): support_vectors[i][n.index - 1] = n.value support_vectors = support_vectors.flatten() else: - raise ValueError('Unsupported support vector type: %s' % support_type) + raise ValueError("Unsupported support vector type: %s" % support_type) return len(vectors), support_vectors def convert_svm_classifier(scope, operator, container): params = operator.raw_operator.supportVectorClassifier - kernel_enum = {'linearKernel': 'LINEAR', 'polyKernel': 'POLY', - 'rbfKernel': 'RBF', 'sigmoidKernel': 'SIGMOID', 'precomputedKernel': 'PRECOMPUTED'} + kernel_enum = { + "linearKernel": "LINEAR", + "polyKernel": "POLY", + "rbfKernel": "RBF", + "sigmoidKernel": "SIGMOID", + "precomputedKernel": "PRECOMPUTED", + } kernel = params.kernel - kernel_val = kernel.WhichOneof('kernel') + kernel_val = kernel.WhichOneof("kernel") svc_kernel = kernel_enum[kernel_val] - if kernel_val == 'rbfKernel': + if kernel_val == "rbfKernel": svc_kernel_params = [kernel.rbfKernel.gamma, 0.0, 0.0] - elif kernel_val == 'polyKernel': - svc_kernel_params = [kernel.polyKernel.gamma, - kernel.polyKernel.coef0, kernel.polyKernel.degree] - elif kernel_val == 'sigmoidKernel': - svc_kernel_params = [kernel.sigmoidKernel.gamma, - kernel.sigmoidKernel.coef0, 0.0] - elif kernel_val == 'linearKernel': + elif kernel_val == "polyKernel": + svc_kernel_params = [ + kernel.polyKernel.gamma, + kernel.polyKernel.coef0, + kernel.polyKernel.degree, + ] + elif kernel_val == "sigmoidKernel": + svc_kernel_params = [ + kernel.sigmoidKernel.gamma, + kernel.sigmoidKernel.coef0, + 0.0, + ] + elif kernel_val == "linearKernel": svc_kernel_params = [0.0, 0.0, 0.0] prob_a = params.probA prob_b = params.probB support_vectors_per_class = params.numberOfSupportVectorsPerClass n_supports, svc_support_vectors = extract_support_vectors_as_dense_tensor( - operator.raw_operator.supportVectorClassifier) - chain_coef = list(itertools.chain.from_iterable([coef.alpha for coef in params.coefficients])) + operator.raw_operator.supportVectorClassifier + ) + chain_coef = list( + itertools.chain.from_iterable([coef.alpha for coef in params.coefficients]) + ) svc_coefficients = chain_coef svc_rho = [-x for x in params.rho] - op_type = 'SVMClassifier' + op_type = "SVMClassifier" op_name = scope.get_unique_operator_name(op_type) - attrs = {'name': op_name} - attrs['kernel_type'] = svc_kernel - attrs['kernel_params'] = svc_kernel_params + attrs = {"name": op_name} + attrs["kernel_type"] = svc_kernel + attrs["kernel_params"] = svc_kernel_params if prob_a: - attrs['prob_a'] = prob_a + attrs["prob_a"] = prob_a if prob_b: - attrs['prob_b'] = prob_b - attrs['vectors_per_class'] = support_vectors_per_class - attrs['support_vectors'] = svc_support_vectors - attrs['coefficients'] = svc_coefficients - attrs['rho'] = svc_rho - zipmap_attrs = {'name': scope.get_unique_operator_name('ZipMap')} - svc_classes = params.WhichOneof('ClassLabels') - if svc_classes == 'int64ClassLabels': + attrs["prob_b"] = prob_b + attrs["vectors_per_class"] = support_vectors_per_class + attrs["support_vectors"] = svc_support_vectors + attrs["coefficients"] = svc_coefficients + attrs["rho"] = svc_rho + zipmap_attrs = {"name": scope.get_unique_operator_name("ZipMap")} + svc_classes = params.WhichOneof("ClassLabels") + if svc_classes == "int64ClassLabels": class_labels = list(int(i) for i in params.int64ClassLabels.vector) - attrs['classlabels_ints'] = class_labels - zipmap_attrs['classlabels_int64s'] = class_labels - elif svc_classes == 'stringClassLabels': - class_labels = list(str(s).encode('utf-8') for s in params.stringClassLabels.vector) - attrs['classlabels_strings'] = class_labels - zipmap_attrs['classlabels_strings'] = class_labels + attrs["classlabels_ints"] = class_labels + zipmap_attrs["classlabels_int64s"] = class_labels + elif svc_classes == "stringClassLabels": + class_labels = list( + str(s).encode("utf-8") for s in params.stringClassLabels.vector + ) + attrs["classlabels_strings"] = class_labels + zipmap_attrs["classlabels_strings"] = class_labels else: - raise ValueError('Unknown class label type') + raise ValueError("Unknown class label type") - # For classifiers, due to the different representation of classes' probabilities, we need to add some - # operators for type conversion. It turns out that we have the following topology. + # For classifiers, due to the different representation + # of classes' probabilities, we need to add some + # operators for type conversion. It turns out that we + # have the following topology. # input features ---> SupportVectorClassifier ---> label (must present) # | - # '--> probability tensor ---> ZipMap ---> probability map (optional) + # '--> probability tensor raw_model = operator.raw_operator # Find label name and probability name @@ -97,23 +115,42 @@ def convert_svm_classifier(scope, operator, container): for variable in operator.outputs: if raw_model.description.predictedFeatureName == variable.raw_name: label_output_name = variable.full_name - if raw_model.description.predictedProbabilitiesName != '' and \ - raw_model.description.predictedProbabilitiesName == variable.raw_name: + if ( + raw_model.description.predictedProbabilitiesName != "" + and raw_model.description.predictedProbabilitiesName == variable.raw_name + ): proba_output_name = variable.full_name inputs = [variable.full_name for variable in operator.inputs] - proba_tensor_name = scope.get_unique_variable_name('ProbabilityTensor') + proba_tensor_name = scope.get_unique_variable_name("ProbabilityTensor") if proba_output_name is not None: # Add support vector classifier in terms of ONNX node with probability output - container.add_node(op_type, inputs, [label_output_name, proba_tensor_name], op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + inputs, + [label_output_name, proba_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) # Add ZipMap to convert probability tensor into probability map - container.add_node('ZipMap', [proba_tensor_name], [proba_output_name], - op_domain='ai.onnx.ml', **zipmap_attrs) + container.add_node( + "ZipMap", + [proba_tensor_name], + [proba_output_name], + op_domain="ai.onnx.ml", + **zipmap_attrs + ) else: # Add support vector classifier in terms of ONNX node - container.add_node(op_type, inputs, [label_output_name, proba_tensor_name], op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + inputs, + [label_output_name, proba_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('supportVectorClassifier', convert_svm_classifier) +register_converter("supportVectorClassifier", convert_svm_classifier) diff --git a/onnxmltools/convert/coreml/operator_converters/SVR.py b/onnxmltools/convert/coreml/operator_converters/SVR.py index ae653ed0..25b00d3e 100644 --- a/onnxmltools/convert/coreml/operator_converters/SVR.py +++ b/onnxmltools/convert/coreml/operator_converters/SVR.py @@ -7,24 +7,37 @@ def convert_svm_regressor(scope, operator, container): params = operator.raw_operator.supportVectorRegressor - kernel_enum = {'linearKernel': 'LINEAR', 'polyKernel': 'POLY', - 'rbfKernel': 'RBF', 'sigmoidKernel': 'SIGMOID', 'precomputedKernel': 'PRECOMPUTED'} + kernel_enum = { + "linearKernel": "LINEAR", + "polyKernel": "POLY", + "rbfKernel": "RBF", + "sigmoidKernel": "SIGMOID", + "precomputedKernel": "PRECOMPUTED", + } kernel = params.kernel - kernel_val = kernel.WhichOneof('kernel') + kernel_val = kernel.WhichOneof("kernel") svr_kernel = kernel_enum[kernel_val] - if kernel_val == 'rbfKernel': + if kernel_val == "rbfKernel": svr_kernel_params = [kernel.rbfKernel.gamma, 0.0, 0.0] - elif kernel_val == 'polyKernel': - svr_kernel_params = [kernel.polyKernel.gamma, - kernel.polyKernel.coef0, kernel.polyKernel.degree] - elif kernel_val == 'sigmoidKernel': - svr_kernel_params = [kernel.sigmoidKernel.gamma, - kernel.sigmoidKernel.coef0, 0.0] - elif kernel_val == 'linearKernel': + elif kernel_val == "polyKernel": + svr_kernel_params = [ + kernel.polyKernel.gamma, + kernel.polyKernel.coef0, + kernel.polyKernel.degree, + ] + elif kernel_val == "sigmoidKernel": + svr_kernel_params = [ + kernel.sigmoidKernel.gamma, + kernel.sigmoidKernel.coef0, + 0.0, + ] + elif kernel_val == "linearKernel": svr_kernel_params = [0.0, 0.0, 0.0] - n_supports, support_vectors = extract_support_vectors_as_dense_tensor(operator.raw_operator.supportVectorRegressor) + n_supports, support_vectors = extract_support_vectors_as_dense_tensor( + operator.raw_operator.supportVectorRegressor + ) svr_coefficients = params.coefficients.alpha if isinstance(params.rho, list): @@ -32,17 +45,23 @@ def convert_svm_regressor(scope, operator, container): else: svr_rho = [-params.rho] - op_type = 'SVMRegressor' + op_type = "SVMRegressor" op_name = scope.get_unique_operator_name(op_type) - attrs = {'name': op_name} - attrs['kernel_type'] = svr_kernel - attrs['kernel_params'] = svr_kernel_params - attrs['support_vectors'] = support_vectors - attrs['n_supports'] = n_supports - attrs['coefficients'] = svr_coefficients - attrs['rho'] = svr_rho + attrs = {"name": op_name} + attrs["kernel_type"] = svr_kernel + attrs["kernel_params"] = svr_kernel_params + attrs["support_vectors"] = support_vectors + attrs["n_supports"] = n_supports + attrs["coefficients"] = svr_coefficients + attrs["rho"] = svr_rho - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('supportVectorRegressor', convert_svm_regressor) +register_converter("supportVectorRegressor", convert_svm_regressor) diff --git a/onnxmltools/convert/coreml/operator_converters/Scaler.py b/onnxmltools/convert/coreml/operator_converters/Scaler.py index 2d68aa56..d6e8eb99 100644 --- a/onnxmltools/convert/coreml/operator_converters/Scaler.py +++ b/onnxmltools/convert/coreml/operator_converters/Scaler.py @@ -4,17 +4,23 @@ def convert_scaler(scope, operator, container): - op_type = 'Scaler' - attrs = {'name': operator.full_name} + op_type = "Scaler" + attrs = {"name": operator.full_name} scaler = operator.raw_operator.scaler scale = [x for x in scaler.scaleValue] offset = [-x for x in scaler.shiftValue] - attrs['scale'] = scale - attrs['offset'] = offset + attrs["scale"] = scale + attrs["offset"] = offset - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('scaler', convert_scaler) +register_converter("scaler", convert_scaler) diff --git a/onnxmltools/convert/coreml/operator_converters/TensorToLabel.py b/onnxmltools/convert/coreml/operator_converters/TensorToLabel.py index 0edecb32..fbfd98f2 100644 --- a/onnxmltools/convert/coreml/operator_converters/TensorToLabel.py +++ b/onnxmltools/convert/coreml/operator_converters/TensorToLabel.py @@ -5,8 +5,9 @@ from ...common._registration import register_converter from ...common._apply_operation import apply_constant + def convert_tensor_to_label(scope, operator, container): - ''' + """ This converter tries to convert a dummy operator 'TensorToLabel' into a sequence of some ONNX operators. Those operators are used to extract the label with the highest probability for doing a prediction. We assume that the elements in the given probability tensor are aligned with the class labels specified in the CoreML model. That is, @@ -32,51 +33,73 @@ def convert_tensor_to_label(scope, operator, container): | v predicted label [1] - ''' - model_type = operator.raw_operator.WhichOneof('Type') - if model_type == 'neuralNetworkClassifier': + """ + model_type = operator.raw_operator.WhichOneof("Type") + if model_type == "neuralNetworkClassifier": model = operator.raw_operator.neuralNetworkClassifier - if model.WhichOneof('ClassLabels') == 'stringClassLabels': - labels = list(s.encode('utf-8') for s in model.stringClassLabels.vector) + if model.WhichOneof("ClassLabels") == "stringClassLabels": + labels = list(s.encode("utf-8") for s in model.stringClassLabels.vector) label_type = onnx_proto.TensorProto.STRING - elif model.WhichOneof('ClassLabels') == 'int64ClassLabels': + elif model.WhichOneof("ClassLabels") == "int64ClassLabels": labels = list(int(i) for i in model.int64ClassLabels.vector) label_type = onnx_proto.TensorProto.INT64 else: - raise ValueError('Unknown label type found') - elif model_type == 'pipelineClassifier': + raise ValueError("Unknown label type found") + elif model_type == "pipelineClassifier": model = operator.raw_operator.pipelineClassifier - if model.WhichOneof('ClassLabels') == 'stringClassLabels': - labels = list(s.encode('utf-8') for s in model.pipelineClassifier.stringClassLabels.vector) + if model.WhichOneof("ClassLabels") == "stringClassLabels": + labels = list( + s.encode("utf-8") + for s in model.pipelineClassifier.stringClassLabels.vector + ) label_type = onnx_proto.TensorProto.STRING - elif model.WhichOneof('ClassLabels') == 'int64ClassLabels': + elif model.WhichOneof("ClassLabels") == "int64ClassLabels": labels = list(int(i) for i in model.int64ClassLabels.vector) label_type = onnx_proto.TensorProto.INT64 else: - raise ValueError('Unknown label type found') + raise ValueError("Unknown label type found") else: - raise ValueError('Only neural network classifiers and pipeline classifiers are supported') + raise ValueError( + "Only neural network classifiers and pipeline classifiers are supported" + ) # Use a Constant operator to load and output all labels as a tensor - label_loader_name = scope.get_unique_operator_name('LabelLoader') - label_buffer_name = scope.get_unique_variable_name('ClassLabels') - label_loader_value = helper.make_tensor(label_buffer_name, label_type, [len(labels)], labels) - apply_constant(scope, [label_buffer_name], container, - operator_name=label_loader_name, value=label_loader_value) + label_loader_name = scope.get_unique_operator_name("LabelLoader") + label_buffer_name = scope.get_unique_variable_name("ClassLabels") + label_loader_value = helper.make_tensor( + label_buffer_name, label_type, [len(labels)], labels + ) + apply_constant( + scope, + [label_buffer_name], + container, + operator_name=label_loader_name, + value=label_loader_value, + ) # Extract most possible label index - label_id_extractor_name = scope.get_unique_operator_name('LabelIndexExtractor') - label_id_extractor_attrs = {'name': label_id_extractor_name} - label_id_extractor_attrs['axis'] = 1 - label_id_extractor_attrs['keepdims'] = 1 - extracted_id_name = scope.get_unique_variable_name('LabelId') - container.add_node('ArgMax', [operator.inputs[0].full_name], [extracted_id_name], **label_id_extractor_attrs) + label_id_extractor_name = scope.get_unique_operator_name("LabelIndexExtractor") + label_id_extractor_attrs = {"name": label_id_extractor_name} + label_id_extractor_attrs["axis"] = 1 + label_id_extractor_attrs["keepdims"] = 1 + extracted_id_name = scope.get_unique_variable_name("LabelId") + container.add_node( + "ArgMax", + [operator.inputs[0].full_name], + [extracted_id_name], + **label_id_extractor_attrs + ) # Pick up the label indicated by the selected ID - label_selector_name = scope.get_unique_operator_name('LabelSelector') - label_selector_attrs = {'name': label_selector_name} - container.add_node('ArrayFeatureExtractor', [label_buffer_name, extracted_id_name], [operator.outputs[0].full_name], - op_domain='ai.onnx.ml', **label_selector_attrs) + label_selector_name = scope.get_unique_operator_name("LabelSelector") + label_selector_attrs = {"name": label_selector_name} + container.add_node( + "ArrayFeatureExtractor", + [label_buffer_name, extracted_id_name], + [operator.outputs[0].full_name], + op_domain="ai.onnx.ml", + **label_selector_attrs + ) -register_converter('tensorToLabel', convert_tensor_to_label) +register_converter("tensorToLabel", convert_tensor_to_label) diff --git a/onnxmltools/convert/coreml/operator_converters/TensorToProbabilityMap.py b/onnxmltools/convert/coreml/operator_converters/TensorToProbabilityMap.py index 1fdaada2..ca0f41d3 100644 --- a/onnxmltools/convert/coreml/operator_converters/TensorToProbabilityMap.py +++ b/onnxmltools/convert/coreml/operator_converters/TensorToProbabilityMap.py @@ -6,61 +6,92 @@ def convert_tensor_to_probability_map(scope, operator, container): - ''' - This converter tries to convert a special operator 'TensorToProbabilityMap' into a sequence of some ONNX operators. - Those operators are used to create a dictionary in which keys are class labels and values are the associated - probabilities. We assume that the elements in the given probability tensor are aligned with the class labels + """ + This converter tries to convert a special operator + 'TensorToProbabilityMap' into a sequence of some ONNX operators. + Those operators are used to create a dictionary in which keys + are class labels and values are the associated + probabilities. We assume that the elements in the given probability + tensor are aligned with the class labels specified in the CoreML model. - Notice that ONNX<1.2 doesn't support a CoreML classifier with a batch size larger than one because old ONNX ZipMap - is not able to produce a sequence of dictionaries. This issue has been fixed in ONNX-1.2. - ''' - attrs = {'name': scope.get_unique_operator_name('ZipMap')} + Notice that ONNX<1.2 doesn't support a CoreML classifier + with a batch size larger than one because old ONNX ZipMap + is not able to produce a sequence of dictionaries. + This issue has been fixed in ONNX-1.2. + """ + attrs = {"name": scope.get_unique_operator_name("ZipMap")} - model_type = operator.raw_operator.WhichOneof('Type') - if model_type == 'neuralNetworkClassifier': + model_type = operator.raw_operator.WhichOneof("Type") + if model_type == "neuralNetworkClassifier": model = operator.raw_operator.neuralNetworkClassifier - if model.WhichOneof('ClassLabels') == 'stringClassLabels': - attrs['classlabels_strings'] = list(s.encode('utf-8') for s in model.stringClassLabels.vector) - elif model.WhichOneof('ClassLabels') == 'int64ClassLabels': - attrs['classlabels_int64s'] = list(int(i) for i in model.int64ClassLabels.vector) + if model.WhichOneof("ClassLabels") == "stringClassLabels": + attrs["classlabels_strings"] = list( + s.encode("utf-8") for s in model.stringClassLabels.vector + ) + elif model.WhichOneof("ClassLabels") == "int64ClassLabels": + attrs["classlabels_int64s"] = list( + int(i) for i in model.int64ClassLabels.vector + ) else: - raise ValueError('Unknown label type found') - elif model_type == 'pipelineClassifier': + raise ValueError("Unknown label type found") + elif model_type == "pipelineClassifier": model = operator.raw_operator.pipelineClassifier - if model.WhichOneof('ClassLabels') == 'stringClassLabels': - attrs['classlabels_strings'] = list(s.encode('utf-8') for s in model.stringClassLabels.vector) - elif model.WhichOneof('ClassLabels') == 'int64ClassLabels': - attrs['classlabels_int64s'] = list(int(i) for i in model.int64ClassLabels.vector) + if model.WhichOneof("ClassLabels") == "stringClassLabels": + attrs["classlabels_strings"] = list( + s.encode("utf-8") for s in model.stringClassLabels.vector + ) + elif model.WhichOneof("ClassLabels") == "int64ClassLabels": + attrs["classlabels_int64s"] = list( + int(i) for i in model.int64ClassLabels.vector + ) else: - raise ValueError('Unknown label type found') + raise ValueError("Unknown label type found") else: - raise TypeError('Only neural network classifiers and pipeline classifiers are supported') + raise TypeError( + "Only neural network classifiers and pipeline classifiers are supported" + ) input_shape = operator.inputs[0].type.shape if len(operator.inputs[0].type.shape) != 2: # Calculate the shape attribute of ONNX Reshape - if input_shape[0] != 'None': + if input_shape[0] != "None": N = input_shape[0] else: - N = -1 # -1 means that this dimension is automatically determined in runtime and unknown in conversion time + # -1 means that this dimension is automatically + # determined in runtime and unknown in conversion time + N = -1 if all(isinstance(i, numbers.Integral) for i in input_shape[1:]): C = 1 for i in input_shape[1:]: C *= int(i) else: - C = -1 # -1 means that this dimension is automatically determined in runtime and unknown in conversion time + # -1 means that this dimension is automatically determined + # in runtime and unknown in conversion time + C = -1 - # ZipMap in ONNX only accepts [C] and [N, C] inputs. In cases of [N, C, 1, 1], we reshape the probability tensor + # ZipMap in ONNX only accepts [C] and [N, C] inputs. + # In cases of [N, C, 1, 1], we reshape the probability tensor # into [N, C] before feeding it into ZipMap. - buffer_name = scope.get_unique_variable_name('buffer') - apply_reshape(scope, operator.inputs[0].full_name, buffer_name, container, desired_shape=[N, C]) + buffer_name = scope.get_unique_variable_name("buffer") + apply_reshape( + scope, + operator.inputs[0].full_name, + buffer_name, + container, + desired_shape=[N, C], + ) else: buffer_name = operator.inputs[0].full_name - container.add_node('ZipMap', buffer_name, operator.outputs[0].full_name, - op_domain='ai.onnx.ml', **attrs) + container.add_node( + "ZipMap", + buffer_name, + operator.outputs[0].full_name, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('tensorToProbabilityMap', convert_tensor_to_probability_map) +register_converter("tensorToProbabilityMap", convert_tensor_to_probability_map) diff --git a/onnxmltools/convert/coreml/operator_converters/TreeEnsemble.py b/onnxmltools/convert/coreml/operator_converters/TreeEnsemble.py index 33e6ab96..db52cc01 100644 --- a/onnxmltools/convert/coreml/operator_converters/TreeEnsemble.py +++ b/onnxmltools/convert/coreml/operator_converters/TreeEnsemble.py @@ -3,73 +3,116 @@ from ...common._registration import register_converter COREML_TREE_NODE_BEHAVIOR_TO_ONNX_TREE_NODE_MODE = { - 0: 'BRANCH_LEQ', - 1: 'BRANCH_LT', - 2: 'BRANCH_GTE', - 3: 'BRANCH_GT', - 4: 'BRANCH_EQ', - 5: 'BRANCH_NEQ', - 6: 'LEAF' + 0: "BRANCH_LEQ", + 1: "BRANCH_LT", + 2: "BRANCH_GTE", + 3: "BRANCH_GT", + 4: "BRANCH_EQ", + 5: "BRANCH_NEQ", + 6: "LEAF", } COREML_TREE_POST_TRANSFORM_TO_ONNX_TREE_POST_TRANSFORM = { - 0: 'NONE', - 1: 'SOFTMAX', - 2: 'LOGISTIC', - 3: 'SOFTMAX_ZERO' + 0: "NONE", + 1: "SOFTMAX", + 2: "LOGISTIC", + 3: "SOFTMAX_ZERO", } def get_onnx_tree_mode(cm_tree_behavior): if cm_tree_behavior in COREML_TREE_NODE_BEHAVIOR_TO_ONNX_TREE_NODE_MODE: return COREML_TREE_NODE_BEHAVIOR_TO_ONNX_TREE_NODE_MODE[cm_tree_behavior] - raise ValueError('CoreML tree node behavior not supported {0}'.format(cm_tree_behavior)) + raise ValueError( + "CoreML tree node behavior not supported {0}".format(cm_tree_behavior) + ) def get_onnx_tree_post_transform(cm_tree_post_transform): if cm_tree_post_transform in COREML_TREE_POST_TRANSFORM_TO_ONNX_TREE_POST_TRANSFORM: - return COREML_TREE_POST_TRANSFORM_TO_ONNX_TREE_POST_TRANSFORM[cm_tree_post_transform] - raise ValueError('CoreML tree post transform not supported {0}'.format(cm_tree_post_transform)) + return COREML_TREE_POST_TRANSFORM_TO_ONNX_TREE_POST_TRANSFORM[ + cm_tree_post_transform + ] + raise ValueError( + "CoreML tree post transform not supported {0}".format(cm_tree_post_transform) + ) def convert_tree_ensemble_model(scope, operator, container): raw_model = operator.raw_operator - attrs = {'name': operator.full_name} - if raw_model.WhichOneof('Type') == 'treeEnsembleClassifier': - op_type = 'TreeEnsembleClassifier' - prefix = 'class' + attrs = {"name": operator.full_name} + if raw_model.WhichOneof("Type") == "treeEnsembleClassifier": + op_type = "TreeEnsembleClassifier" + prefix = "class" nodes = raw_model.treeEnsembleClassifier.treeEnsemble.nodes - attrs['base_values'] = raw_model.treeEnsembleClassifier.treeEnsemble.basePredictionValue - attrs['post_transform'] = get_onnx_tree_post_transform(raw_model.treeEnsembleClassifier.postEvaluationTransform) - zipmap_attrs = {'name': scope.get_unique_operator_name('ZipMap')} - if raw_model.treeEnsembleClassifier.WhichOneof('ClassLabels') == 'int64ClassLabels': - class_labels = list(int(i) for i in raw_model.treeEnsembleClassifier.int64ClassLabels.vector) - attrs['classlabels_int64s'] = class_labels - zipmap_attrs['classlabels_int64s'] = class_labels + attrs[ + "base_values" + ] = raw_model.treeEnsembleClassifier.treeEnsemble.basePredictionValue + attrs["post_transform"] = get_onnx_tree_post_transform( + raw_model.treeEnsembleClassifier.postEvaluationTransform + ) + zipmap_attrs = {"name": scope.get_unique_operator_name("ZipMap")} + if ( + raw_model.treeEnsembleClassifier.WhichOneof("ClassLabels") + == "int64ClassLabels" + ): + class_labels = list( + int(i) for i in raw_model.treeEnsembleClassifier.int64ClassLabels.vector + ) + attrs["classlabels_int64s"] = class_labels + zipmap_attrs["classlabels_int64s"] = class_labels else: - class_labels = list(s.encode('utf-8') for s in raw_model.treeEnsembleClassifier.stringClassLabels.vector) - attrs['classlabels_strings'] = class_labels - zipmap_attrs['classlabels_strings'] = class_labels - elif raw_model.WhichOneof('Type') == 'treeEnsembleRegressor': - op_type = 'TreeEnsembleRegressor' - prefix = 'target' + class_labels = list( + s.encode("utf-8") + for s in raw_model.treeEnsembleClassifier.stringClassLabels.vector + ) + attrs["classlabels_strings"] = class_labels + zipmap_attrs["classlabels_strings"] = class_labels + elif raw_model.WhichOneof("Type") == "treeEnsembleRegressor": + op_type = "TreeEnsembleRegressor" + prefix = "target" nodes = raw_model.treeEnsembleRegressor.treeEnsemble.nodes - attrs['base_values'] = raw_model.treeEnsembleRegressor.treeEnsemble.basePredictionValue - attrs['n_targets'] = raw_model.treeEnsembleRegressor.treeEnsemble.numPredictionDimensions - attrs['post_transform'] = get_onnx_tree_post_transform(raw_model.treeEnsembleRegressor.postEvaluationTransform) + attrs[ + "base_values" + ] = raw_model.treeEnsembleRegressor.treeEnsemble.basePredictionValue + attrs[ + "n_targets" + ] = raw_model.treeEnsembleRegressor.treeEnsemble.numPredictionDimensions + attrs["post_transform"] = get_onnx_tree_post_transform( + raw_model.treeEnsembleRegressor.postEvaluationTransform + ) else: - raise ValueError('Unknown tree model type') + raise ValueError("Unknown tree model type") - leaf_treeids = [node.treeId for node in nodes if 6 == node.nodeBehavior for _ in node.evaluationInfo] - leaf_nodeids = [node.nodeId for node in nodes if 6 == node.nodeBehavior for _ in node.evaluationInfo] - leaf_ids = [weight.evaluationIndex for node in nodes if 6 == node.nodeBehavior for weight in node.evaluationInfo] + leaf_treeids = [ + node.treeId + for node in nodes + if 6 == node.nodeBehavior + for _ in node.evaluationInfo + ] + leaf_nodeids = [ + node.nodeId + for node in nodes + if 6 == node.nodeBehavior + for _ in node.evaluationInfo + ] + leaf_ids = [ + weight.evaluationIndex + for node in nodes + if 6 == node.nodeBehavior + for weight in node.evaluationInfo + ] - leaf_weights = [weight.evaluationValue for node in nodes if 6 == node.nodeBehavior for weight in - node.evaluationInfo] + leaf_weights = [ + weight.evaluationValue + for node in nodes + if 6 == node.nodeBehavior + for weight in node.evaluationInfo + ] - assert (len(leaf_ids) == len(leaf_weights)) - assert (len(leaf_weights) == len(leaf_nodeids)) - assert (len(leaf_nodeids) == len(leaf_treeids)) + assert len(leaf_ids) == len(leaf_weights) + assert len(leaf_weights) == len(leaf_nodeids) + assert len(leaf_nodeids) == len(leaf_treeids) nodes_nodeids = [x.nodeId for x in nodes] nodes_treeids = [x.treeId for x in nodes] @@ -81,23 +124,23 @@ def convert_tree_ensemble_model(scope, operator, container): nodes_hitrates = [float(x.relativeHitRate) for x in nodes] nodes_modes = [get_onnx_tree_mode(x.nodeBehavior) for x in nodes] - attrs['nodes_treeids'] = nodes_treeids - attrs['nodes_nodeids'] = nodes_nodeids - attrs['nodes_featureids'] = nodes_featureids - attrs['nodes_values'] = nodes_values - attrs['nodes_hitrates'] = nodes_hitrates - attrs['nodes_modes'] = nodes_modes - attrs['nodes_truenodeids'] = nodes_truenodeids - attrs['nodes_falsenodeids'] = nodes_falsenodeids - attrs['nodes_missing_value_tracks_true'] = nodes_missing_value_tracks_true - attrs[prefix + '_treeids'] = leaf_treeids - attrs[prefix + '_nodeids'] = leaf_nodeids - attrs[prefix + '_ids'] = leaf_ids - attrs[prefix + '_weights'] = leaf_weights + attrs["nodes_treeids"] = nodes_treeids + attrs["nodes_nodeids"] = nodes_nodeids + attrs["nodes_featureids"] = nodes_featureids + attrs["nodes_values"] = nodes_values + attrs["nodes_hitrates"] = nodes_hitrates + attrs["nodes_modes"] = nodes_modes + attrs["nodes_truenodeids"] = nodes_truenodeids + attrs["nodes_falsenodeids"] = nodes_falsenodeids + attrs["nodes_missing_value_tracks_true"] = nodes_missing_value_tracks_true + attrs[prefix + "_treeids"] = leaf_treeids + attrs[prefix + "_nodeids"] = leaf_nodeids + attrs[prefix + "_ids"] = leaf_ids + attrs[prefix + "_weights"] = leaf_weights # For regression, we can simply construct a model. For classifier, due to the different representation of # classes' probabilities, we need to add some operators for type conversion. - if raw_model.WhichOneof('Type') == 'treeEnsembleRegressor': + if raw_model.WhichOneof("Type") == "treeEnsembleRegressor": # Create ONNX representation of this operator. If there is only one input, its full topology is # # input features ---> TreeEnsembleRegressor ---> output @@ -113,15 +156,32 @@ def convert_tree_ensemble_model(scope, operator, container): # ... | # input feature N -----' if len(operator.inputs) > 1: - feature_vector_name = scope.get_unique_variable_name('feature_vector') - container.add_node('FeatureVectorizer', operator.input_full_names, feature_vector_name, - op_domain='ai.onnx.ml', name=scope.get_unique_operator_name('FeatureVectorizer'), - inputdimensions=[variable.type.shape[1] for variable in operator.inputs]) - container.add_node(op_type, feature_vector_name, operator.output_full_names, - op_domain='ai.onnx.ml', **attrs) + feature_vector_name = scope.get_unique_variable_name("feature_vector") + container.add_node( + "FeatureVectorizer", + operator.input_full_names, + feature_vector_name, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("FeatureVectorizer"), + inputdimensions=[ + variable.type.shape[1] for variable in operator.inputs + ], + ) + container.add_node( + op_type, + feature_vector_name, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) else: - container.add_node(op_type, operator.input_full_names, operator.output_full_names, - op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) else: # For classifiers, due to the different representation of classes' probabilities, we need to add some # operators for type conversion. It turns out that we have the following topology. @@ -142,10 +202,17 @@ def convert_tree_ensemble_model(scope, operator, container): # Set up input feature(s) if len(operator.inputs) > 1: - feature_vector_name = scope.get_unique_variable_name('feature_vector') - container.add_node('FeatureVectorizer', operator.input_full_names, feature_vector_name, - op_domain='ai.onnx.ml', name=scope.get_unique_operator_name('FeatureVectorizer'), - inputdimensions=[variable.type.shape[1] for variable in operator.inputs]) + feature_vector_name = scope.get_unique_variable_name("feature_vector") + container.add_node( + "FeatureVectorizer", + operator.input_full_names, + feature_vector_name, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("FeatureVectorizer"), + inputdimensions=[ + variable.type.shape[1] for variable in operator.inputs + ], + ) else: feature_vector_name = operator.inputs[0].full_name @@ -154,23 +221,42 @@ def convert_tree_ensemble_model(scope, operator, container): for variable in operator.outputs: if raw_model.description.predictedFeatureName == variable.raw_name: label_output_name = variable.full_name - if raw_model.description.predictedProbabilitiesName != '' and raw_model.description.predictedProbabilitiesName == variable.raw_name: + if ( + raw_model.description.predictedProbabilitiesName != "" + and raw_model.description.predictedProbabilitiesName + == variable.raw_name + ): proba_output_name = variable.full_name - proba_tensor_name = scope.get_unique_variable_name('ProbabilityTensor') + proba_tensor_name = scope.get_unique_variable_name("ProbabilityTensor") if proba_output_name is not None: # Add tree model ONNX node with probability output - container.add_node(op_type, feature_vector_name, [label_output_name, proba_tensor_name], - op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + feature_vector_name, + [label_output_name, proba_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) # Add ZipMap to convert probability tensor into probability map - container.add_node('ZipMap', [proba_tensor_name], [proba_output_name], - op_domain='ai.onnx.ml', **zipmap_attrs) + container.add_node( + "ZipMap", + [proba_tensor_name], + [proba_output_name], + op_domain="ai.onnx.ml", + **zipmap_attrs + ) else: # Add support vector classifier without probability output - container.add_node(op_type, feature_vector_name, [label_output_name, proba_tensor_name], - op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + feature_vector_name, + [label_output_name, proba_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) register_converter("treeEnsembleClassifier", convert_tree_ensemble_model) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Activation.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Activation.py index e4c81910..f2a8637b 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Activation.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Activation.py @@ -2,57 +2,123 @@ import numpy as np -from ....common._apply_operation import apply_elu, apply_hard_sigmoid, apply_leaky_relu, apply_prelu, apply_relu, \ - apply_sigmoid, apply_tanh, apply_affine, apply_parametric_softplus, apply_scaled_tanh, apply_thresholded_relu +from ....common._apply_operation import ( + apply_elu, + apply_hard_sigmoid, + apply_leaky_relu, + apply_prelu, + apply_relu, + apply_sigmoid, + apply_tanh, + apply_affine, + apply_parametric_softplus, + apply_scaled_tanh, + apply_thresholded_relu, +) from ....common._registration import register_converter def convert_activation(scope, operator, container): inputs = [variable.full_name for variable in operator.inputs] outputs = [variable.full_name for variable in operator.outputs] - attrs = {'name': operator.full_name} + attrs = {"name": operator.full_name} params = operator.raw_operator.activation - activation_type = params.WhichOneof('NonlinearityType') + activation_type = params.WhichOneof("NonlinearityType") # Create ONNX objects by high-level APIs such as apply_relu(...) if possible. # Otherwise, we use low-level APIs such as add_node(...) - if activation_type == 'leakyReLU': - apply_leaky_relu(scope, inputs, outputs, container, operator_name=attrs['name'], alpha=params.leakyReLU.alpha) - elif activation_type == 'ReLU': - apply_relu(scope, inputs, outputs, container, operator_name=attrs['name']) - elif activation_type == 'PReLU': - apply_prelu(scope, inputs[0], outputs, container, operator_name=attrs['name'], slope=np.asarray([params.PReLU.alpha.floatValue])) - elif activation_type == 'ELU': - apply_elu(scope, inputs, outputs, container, operator_name=attrs['name'], alpha=params.ELU.alpha) - elif activation_type == 'tanh': - apply_tanh(scope, inputs, outputs, container, operator_name=attrs['name']) - elif activation_type == 'sigmoid': - apply_sigmoid(scope, inputs, outputs, container, operator_name=attrs['name']) - elif activation_type == 'sigmoidHard': - apply_hard_sigmoid(scope, inputs, outputs, container, operator_name=attrs['name'], - alpha=params.sigmoidHard.alpha, beta=params.sigmoidHard.beta) - elif activation_type == 'linear': - apply_affine(scope, inputs[0], outputs[0], container, operator_name=attrs['name'], - alpha=params.linear.alpha, beta=params.linear.beta) - elif activation_type == 'parametricSoftplus': - apply_parametric_softplus(scope, inputs[0], outputs[0], container, operator_name=attrs['name'], - alpha=params.parametricSoftplus.alpha.floatValue, beta=params.parametricSoftplus.beta.floatValue) - elif activation_type =='scaledTanh': - apply_scaled_tanh(scope, inputs[0], outputs[0], container, operator_name=attrs['name'], - alpha=params.scaledTanh.alpha, beta=params.scaledTanh.beta) - elif activation_type == 'thresholdedReLU': - apply_thresholded_relu(scope, inputs, outputs, container, operator_name=attrs['name'], - alpha=params.thresholdedReLU.alpha) + if activation_type == "leakyReLU": + apply_leaky_relu( + scope, + inputs, + outputs, + container, + operator_name=attrs["name"], + alpha=params.leakyReLU.alpha, + ) + elif activation_type == "ReLU": + apply_relu(scope, inputs, outputs, container, operator_name=attrs["name"]) + elif activation_type == "PReLU": + apply_prelu( + scope, + inputs[0], + outputs, + container, + operator_name=attrs["name"], + slope=np.asarray([params.PReLU.alpha.floatValue]), + ) + elif activation_type == "ELU": + apply_elu( + scope, + inputs, + outputs, + container, + operator_name=attrs["name"], + alpha=params.ELU.alpha, + ) + elif activation_type == "tanh": + apply_tanh(scope, inputs, outputs, container, operator_name=attrs["name"]) + elif activation_type == "sigmoid": + apply_sigmoid(scope, inputs, outputs, container, operator_name=attrs["name"]) + elif activation_type == "sigmoidHard": + apply_hard_sigmoid( + scope, + inputs, + outputs, + container, + operator_name=attrs["name"], + alpha=params.sigmoidHard.alpha, + beta=params.sigmoidHard.beta, + ) + elif activation_type == "linear": + apply_affine( + scope, + inputs[0], + outputs[0], + container, + operator_name=attrs["name"], + alpha=params.linear.alpha, + beta=params.linear.beta, + ) + elif activation_type == "parametricSoftplus": + apply_parametric_softplus( + scope, + inputs[0], + outputs[0], + container, + operator_name=attrs["name"], + alpha=params.parametricSoftplus.alpha.floatValue, + beta=params.parametricSoftplus.beta.floatValue, + ) + elif activation_type == "scaledTanh": + apply_scaled_tanh( + scope, + inputs[0], + outputs[0], + container, + operator_name=attrs["name"], + alpha=params.scaledTanh.alpha, + beta=params.scaledTanh.beta, + ) + elif activation_type == "thresholdedReLU": + apply_thresholded_relu( + scope, + inputs, + outputs, + container, + operator_name=attrs["name"], + alpha=params.thresholdedReLU.alpha, + ) else: - if activation_type == 'softsign': - op_type = 'Softsign' - elif activation_type == 'softplus': - op_type = 'Softplus' + if activation_type == "softsign": + op_type = "Softsign" + elif activation_type == "softplus": + op_type = "Softplus" else: - raise TypeError('Unsupported activation layer {0}'.format(activation_type)) + raise TypeError("Unsupported activation layer {0}".format(activation_type)) container.add_node(op_type, inputs, outputs, **attrs) -register_converter('activation', convert_activation) +register_converter("activation", convert_activation) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Add.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Add.py index 8293f86f..a4af7669 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Add.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Add.py @@ -7,12 +7,23 @@ def convert_add(scope, operator, container): if len(operator.input_full_names) == 1: - scaler_name = scope.get_unique_variable_name(operator.full_name + '_B') - container.add_initializer(scaler_name, onnx_proto.TensorProto.FLOAT, [], [operator.raw_operator.add.alpha]) + scaler_name = scope.get_unique_variable_name(operator.full_name + "_B") + container.add_initializer( + scaler_name, + onnx_proto.TensorProto.FLOAT, + [], + [operator.raw_operator.add.alpha], + ) inputs = [operator.inputs[0].full_name, scaler_name] broadcast = 1 - apply_add(scope, inputs, operator.output_full_names, container, operator_name=operator.full_name, - broadcast=broadcast) + apply_add( + scope, + inputs, + operator.output_full_names, + container, + operator_name=operator.full_name, + broadcast=broadcast, + ) else: inputs = operator.input_full_names @@ -24,27 +35,48 @@ def convert_add(scope, operator, container): else: broadcast = 0 - apply_add(scope, [left_tensor, right_tensor], operator.outputs[0].full_name, container, - operator_name=operator.full_name, broadcast=broadcast) + apply_add( + scope, + [left_tensor, right_tensor], + operator.outputs[0].full_name, + container, + operator_name=operator.full_name, + broadcast=broadcast, + ) else: left_tensor = operator.inputs[0].full_name right_tensor = operator.inputs[1].full_name # Sum up the first two inputs - intermediate_tensor_name = scope.get_unique_variable_name('buffer_tensor') - apply_add(scope, [left_tensor, right_tensor], intermediate_tensor_name, container, - operator_name=operator.full_name, broadcast=1) + intermediate_tensor_name = scope.get_unique_variable_name("buffer_tensor") + apply_add( + scope, + [left_tensor, right_tensor], + intermediate_tensor_name, + container, + operator_name=operator.full_name, + broadcast=1, + ) - # Accumulate other inputs onto intermediate tensors. Note that we may use the original operator's output as + # Accumulate other inputs onto intermediate tensors. + # Note that we may use the original operator's output as # the last intermediate tensor. for i in range(2, len(inputs)): left_tensor = intermediate_tensor_name right_tensor = inputs[i].full_name if i != len(inputs) - 1: - intermediate_tensor_name = scope.get_unique_variable_name('buffer_tensor') + intermediate_tensor_name = scope.get_unique_variable_name( + "buffer_tensor" + ) else: intermediate_tensor_name = operator.outputs[0].full_name - apply_add(scope, [left_tensor, right_tensor], intermediate_tensor_name, container, broadcast=1) + apply_add( + scope, + [left_tensor, right_tensor], + intermediate_tensor_name, + container, + broadcast=1, + ) -register_converter('add', convert_add) +register_converter("add", convert_add) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Average.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Average.py index 60ce087e..700143cc 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Average.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Average.py @@ -5,8 +5,13 @@ def convert_average(scope, operator, container): - apply_mean(scope, operator.input_full_names, operator.output_full_names, container, - operator_name=operator.full_name) + apply_mean( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) -register_converter('average', convert_average) +register_converter("average", convert_average) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/BatchNorm.py b/onnxmltools/convert/coreml/operator_converters/neural_network/BatchNorm.py index 630523b4..48d81f3c 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/BatchNorm.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/BatchNorm.py @@ -9,61 +9,103 @@ def convert_batch_normalization(scope, operator, container): params = operator.raw_operator.batchnorm if params.instanceNormalization and not params.computeMeanVar: - raise ValueError('It is impossible to do instance normalization without re-computing mean and variance') + raise ValueError( + "It is impossible to do instance normalization " + "without re-computing mean and variance" + ) if params.instanceNormalization and params.computeMeanVar: - op_type = 'InstanceNormalization' + op_type = "InstanceNormalization" else: - op_type = 'BatchNormalization' + op_type = "BatchNormalization" inputs = [operator.inputs[0].full_name] outputs = [operator.outputs[0].full_name] - scale_tensor_name = scope.get_unique_variable_name(op_type + '_scale') - container.add_initializer(scale_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels], - params.gamma.floatValue) + scale_tensor_name = scope.get_unique_variable_name(op_type + "_scale") + container.add_initializer( + scale_tensor_name, + onnx_proto.TensorProto.FLOAT, + [params.channels], + params.gamma.floatValue, + ) inputs.append(scale_tensor_name) - bias_tensor_name = scope.get_unique_variable_name(op_type + '_B') - container.add_initializer(bias_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels], params.beta.floatValue) + bias_tensor_name = scope.get_unique_variable_name(op_type + "_B") + container.add_initializer( + bias_tensor_name, + onnx_proto.TensorProto.FLOAT, + [params.channels], + params.beta.floatValue, + ) inputs.append(bias_tensor_name) epsilon = params.epsilon - if op_type == 'BatchNormalization': - mean_tensor_name = scope.get_unique_variable_name(op_type + '_mean') - container.add_initializer(mean_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels], - params.mean.floatValue) + if op_type == "BatchNormalization": + mean_tensor_name = scope.get_unique_variable_name(op_type + "_mean") + container.add_initializer( + mean_tensor_name, + onnx_proto.TensorProto.FLOAT, + [params.channels], + params.mean.floatValue, + ) inputs.append(mean_tensor_name) - variance_tensor_name = scope.get_unique_variable_name(op_type + '_variance') - container.add_initializer(variance_tensor_name, onnx_proto.TensorProto.FLOAT, [params.channels], - params.variance.floatValue) + variance_tensor_name = scope.get_unique_variable_name(op_type + "_variance") + container.add_initializer( + variance_tensor_name, + onnx_proto.TensorProto.FLOAT, + [params.channels], + params.variance.floatValue, + ) inputs.append(variance_tensor_name) - momentum = 0. + momentum = 0.0 spatial = 1 # True if not params.instanceNormalization and params.computeMeanVar: - # In this case, we apply batch normalization and adjust the statistics stored according the the batch + # In this case, we apply batch normalization and adjust + # the statistics stored according the the batch # being processed. - # To update "mean" and "var," we put their updated results back to the associated input tensors. + # To update "mean" and "var," we put their updated results + # back to the associated input tensors. outputs += inputs[1:3] - # We also allocate two extra output buffers to store some intermediate results, but they are not used + # We also allocate two extra output buffers to store some + # intermediate results, but they are not used # in CoreML model. - outputs.append(scope.get_unique_variable_name('saved_mean')) - outputs.append(scope.get_unique_variable_name('saved_var')) + outputs.append(scope.get_unique_variable_name("saved_mean")) + outputs.append(scope.get_unique_variable_name("saved_var")) # We choose "training" mode because some variables need to be updated. is_test = 0 # False elif not params.instanceNormalization and not params.computeMeanVar: - # In this case, batch normalization is applied without updating mean, variance, etc. according to - # the batches being processed. It means this operator works under testing model. Because there is no - # variable update, we don't need to specify extra inputs and outputs like in previous code block. + # In this case, batch normalization is applied without + # updating mean, variance, etc. according to + # the batches being processed. It means this operator + # works under testing model. Because there is no + # variable update, we don't need to specify extra + # inputs and outputs like in previous code block. is_test = 1 # True else: - raise ValueError('Unsupported operation mode') + raise ValueError("Unsupported operation mode") - apply_batch_norm(scope, inputs, outputs, container, operator_name=operator.full_name, epsilon=epsilon, - is_test=is_test, momentum=momentum, spatial=spatial) + apply_batch_norm( + scope, + inputs, + outputs, + container, + operator_name=operator.full_name, + epsilon=epsilon, + is_test=is_test, + momentum=momentum, + spatial=spatial, + ) else: - apply_instance_norm(scope, inputs, outputs, container, operator_name=operator.full_name, epsilon=epsilon) + apply_instance_norm( + scope, + inputs, + outputs, + container, + operator_name=operator.full_name, + epsilon=epsilon, + ) -register_converter('batchnorm', convert_batch_normalization) +register_converter("batchnorm", convert_batch_normalization) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Bias.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Bias.py index a50e7ddd..6841991b 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Bias.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Bias.py @@ -7,22 +7,33 @@ def convert_bias(scope, operator, container): - # Feed the input (which we are going to add a bias onto) into Add operator. Its shape is [C, H, W] in CoreML but + # Feed the input (which we are going to add a bias onto) + # into Add operator. Its shape is [C, H, W] in CoreML but # [N, C, H, W] in ONNX. params = operator.raw_operator.bias # Adjust CoreML's bias shape and find a proper axis for broadcasting axis, shape = deduce_broadcast_axis_and_shape(container.target_opset, params.shape) - # No matter what shape it is, we need "broadcast" on because input shape is 4-D while bias is at most 3-D. + # No matter what shape it is, we need "broadcast" on + # because input shape is 4-D while bias is at most 3-D. broadcast = 1 # True # Create bias vector as an ONNX tensor - bias_tensor_name = scope.get_unique_variable_name(operator.full_name + '_B') - container.add_initializer(bias_tensor_name, onnx_proto.TensorProto.FLOAT, shape, params.bias.floatValue) - - apply_add(scope, [operator.inputs[0].full_name, bias_tensor_name], operator.output_full_names, container, - operator_name=operator.full_name, axis=axis, broadcast=broadcast) - - -register_converter('bias', convert_bias) + bias_tensor_name = scope.get_unique_variable_name(operator.full_name + "_B") + container.add_initializer( + bias_tensor_name, onnx_proto.TensorProto.FLOAT, shape, params.bias.floatValue + ) + + apply_add( + scope, + [operator.inputs[0].full_name, bias_tensor_name], + operator.output_full_names, + container, + operator_name=operator.full_name, + axis=axis, + broadcast=broadcast, + ) + + +register_converter("bias", convert_bias) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/BidirectionalLSTM.py b/onnxmltools/convert/coreml/operator_converters/neural_network/BidirectionalLSTM.py index 03e6252f..723805eb 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/BidirectionalLSTM.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/BidirectionalLSTM.py @@ -170,43 +170,72 @@ def convert_bidirectional_lstm(scope, operator, container): input_size = params.inputVectorSize hidden_size = params.outputVectorSize - lstm_op_name = scope.get_unique_operator_name('LSTM') - lstm_attrs = {'name': lstm_op_name} + lstm_op_name = scope.get_unique_operator_name("LSTM") + lstm_attrs = {"name": lstm_op_name} lstm_inputs = [] lstm_outputs = [] - lstm_x_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_X_reshape') + lstm_x_reshape_name = scope.get_unique_variable_name(lstm_op_name + "_X_reshape") - apply_reshape(scope, operator.inputs[0].full_name, lstm_x_reshape_name, container, - desired_shape=[-1, 1, input_size]) + apply_reshape( + scope, + operator.inputs[0].full_name, + lstm_x_reshape_name, + container, + desired_shape=[-1, 1, input_size], + ) lstm_inputs.append(lstm_x_reshape_name) # Handle LSTM's weight matrices - matrices_w_forward = np.concatenate([lstm_weights[0].inputGateWeightMatrix.floatValue, - lstm_weights[0].outputGateWeightMatrix.floatValue, - lstm_weights[0].forgetGateWeightMatrix.floatValue, - lstm_weights[0].blockInputWeightMatrix.floatValue]) - matrices_w_backward = np.concatenate([lstm_weights[1].inputGateWeightMatrix.floatValue, - lstm_weights[1].outputGateWeightMatrix.floatValue, - lstm_weights[1].forgetGateWeightMatrix.floatValue, - lstm_weights[1].blockInputWeightMatrix.floatValue]) - matrices_w_name = scope.get_unique_variable_name(lstm_op_name + '_W') - container.add_initializer(matrices_w_name, onnx_proto.TensorProto.FLOAT, [2, 4 * hidden_size, input_size], - np.concatenate([matrices_w_forward, matrices_w_backward])) + matrices_w_forward = np.concatenate( + [ + lstm_weights[0].inputGateWeightMatrix.floatValue, + lstm_weights[0].outputGateWeightMatrix.floatValue, + lstm_weights[0].forgetGateWeightMatrix.floatValue, + lstm_weights[0].blockInputWeightMatrix.floatValue, + ] + ) + matrices_w_backward = np.concatenate( + [ + lstm_weights[1].inputGateWeightMatrix.floatValue, + lstm_weights[1].outputGateWeightMatrix.floatValue, + lstm_weights[1].forgetGateWeightMatrix.floatValue, + lstm_weights[1].blockInputWeightMatrix.floatValue, + ] + ) + matrices_w_name = scope.get_unique_variable_name(lstm_op_name + "_W") + container.add_initializer( + matrices_w_name, + onnx_proto.TensorProto.FLOAT, + [2, 4 * hidden_size, input_size], + np.concatenate([matrices_w_forward, matrices_w_backward]), + ) lstm_inputs.append(matrices_w_name) # Handle LSTM's recursion matrices - matrices_r_forward = np.concatenate([lstm_weights[0].inputGateRecursionMatrix.floatValue, - lstm_weights[0].outputGateRecursionMatrix.floatValue, - lstm_weights[0].forgetGateRecursionMatrix.floatValue, - lstm_weights[0].blockInputRecursionMatrix.floatValue]) - matrices_r_backward = np.concatenate([lstm_weights[1].inputGateRecursionMatrix.floatValue, - lstm_weights[1].outputGateRecursionMatrix.floatValue, - lstm_weights[1].forgetGateRecursionMatrix.floatValue, - lstm_weights[1].blockInputRecursionMatrix.floatValue]) - matrices_r_name = scope.get_unique_variable_name(lstm_op_name + '_R') - container.add_initializer(matrices_r_name, onnx_proto.TensorProto.FLOAT, [2, 4 * hidden_size, hidden_size], - np.concatenate([matrices_r_forward, matrices_r_backward])) + matrices_r_forward = np.concatenate( + [ + lstm_weights[0].inputGateRecursionMatrix.floatValue, + lstm_weights[0].outputGateRecursionMatrix.floatValue, + lstm_weights[0].forgetGateRecursionMatrix.floatValue, + lstm_weights[0].blockInputRecursionMatrix.floatValue, + ] + ) + matrices_r_backward = np.concatenate( + [ + lstm_weights[1].inputGateRecursionMatrix.floatValue, + lstm_weights[1].outputGateRecursionMatrix.floatValue, + lstm_weights[1].forgetGateRecursionMatrix.floatValue, + lstm_weights[1].blockInputRecursionMatrix.floatValue, + ] + ) + matrices_r_name = scope.get_unique_variable_name(lstm_op_name + "_R") + container.add_initializer( + matrices_r_name, + onnx_proto.TensorProto.FLOAT, + [2, 4 * hidden_size, hidden_size], + np.concatenate([matrices_r_forward, matrices_r_backward]), + ) lstm_inputs.append(matrices_r_name) # Handle bias vectors @@ -225,72 +254,129 @@ def convert_bidirectional_lstm(scope, operator, container): # but it's not correct as CoreML has added 1 into those bias vectors. pass if lstm_params.hasBiasVectors or lstm_params.forgetBias: - vectors_b_name = scope.get_unique_variable_name(lstm_op_name + '_B') - container.add_initializer(vectors_b_name, onnx_proto.TensorProto.FLOAT, - [2, 8 * hidden_size], vectors_b.flatten()) + vectors_b_name = scope.get_unique_variable_name(lstm_op_name + "_B") + container.add_initializer( + vectors_b_name, + onnx_proto.TensorProto.FLOAT, + [2, 8 * hidden_size], + vectors_b.flatten(), + ) lstm_inputs.append(vectors_b_name) else: - lstm_inputs.append('') + lstm_inputs.append("") - # Due to the position sensitivity in ONNX argument parsing, we add an empty string for the non-existing + # Due to the position sensitivity in ONNX argument parsing, + # we add an empty string for the non-existing # sequence length - lstm_inputs.append('') + lstm_inputs.append("") # Handle initial hidden state if necessary if len(operator.inputs) > 1: - lstm_h_init_name = scope.get_unique_variable_name(lstm_op_name + '_h_init') - apply_concat(scope, [operator.inputs[1].full_name, operator.inputs[3].full_name], lstm_h_init_name, - container, axis=0) - - lstm_h_init_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_h_init_reshape') - apply_reshape(scope, lstm_h_init_name, lstm_h_init_reshape_name, container, desired_shape=[2, 1, hidden_size]) - - # Add zero initializers to forward and backward initial hidden states so that they become optional - container.add_initializer(operator.inputs[1].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[1].type.shape, - np.zeros(shape=operator.inputs[1].type.shape).flatten()) - container.add_initializer(operator.inputs[3].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[3].type.shape, - np.zeros(shape=operator.inputs[3].type.shape).flatten()) + lstm_h_init_name = scope.get_unique_variable_name(lstm_op_name + "_h_init") + apply_concat( + scope, + [operator.inputs[1].full_name, operator.inputs[3].full_name], + lstm_h_init_name, + container, + axis=0, + ) + + lstm_h_init_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_h_init_reshape" + ) + apply_reshape( + scope, + lstm_h_init_name, + lstm_h_init_reshape_name, + container, + desired_shape=[2, 1, hidden_size], + ) + + # Add zero initializers to forward and backward initial + # hidden states so that they become optional + container.add_initializer( + operator.inputs[1].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[1].type.shape, + np.zeros(shape=operator.inputs[1].type.shape).flatten(), + ) + container.add_initializer( + operator.inputs[3].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[3].type.shape, + np.zeros(shape=operator.inputs[3].type.shape).flatten(), + ) lstm_inputs.append(lstm_h_init_reshape_name) else: - lstm_inputs.append('') + lstm_inputs.append("") # Handle initial cell state if needed if len(operator.inputs) > 2: - lstm_c_init_name = scope.get_unique_variable_name(lstm_op_name + '_c_init') - apply_concat(scope, [operator.inputs[2].full_name, operator.inputs[4].full_name], lstm_c_init_name, container, - axis=0) + lstm_c_init_name = scope.get_unique_variable_name(lstm_op_name + "_c_init") + apply_concat( + scope, + [operator.inputs[2].full_name, operator.inputs[4].full_name], + lstm_c_init_name, + container, + axis=0, + ) # Reshape the Cell state so that ONNX LSTM can accept it - lstm_c_init_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_c_init_reshape') - apply_reshape(scope, lstm_c_init_name, lstm_c_init_reshape_name, container, desired_shape=[2, 1, hidden_size]) + lstm_c_init_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_c_init_reshape" + ) + apply_reshape( + scope, + lstm_c_init_name, + lstm_c_init_reshape_name, + container, + desired_shape=[2, 1, hidden_size], + ) lstm_inputs.append(lstm_c_init_reshape_name) - # Add zero initializers to forward and backward initial cell states so that they become optional - container.add_initializer(operator.inputs[2].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[2].type.shape, - np.zeros(shape=operator.inputs[2].type.shape).flatten()) - container.add_initializer(operator.inputs[4].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[4].type.shape, - np.zeros(shape=operator.inputs[4].type.shape).flatten()) + # Add zero initializers to forward and backward initial + # cell states so that they become optional + container.add_initializer( + operator.inputs[2].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[2].type.shape, + np.zeros(shape=operator.inputs[2].type.shape).flatten(), + ) + container.add_initializer( + operator.inputs[4].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[4].type.shape, + np.zeros(shape=operator.inputs[4].type.shape).flatten(), + ) else: - lstm_inputs.append('') + lstm_inputs.append("") # Handle peephole vectors if necessary if lstm_params.hasPeepholeVectors: - p_forward = np.concatenate([lstm_weights[0].inputGatePeepholeVector.floatValue, - lstm_weights[0].outputGatePeepholeVector.floatValue, - lstm_weights[0].forgetGatePeepholeVector.floatValue]) - p_backward = np.concatenate([lstm_weights[1].inputGatePeepholeVector.floatValue, - lstm_weights[1].outputGatePeepholeVector.floatValue, - lstm_weights[1].forgetGatePeepholeVector.floatValue]) - p_name = scope.get_unique_variable_name(lstm_op_name + '_P') - container.add_initializer(p_name, onnx_proto.TensorProto.FLOAT, - [2, 3 * hidden_size], np.concatenate([p_forward, p_backward])) + p_forward = np.concatenate( + [ + lstm_weights[0].inputGatePeepholeVector.floatValue, + lstm_weights[0].outputGatePeepholeVector.floatValue, + lstm_weights[0].forgetGatePeepholeVector.floatValue, + ] + ) + p_backward = np.concatenate( + [ + lstm_weights[1].inputGatePeepholeVector.floatValue, + lstm_weights[1].outputGatePeepholeVector.floatValue, + lstm_weights[1].forgetGatePeepholeVector.floatValue, + ] + ) + p_name = scope.get_unique_variable_name(lstm_op_name + "_P") + container.add_initializer( + p_name, + onnx_proto.TensorProto.FLOAT, + [2, 3 * hidden_size], + np.concatenate([p_forward, p_backward]), + ) lstm_inputs.append(p_name) else: - lstm_inputs.append('') + lstm_inputs.append("") # Parse activation functions and add them into ONNX LSTM's attribute dictionary activation_types = [] @@ -298,78 +384,135 @@ def convert_bidirectional_lstm(scope, operator, container): betas = [] for activation in params.activationsForwardLSTM: activation_type, alpha, beta = extract_rnn_activation_info(activation) - activation_types.append(activation_type.encode('utf-8')) + activation_types.append(activation_type.encode("utf-8")) if alpha is not None: alphas.append(alpha) if beta is not None: betas.append(beta) for activation in params.activationsBackwardLSTM: activation_type, alpha, beta = extract_rnn_activation_info(activation) - activation_types.append(activation_type.encode('utf-8')) + activation_types.append(activation_type.encode("utf-8")) if alpha is not None: alphas.append(alpha) if beta is not None: betas.append(beta) - lstm_attrs['activations'] = activation_types + lstm_attrs["activations"] = activation_types if alphas: - lstm_attrs['activation_alpha'] = alphas + lstm_attrs["activation_alpha"] = alphas if betas: - lstm_attrs['activation_beta'] = betas + lstm_attrs["activation_beta"] = betas # Add more attributes - lstm_attrs['direction'] = 'bidirectional' - lstm_attrs['hidden_size'] = hidden_size - lstm_attrs['clip'] = float(lstm_params.cellClipThreshold) - lstm_attrs['input_forget'] = lstm_params.coupledInputAndForgetGate + lstm_attrs["direction"] = "bidirectional" + lstm_attrs["hidden_size"] = hidden_size + lstm_attrs["clip"] = float(lstm_params.cellClipThreshold) + lstm_attrs["input_forget"] = lstm_params.coupledInputAndForgetGate # Set up version-dependent attributes if container.target_opset < 7: - lstm_attrs['output_sequence'] = lstm_params.sequenceOutput + lstm_attrs["output_sequence"] = lstm_params.sequenceOutput op_version = 1 else: op_version = 7 - # Create the major ONNX LSTM operator. We assign a tensor name to each output of LSTM. However, variables can be - # undefined in some cases. For example, when output_sequence=False, the first output is not meaningful. - lstm_y_name = scope.get_unique_variable_name(lstm_op_name + '_Y') - lstm_y_h_name = scope.get_unique_variable_name(lstm_op_name + '_Y_h') - lstm_y_c_name = scope.get_unique_variable_name(lstm_op_name + '_Y_c') + # Create the major ONNX LSTM operator. We assign a tensor name + # to each output of LSTM. However, variables can be + # undefined in some cases. For example, when output_sequence=False, + # the first output is not meaningful. + lstm_y_name = scope.get_unique_variable_name(lstm_op_name + "_Y") + lstm_y_h_name = scope.get_unique_variable_name(lstm_op_name + "_Y_h") + lstm_y_c_name = scope.get_unique_variable_name(lstm_op_name + "_Y_c") lstm_outputs.extend([lstm_y_name, lstm_y_h_name, lstm_y_c_name]) - container.add_node('LSTM', lstm_inputs, lstm_outputs, op_version=op_version, **lstm_attrs) + container.add_node( + "LSTM", lstm_inputs, lstm_outputs, op_version=op_version, **lstm_attrs + ) # Create post-processing operators for converting ONNX LSTM outputs to CoreML ones if lstm_params.sequenceOutput: - apply_reshape(scope, lstm_y_name, operator.outputs[0].full_name, container, desired_shape=[-1, 2 * hidden_size]) + apply_reshape( + scope, + lstm_y_name, + operator.outputs[0].full_name, + container, + desired_shape=[-1, 2 * hidden_size], + ) if len(operator.outputs) > 1: - lstm_y_h_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_Y_h_reshape') - apply_reshape(scope, lstm_y_h_name, lstm_y_h_reshape_name, container, desired_shape=[2, hidden_size]) - - apply_split(scope, lstm_y_h_reshape_name, [operator.outputs[1].full_name, operator.outputs[3].full_name], - container, split=[1, 1], axis=0) + lstm_y_h_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_Y_h_reshape" + ) + apply_reshape( + scope, + lstm_y_h_name, + lstm_y_h_reshape_name, + container, + desired_shape=[2, hidden_size], + ) + + apply_split( + scope, + lstm_y_h_reshape_name, + [operator.outputs[1].full_name, operator.outputs[3].full_name], + container, + split=[1, 1], + axis=0, + ) else: - # Here we ignore ONNX RNN's first output because it's useless. The second output of ONNX LSTM will be used to + # Here we ignore ONNX RNN's first output because it's useless. + # The second output of ONNX LSTM will be used to # generate the first and the second outputs of CoreML LSTM. # Directly reshape ONNX LSTM's 2nd output to CoreML LSTM's 1st output. - apply_reshape(scope, lstm_y_h_name, operator.outputs[0].full_name, container, - desired_shape=[1, 2 * hidden_size]) + apply_reshape( + scope, + lstm_y_h_name, + operator.outputs[0].full_name, + container, + desired_shape=[1, 2 * hidden_size], + ) if len(operator.outputs) > 1: - lstm_y_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_Y_reshape') - apply_reshape(scope, lstm_y_h_name, lstm_y_reshape_name, container, desired_shape=[2, hidden_size]) - - apply_split(scope, lstm_y_reshape_name, [operator.outputs[1].full_name, operator.outputs[3].full_name], - container, split=[1, 1], axis=0) + lstm_y_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_Y_reshape" + ) + apply_reshape( + scope, + lstm_y_h_name, + lstm_y_reshape_name, + container, + desired_shape=[2, hidden_size], + ) + + apply_split( + scope, + lstm_y_reshape_name, + [operator.outputs[1].full_name, operator.outputs[3].full_name], + container, + split=[1, 1], + axis=0, + ) # Output cell state if necessary if len(operator.outputs) > 2: - lstm_y_c_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_Y_c_reshape') - apply_reshape(scope, lstm_y_c_name, lstm_y_c_reshape_name, container, desired_shape=[2, hidden_size]) - - apply_split(scope, lstm_y_c_reshape_name, [operator.outputs[2].full_name, operator.outputs[4].full_name], - container, split=[1, 1], axis=0) - - -register_converter('biDirectionalLSTM', convert_bidirectional_lstm) - + lstm_y_c_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_Y_c_reshape" + ) + apply_reshape( + scope, + lstm_y_c_name, + lstm_y_c_reshape_name, + container, + desired_shape=[2, hidden_size], + ) + + apply_split( + scope, + lstm_y_c_reshape_name, + [operator.outputs[2].full_name, operator.outputs[4].full_name], + container, + split=[1, 1], + axis=0, + ) + + +register_converter("biDirectionalLSTM", convert_bidirectional_lstm) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Concat.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Concat.py index c12643d5..281a90a9 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Concat.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Concat.py @@ -10,8 +10,14 @@ def convert_concat(scope, operator, container): else: axis = 1 - apply_concat(scope, operator.input_full_names, operator.output_full_names, container, - operator_name=operator.full_name, axis=axis) + apply_concat( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator_name=operator.full_name, + axis=axis, + ) -register_converter('concat', convert_concat) +register_converter("concat", convert_concat) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Convolution.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Convolution.py index a92cb8ca..c83ea756 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Convolution.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Convolution.py @@ -8,26 +8,35 @@ def convert_convolution(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import SamePadding params = operator.raw_operator.convolution - op_type = 'ConvTranspose' if params.isDeconvolution else 'Conv' + op_type = "ConvTranspose" if params.isDeconvolution else "Conv" inputs = [operator.inputs[0].full_name] outputs = [operator.outputs[0].full_name] - attrs = {'name': operator.full_name} + attrs = {"name": operator.full_name} n_groups = 1 if params.nGroups == 0 else params.nGroups - shape_w = [params.outputChannels, params.kernelChannels, params.kernelSize[0], params.kernelSize[1]] + shape_w = [ + params.outputChannels, + params.kernelChannels, + params.kernelSize[0], + params.kernelSize[1], + ] if params.isDeconvolution: shape_w[0] = params.kernelChannels shape_w[1] = int(params.outputChannels / n_groups) - name_w = scope.get_unique_variable_name(operator.full_name + '_W') + name_w = scope.get_unique_variable_name(operator.full_name + "_W") inputs.append(name_w) - container.add_initializer(name_w, onnx_proto.TensorProto.FLOAT, shape_w, params.weights.floatValue) + container.add_initializer( + name_w, onnx_proto.TensorProto.FLOAT, shape_w, params.weights.floatValue + ) if params.hasBias: shape_b = [len(params.bias.floatValue)] - name_b = scope.get_unique_variable_name(operator.full_name + '_B') + name_b = scope.get_unique_variable_name(operator.full_name + "_B") inputs.append(name_b) - container.add_initializer(name_b, onnx_proto.TensorProto.FLOAT, shape_b, params.bias.floatValue) + container.add_initializer( + name_b, onnx_proto.TensorProto.FLOAT, shape_b, params.bias.floatValue + ) dilations = [1, 1] if len(params.dilationFactor) > 0: @@ -38,15 +47,15 @@ def convert_convolution(scope, operator, container): strides = [1, 1] if len(params.stride) > 0: strides = params.stride - attrs['dilations'] = dilations - attrs['group'] = n_groups - attrs['kernel_shape'] = kernel_shape - attrs['strides'] = strides + attrs["dilations"] = dilations + attrs["group"] = n_groups + attrs["kernel_shape"] = kernel_shape + attrs["strides"] = strides pads = None auto_pad = None - pad_type = params.WhichOneof('ConvolutionPaddingType') - if pad_type == 'valid': + pad_type = params.WhichOneof("ConvolutionPaddingType") + if pad_type == "valid": if len(params.valid.paddingAmounts.borderAmounts) > 0: pads = [0, 0, 0, 0] pads[0] = params.valid.paddingAmounts.borderAmounts[0].startEdgeSize @@ -56,30 +65,31 @@ def convert_convolution(scope, operator, container): # If padding amounts are all zero, there should be no padding list. if all(pad == 0 for pad in pads): pads = None - auto_pad = 'VALID' + auto_pad = "VALID" else: - auto_pad = 'VALID' - elif pad_type == 'same': + auto_pad = "VALID" + elif pad_type == "same": if params.same.asymmetryMode == SamePadding.BOTTOM_RIGHT_HEAVY: - auto_pad = 'SAME_UPPER' + auto_pad = "SAME_UPPER" elif params.same.asymmetryMode == SamePadding.TOP_LEFT_HEAVY: - auto_pad = 'SAME_LOWER' + auto_pad = "SAME_LOWER" else: - raise ValueError('Unknown asymmetric mode: {}'.format( - params.same.asymmetryMode)) + raise ValueError( + "Unknown asymmetric mode: {}".format(params.same.asymmetryMode) + ) else: - raise ValueError('Unsupported padding mode: {}'.format(pad_type)) + raise ValueError("Unsupported padding mode: {}".format(pad_type)) if params.isDeconvolution and len(params.outputShape) > 0: - attrs['output_shape'] = params.outputShape + attrs["output_shape"] = params.outputShape if pads is not None: - attrs['pads'] = pads + attrs["pads"] = pads if auto_pad is not None: - attrs['auto_pad'] = auto_pad + attrs["auto_pad"] = auto_pad container.add_node(op_type, inputs, outputs, **attrs) -register_converter('convolution', convert_convolution) +register_converter("convolution", convert_convolution) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Crop.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Crop.py index 4e3510f1..e4e4d754 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Crop.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Crop.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -import numpy as np -from .....proto import onnx_proto from ....common._apply_operation import apply_crop_height_width from ....common._registration import register_converter + def convert_crop(scope, operator, container): # Extract number of pixels cropped in CoreML operator. border = operator.raw_operator.crop.cropAmounts.borderAmounts @@ -27,12 +26,17 @@ def convert_crop(scope, operator, container): bottom_border = in_shape[2] - top_border - out_shape[2] # Delegate the selection of ONNX operator to a version-dependent function. - apply_crop_height_width(scope, - operator.input_full_names[0], operator.output_full_names[0], - container, operator_name=operator.full_name, - top_border=top_border, bottom_border=bottom_border, - left_border=left_border, right_border=right_border) - + apply_crop_height_width( + scope, + operator.input_full_names[0], + operator.output_full_names[0], + container, + operator_name=operator.full_name, + top_border=top_border, + bottom_border=bottom_border, + left_border=left_border, + right_border=right_border, + ) -register_converter('crop', convert_crop) +register_converter("crop", convert_crop) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Dot.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Dot.py index 74299e9e..6b9d7e4c 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Dot.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Dot.py @@ -5,36 +5,63 @@ def convert_dot(scope, operator, container): - # To calculate cosine similarity, we first use LpNormalization to make the two input vectors unit-length. - # Then, we calculate element-wise product of the two unit-length vectors. Finally, the similarity is the - # sum of all the product's elements. Notice that we carefully specify the axis of the subsequent operators, + # To calculate cosine similarity, we first use + # LpNormalization to make the two input vectors unit-length. + # Then, we calculate element-wise product of the two + # unit-length vectors. Finally, the similarity is the + # sum of all the product's elements. Notice that + # we carefully specify the axis of the subsequent operators, # so they can work properly with a batch of vectors. if operator.raw_operator.dot.cosineSimilarity: # Normalize the first input and store the result on a temporal variable - intra_variable_name1 = scope.get_unique_variable_name(operator.inputs[0].full_name + '_normalized') - normalizer_name1 = scope.get_unique_operator_name('L2NormNormalizer') - attrs1 = {'name': normalizer_name1, 'p': 2, 'aixs': 1} - container.add_node('LpNormalization', [operator.inputs[0].full_name], [intra_variable_name1], **attrs1) + intra_variable_name1 = scope.get_unique_variable_name( + operator.inputs[0].full_name + "_normalized" + ) + normalizer_name1 = scope.get_unique_operator_name("L2NormNormalizer") + attrs1 = {"name": normalizer_name1, "p": 2, "aixs": 1} + container.add_node( + "LpNormalization", + [operator.inputs[0].full_name], + [intra_variable_name1], + **attrs1 + ) # Normalize the second input and store the result on a temporal variable - intra_variable_name2 = scope.get_unique_variable_name(operator.inputs[1].full_name + '_normalized') - normalizer_name2 = scope.get_unique_operator_name('L2NormNormalizer') - attrs2 = {'name': normalizer_name2, 'p': 2, 'aixs': 1} - container.add_node('LpNormalization', [operator.inputs[1].full_name], [intra_variable_name2], **attrs2) + intra_variable_name2 = scope.get_unique_variable_name( + operator.inputs[1].full_name + "_normalized" + ) + normalizer_name2 = scope.get_unique_operator_name("L2NormNormalizer") + attrs2 = {"name": normalizer_name2, "p": 2, "aixs": 1} + container.add_node( + "LpNormalization", + [operator.inputs[1].full_name], + [intra_variable_name2], + **attrs2 + ) else: # This case is a simple dot product; no normalization is required. intra_variable_name1 = operator.inputs[0].full_name intra_variable_name2 = operator.inputs[1].full_name # Do element-wise product of the two unit-length tensors - product_name = scope.get_unique_variable_name(intra_variable_name1 + '_multiply_' + intra_variable_name2) - apply_mul(scope, [intra_variable_name1, intra_variable_name2], product_name, container, broadcast=0) + product_name = scope.get_unique_variable_name( + intra_variable_name1 + "_multiply_" + intra_variable_name2 + ) + apply_mul( + scope, + [intra_variable_name1, intra_variable_name2], + product_name, + container, + broadcast=0, + ) # Sum up results from different dimensions to get the final cosine similarity - reducer_name = scope.get_unique_operator_name('ReduceSum') - reducer_attrs = {'name': reducer_name, 'axes': [1], 'keepdims': 1} - container.add_node('ReduceSum', [product_name], operator.output_full_names, **reducer_attrs) + reducer_name = scope.get_unique_operator_name("ReduceSum") + reducer_attrs = {"name": reducer_name, "axes": [1], "keepdims": 1} + container.add_node( + "ReduceSum", [product_name], operator.output_full_names, **reducer_attrs + ) -register_converter('dot', convert_dot) +register_converter("dot", convert_dot) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Embed.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Embed.py index ca048756..d5ee3a56 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Embed.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Embed.py @@ -9,39 +9,86 @@ def convert_embedding(scope, operator, container): params = operator.raw_operator.embedding - gather_op_name = scope.get_unique_operator_name('Gather') - gather_attrs = {'name': gather_op_name} + gather_op_name = scope.get_unique_operator_name("Gather") + gather_attrs = {"name": gather_op_name} - # Reshape the indexes we want to embed to 1-D tensor. Otherwise, ONNX Gather's output may get wrong shape. - reshaped_input_name = scope.get_unique_variable_name(gather_op_name + 'input_reshaped') # 2nd input of Gather - apply_reshape(scope, operator.inputs[0].full_name, reshaped_input_name, container, desired_shape=[-1]) + # Reshape the indexes we want to embed to 1-D tensor. + # Otherwise, ONNX Gather's output may get wrong shape. + reshaped_input_name = scope.get_unique_variable_name( + gather_op_name + "input_reshaped" + ) # 2nd input of Gather + apply_reshape( + scope, + operator.inputs[0].full_name, + reshaped_input_name, + container, + desired_shape=[-1], + ) - # ONNX Gather accepts integers so we add a Cast to enforce this before feeding input into ONNX Gather. - casted_input_name = scope.get_unique_variable_name(gather_op_name + 'input_casted') # 2nd input of Gather - apply_cast(scope, reshaped_input_name, casted_input_name, container, to=onnx_proto.TensorProto.INT64) + # ONNX Gather accepts integers so we add a Cast to enforce + # this before feeding input into ONNX Gather. + casted_input_name = scope.get_unique_variable_name( + gather_op_name + "input_casted" + ) # 2nd input of Gather + apply_cast( + scope, + reshaped_input_name, + casted_input_name, + container, + to=onnx_proto.TensorProto.INT64, + ) # Load the embedding matrix. Its shape is outputChannels-by-inputDim. - weights = np.array(params.weights.floatValue).reshape(params.outputChannels, params.inputDim) - weights_name = scope.get_unique_variable_name(gather_op_name + '_W') # 1st input of Gather - container.add_initializer(weights_name, onnx_proto.TensorProto.FLOAT, - [params.inputDim, params.outputChannels], weights.transpose().flatten().tolist()) + weights = np.array(params.weights.floatValue).reshape( + params.outputChannels, params.inputDim + ) + weights_name = scope.get_unique_variable_name( + gather_op_name + "_W" + ) # 1st input of Gather + container.add_initializer( + weights_name, + onnx_proto.TensorProto.FLOAT, + [params.inputDim, params.outputChannels], + weights.transpose().flatten().tolist(), + ) # To support the bias term in an embedding (if exists), we need to create one extra node. if params.hasBias: # Put the embedded result onto a temporal tensor - gather_output_name = scope.get_unique_variable_name(gather_op_name + '_output') - container.add_node('Gather', [weights_name, casted_input_name], gather_output_name, **gather_attrs) + gather_output_name = scope.get_unique_variable_name(gather_op_name + "_output") + container.add_node( + "Gather", + [weights_name, casted_input_name], + gather_output_name, + **gather_attrs + ) # Load the bias vector into an initializer - bias_name = scope.get_unique_variable_name(gather_op_name + '_bias') - bias_axis, bias_shape = deduce_broadcast_axis_and_shape(container.target_opset, [params.outputChannels]) - container.add_initializer(bias_name, onnx_proto.TensorProto.FLOAT, - bias_shape, params.bias.floatValue) - # Create an addition operator to add bias (shape: [C]) into Gather's tensor (shape: [N, C]) - apply_add(scope, [gather_output_name, bias_name], operator.outputs[0].full_name, container, axis=1, broadcast=1) + bias_name = scope.get_unique_variable_name(gather_op_name + "_bias") + bias_axis, bias_shape = deduce_broadcast_axis_and_shape( + container.target_opset, [params.outputChannels] + ) + container.add_initializer( + bias_name, onnx_proto.TensorProto.FLOAT, bias_shape, params.bias.floatValue + ) + # Create an addition operator to add bias (shape: [C]) + # into Gather's tensor (shape: [N, C]) + apply_add( + scope, + [gather_output_name, bias_name], + operator.outputs[0].full_name, + container, + axis=1, + broadcast=1, + ) else: # This case has no bias, so we just output the result produced by the embedding node. - container.add_node('Gather', [weights_name, casted_input_name], operator.output_full_names, **gather_attrs) + container.add_node( + "Gather", + [weights_name, casted_input_name], + operator.output_full_names, + **gather_attrs + ) -register_converter('embedding', convert_embedding) +register_converter("embedding", convert_embedding) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Flatten.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Flatten.py index 07eb6da9..42afa775 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Flatten.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Flatten.py @@ -11,12 +11,18 @@ def convert_flatten(scope, operator, container): flattened_variable_name = operator.outputs[0].full_name if operator.raw_operator.flatten.mode == Params.CHANNEL_LAST: - transposed_variable_name = scope.get_unique_variable_name('transposed') - apply_transpose(scope, variable_to_be_flattened_name, transposed_variable_name, container, perm=[0, 2, 3, 1]) + transposed_variable_name = scope.get_unique_variable_name("transposed") + apply_transpose( + scope, + variable_to_be_flattened_name, + transposed_variable_name, + container, + perm=[0, 2, 3, 1], + ) variable_to_be_flattened_name = transposed_variable_name - op_type = 'Flatten' - flatten_attrs = {'name': operator.full_name, 'axis': 1} + op_type = "Flatten" + flatten_attrs = {"name": operator.full_name, "axis": 1} if container.target_opset < 9: target_opset = 1 @@ -25,7 +31,13 @@ def convert_flatten(scope, operator, container): else: target_opset = 11 - container.add_node(op_type, [variable_to_be_flattened_name], [flattened_variable_name], op_version=target_opset, **flatten_attrs) + container.add_node( + op_type, + [variable_to_be_flattened_name], + [flattened_variable_name], + op_version=target_opset, + **flatten_attrs + ) -register_converter('flatten', convert_flatten) +register_converter("flatten", convert_flatten) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/GRU.py b/onnxmltools/convert/coreml/operator_converters/neural_network/GRU.py index a7ec7783..639f745a 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/GRU.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/GRU.py @@ -70,122 +70,188 @@ def convert_gru(scope, operator, container): hidden_size = params.outputVectorSize # Initialize GRU's attributes. They will be used to build GRU in the end of this function. - gru_op_name = scope.get_unique_operator_name('GRU') - gru_attrs = {'name': gru_op_name} + gru_op_name = scope.get_unique_operator_name("GRU") + gru_attrs = {"name": gru_op_name} gru_inputs = [] gru_outputs = [] # Resahpe CoreML variable into ONNX format for feeding it into ONNX GRU - gru_x_reshape_name = scope.get_unique_variable_name(gru_op_name + '_X_reshape') - apply_reshape(scope, operator.inputs[0].full_name, gru_x_reshape_name, container, desired_shape=[-1, 1, input_size]) + gru_x_reshape_name = scope.get_unique_variable_name(gru_op_name + "_X_reshape") + apply_reshape( + scope, + operator.inputs[0].full_name, + gru_x_reshape_name, + container, + desired_shape=[-1, 1, input_size], + ) gru_inputs.append(gru_x_reshape_name) # Create weight matrices of GRU and add it into ONNX GRU's input list - matrices_w = np.concatenate([params.updateGateWeightMatrix.floatValue, - params.resetGateWeightMatrix.floatValue, - params.outputGateWeightMatrix.floatValue]) - matrices_w_name = scope.get_unique_variable_name(gru_op_name + '_W') - container.add_initializer(matrices_w_name, onnx_proto.TensorProto.FLOAT, - [1, 3 * hidden_size, input_size], matrices_w) + matrices_w = np.concatenate( + [ + params.updateGateWeightMatrix.floatValue, + params.resetGateWeightMatrix.floatValue, + params.outputGateWeightMatrix.floatValue, + ] + ) + matrices_w_name = scope.get_unique_variable_name(gru_op_name + "_W") + container.add_initializer( + matrices_w_name, + onnx_proto.TensorProto.FLOAT, + [1, 3 * hidden_size, input_size], + matrices_w, + ) gru_inputs.append(matrices_w_name) # Create recursion matrices of GRU and add it into ONNX GRU's input list - matrices_r = np.concatenate([params.updateGateRecursionMatrix.floatValue, - params.resetGateRecursionMatrix.floatValue, - params.outputGateRecursionMatrix.floatValue]) - matrices_r_name = scope.get_unique_variable_name(gru_op_name + '_R') - container.add_initializer(matrices_r_name, onnx_proto.TensorProto.FLOAT, - [1, 3 * hidden_size, hidden_size], matrices_r) + matrices_r = np.concatenate( + [ + params.updateGateRecursionMatrix.floatValue, + params.resetGateRecursionMatrix.floatValue, + params.outputGateRecursionMatrix.floatValue, + ] + ) + matrices_r_name = scope.get_unique_variable_name(gru_op_name + "_R") + container.add_initializer( + matrices_r_name, + onnx_proto.TensorProto.FLOAT, + [1, 3 * hidden_size, hidden_size], + matrices_r, + ) gru_inputs.append(matrices_r_name) if params.hasBiasVectors: # Create bias vectors of GRU and add them into ONNX GRU's input list - vectors_b = np.concatenate([params.updateGateBiasVector.floatValue, - params.resetGateBiasVector.floatValue, - params.outputGateBiasVector.floatValue, - np.zeros(3 * hidden_size)]) - vectors_b_name = scope.get_unique_variable_name(gru_op_name + '_B') - container.add_initializer(vectors_b_name, onnx_proto.TensorProto.FLOAT, - [1, 6 * hidden_size], vectors_b) + vectors_b = np.concatenate( + [ + params.updateGateBiasVector.floatValue, + params.resetGateBiasVector.floatValue, + params.outputGateBiasVector.floatValue, + np.zeros(3 * hidden_size), + ] + ) + vectors_b_name = scope.get_unique_variable_name(gru_op_name + "_B") + container.add_initializer( + vectors_b_name, + onnx_proto.TensorProto.FLOAT, + [1, 6 * hidden_size], + vectors_b, + ) gru_inputs.append(vectors_b_name) else: # Because operator's arguments are position-sensitive, we need an empty string even if # this variable doesn't exist. - gru_inputs.append('') + gru_inputs.append("") # The argument, sequence length, is always missing when converting CoreML GRU. - gru_inputs.append('') + gru_inputs.append("") # Handle initial hidden state if it exists if len(operator.inputs) == 2: # Change the shape of initial state in CoreML so that ONNX's GRU is willing to take it. - gru_h_init_reshape_name = scope.get_unique_variable_name(gru_op_name + '_h_init') - apply_reshape(scope, operator.inputs[1].full_name, gru_h_init_reshape_name, container, - desired_shape=[1, 1, hidden_size]) + gru_h_init_reshape_name = scope.get_unique_variable_name( + gru_op_name + "_h_init" + ) + apply_reshape( + scope, + operator.inputs[1].full_name, + gru_h_init_reshape_name, + container, + desired_shape=[1, 1, hidden_size], + ) gru_inputs.append(gru_h_init_reshape_name) # Add a zero initializer to initial hidden state so that this variable becomes optional - container.add_initializer(operator.inputs[1].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[1].type.shape, - np.zeros(shape=operator.inputs[1].type.shape).flatten()) + container.add_initializer( + operator.inputs[1].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[1].type.shape, + np.zeros(shape=operator.inputs[1].type.shape).flatten(), + ) else: # Because operator's arguments are position-sensitive, we need an empty string even if # this variable doesn't exist. - gru_inputs.append('') + gru_inputs.append("") activation_types = [] alphas = [] betas = [] for activation in params.activations: activation_type, alpha, beta = extract_rnn_activation_info(activation) - activation_types.append(activation_type.encode('utf-8')) + activation_types.append(activation_type.encode("utf-8")) if alpha is not None: alphas.append(alpha) if beta is not None: betas.append(beta) - gru_attrs['activations'] = activation_types + gru_attrs["activations"] = activation_types if alphas: - gru_attrs['activation_alpha'] = alphas + gru_attrs["activation_alpha"] = alphas if betas: - gru_attrs['activation_beta'] = betas - gru_attrs['direction'] = 'reverse' if params.reverseInput else 'forward' - gru_attrs['hidden_size'] = hidden_size + gru_attrs["activation_beta"] = betas + gru_attrs["direction"] = "reverse" if params.reverseInput else "forward" + gru_attrs["hidden_size"] = hidden_size # Set up version-dependent attributes if container.target_opset < 5: - gru_attrs['output_sequence'] = params.sequenceOutput + gru_attrs["output_sequence"] = params.sequenceOutput op_version = 1 elif container.target_opset < 7: - gru_attrs['linear_before_reset'] = 0 - gru_attrs['output_sequence'] = params.sequenceOutput + gru_attrs["linear_before_reset"] = 0 + gru_attrs["output_sequence"] = params.sequenceOutput op_version = 3 else: - gru_attrs['linear_before_reset'] = 0 + gru_attrs["linear_before_reset"] = 0 op_version = 7 # Create the major GRU operator in ONNX. - gru_y_name = scope.get_unique_variable_name(gru_op_name + '_Y') - gru_y_h_name = scope.get_unique_variable_name(gru_op_name + '_Y_h') + gru_y_name = scope.get_unique_variable_name(gru_op_name + "_Y") + gru_y_h_name = scope.get_unique_variable_name(gru_op_name + "_Y_h") gru_outputs.extend([gru_y_name, gru_y_h_name]) - container.add_node('GRU', gru_inputs, gru_outputs, op_version=op_version, **gru_attrs) + container.add_node( + "GRU", gru_inputs, gru_outputs, op_version=op_version, **gru_attrs + ) - # To simulate CoreML LSTM, we add post-processing operators to adjust ONNX LSTM outputs + # To simulate CoreML LSTM, we add post-processing + # operators to adjust ONNX LSTM outputs if params.sequenceOutput: - # Again, the output shapes in ONNX's GRU is not consistent with that in CoreML, so we need + # Again, the output shapes in ONNX's GRU + # is not consistent with that in CoreML, so we need # to adjust the result produced by ONNX according to CoreML format. - apply_reshape(scope, gru_y_name, operator.outputs[0].full_name, container, desired_shape=[-1, hidden_size]) + apply_reshape( + scope, + gru_y_name, + operator.outputs[0].full_name, + container, + desired_shape=[-1, hidden_size], + ) # Handle the second output, the last hidden state of a sequence, if exists. if len(operator.outputs) == 2: - apply_reshape(scope, gru_y_h_name, operator.outputs[1].full_name, container, desired_shape=[1, hidden_size]) + apply_reshape( + scope, + gru_y_h_name, + operator.outputs[1].full_name, + container, + desired_shape=[1, hidden_size], + ) else: # Recall that when sequence output is false, the first and the second outputs of GRU # are identical. Thus, we can ignore ONNX GRU's first output. - apply_reshape(scope, gru_y_h_name, operator.outputs[0].full_name, container, desired_shape=[1, hidden_size]) + apply_reshape( + scope, + gru_y_h_name, + operator.outputs[0].full_name, + container, + desired_shape=[1, hidden_size], + ) if len(operator.outputs) == 2: - container.add_node('Identity', operator.outputs[0].full_name, operator.outputs[1].full_name, - name=scope.get_unique_operator_name('Identity')) + container.add_node( + "Identity", + operator.outputs[0].full_name, + operator.outputs[1].full_name, + name=scope.get_unique_operator_name("Identity"), + ) -register_converter('gru', convert_gru) +register_converter("gru", convert_gru) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/ImageScaler.py b/onnxmltools/convert/coreml/operator_converters/neural_network/ImageScaler.py index 3c964f78..0902065b 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/ImageScaler.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/ImageScaler.py @@ -7,40 +7,53 @@ def convert_preprocessing_scaler(scope, operator, container): params = operator.raw_operator.scaler - # Specify some of this operator's attribute. The scale parameter in CoreML is always a scalar. + # Specify some of this operator's attribute. + # The scale parameter in CoreML is always a scalar. # We just copy it and let ONNX scaler to broadcast it to all channels. color_space = operator.inputs[0].type.color_space - if color_space == 'Gray8': + if color_space == "Gray8": bias = [params.grayBias] - elif color_space == 'Rgb8': + elif color_space == "Rgb8": bias = [params.redBias, params.greenBias, params.blueBias] - elif color_space == 'Bgr8': + elif color_space == "Bgr8": bias = [params.blueBias, params.greenBias, params.redBias] else: - raise ValueError('Unknown color space for tensor {}'.format(operator.inputs[0].full_name)) + raise ValueError( + "Unknown color space for tensor {}".format(operator.inputs[0].full_name) + ) if container.target_opset < 9: - attrs = {'name': operator.full_name, 'scale': params.channelScale} - attrs['bias'] = bias - container.add_node('ImageScaler', [operator.inputs[0].full_name], [operator.outputs[0].full_name], **attrs) + attrs = {"name": operator.full_name, "scale": params.channelScale} + attrs["bias"] = bias + container.add_node( + "ImageScaler", + [operator.inputs[0].full_name], + [operator.outputs[0].full_name], + **attrs + ) else: - # In comments below, assume input tensor is X, the scale scalar is a, the bias vector is b. + # In comments below, assume input tensor is X, + # the scale scalar is a, the bias vector is b. # Store the scalar, a, used to scale all elements in the input tensor. - aName = scope.get_unique_variable_name(operator.full_name + '_scale') - container.add_initializer(aName, onnx_proto.TensorProto.FLOAT, [1], [params.channelScale]) + aName = scope.get_unique_variable_name(operator.full_name + "_scale") + container.add_initializer( + aName, onnx_proto.TensorProto.FLOAT, [1], [params.channelScale] + ) # Store the bias vector. It will be added into the input tensor. - bName = scope.get_unique_variable_name(operator.full_name + '_bias') - container.add_initializer(bName, onnx_proto.TensorProto.FLOAT, [len(bias), 1, 1], bias) + bName = scope.get_unique_variable_name(operator.full_name + "_bias") + container.add_initializer( + bName, onnx_proto.TensorProto.FLOAT, [len(bias), 1, 1], bias + ) # Compute Z = a * X. - zName = scope.get_unique_variable_name(operator.full_name + '_scaled') + zName = scope.get_unique_variable_name(operator.full_name + "_scaled") apply_mul(scope, [operator.input_full_names[0], aName], zName, container) # Compute Y = Z + b, which is the final output. apply_add(scope, [bName, zName], operator.output_full_names[0], container) -register_converter('scalerPreprocessor', convert_preprocessing_scaler) +register_converter("scalerPreprocessor", convert_preprocessing_scaler) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/InnerProduct.py b/onnxmltools/convert/coreml/operator_converters/neural_network/InnerProduct.py index 374c1292..c51d0326 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/InnerProduct.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/InnerProduct.py @@ -10,41 +10,56 @@ def convert_inner_product(scope, operator, container): # Apply pre-processing step if needed if len(operator.inputs[0].type.shape) == 4: - # Input shape is [N, C, 1, 1]. Adjust input shape because Gemm in ONNX only takes 2-D input - reshaped_tensor_name = scope.get_unique_variable_name(operator.inputs[0].full_name + '_reshaped') - apply_reshape(scope, operator.inputs[0].full_name, reshaped_tensor_name, container, - desired_shape=[-1, int(params.inputChannels)]) + # Input shape is [N, C, 1, 1]. Adjust input shape + # because Gemm in ONNX only takes 2-D input + reshaped_tensor_name = scope.get_unique_variable_name( + operator.inputs[0].full_name + "_reshaped" + ) + apply_reshape( + scope, + operator.inputs[0].full_name, + reshaped_tensor_name, + container, + desired_shape=[-1, int(params.inputChannels)], + ) name_a = reshaped_tensor_name else: # Input shape is [N, C]. There is no pre-processing for applying ONNX operator. name_a = operator.inputs[0].full_name # Allocate the weights of dense layer - name_b = scope.get_unique_variable_name(operator.full_name + '_B') + name_b = scope.get_unique_variable_name(operator.full_name + "_B") shape_b = [params.outputChannels, params.inputChannels] - container.add_initializer(name_b, onnx_proto.TensorProto.FLOAT, shape_b, params.weights.floatValue) + container.add_initializer( + name_b, onnx_proto.TensorProto.FLOAT, shape_b, params.weights.floatValue + ) # Allocate the bias of dense layer - name_c = scope.get_unique_variable_name(operator.full_name + '_C') + name_c = scope.get_unique_variable_name(operator.full_name + "_C") shape_c = [params.outputChannels] if params.hasBias: - container.add_initializer(name_c, onnx_proto.TensorProto.FLOAT, shape_c, params.bias.floatValue) + container.add_initializer( + name_c, onnx_proto.TensorProto.FLOAT, shape_c, params.bias.floatValue + ) else: - container.add_initializer(name_c, onnx_proto.TensorProto.FLOAT, shape_c, [0.] * shape_b[0]) + container.add_initializer( + name_c, onnx_proto.TensorProto.FLOAT, shape_c, [0.0] * shape_b[0] + ) - # Set up attributes for ONNX Gemm which is the counterpart of CoreML inner product layer in ONNX. - attrs = {'name': operator.full_name} - attrs['alpha'] = 1.0 - attrs['beta'] = 1.0 - attrs['transA'] = 0 - attrs['transB'] = 1 + # Set up attributes for ONNX Gemm which is the counterpart + # of CoreML inner product layer in ONNX. + attrs = {"name": operator.full_name} + attrs["alpha"] = 1.0 + attrs["beta"] = 1.0 + attrs["transA"] = 0 + attrs["transB"] = 1 # Get the correct version number for Gemm in ONNX if container.target_opset < 5: - attrs['broadcast'] = 1 + attrs["broadcast"] = 1 op_version = 1 elif container.target_opset < 7: - attrs['broadcast'] = 1 + attrs["broadcast"] = 1 op_version = 6 elif container.target_opset < 9: op_version = 7 @@ -53,16 +68,36 @@ def convert_inner_product(scope, operator, container): else: op_version = 11 - # Create the major ONNX operator, Gemm, to do CoreML inner product and possibly add shape adjustment + # Create the major ONNX operator, Gemm, to do CoreML + # inner product and possibly add shape adjustment if len(operator.inputs[0].type.shape) == 4: # Input shape is [N, C, 1, 1] so we expect output is also 4-D, [N, C', 1, 1]. - buffer_tensor_name = scope.get_unique_variable_name(operator.full_name + '_buffer') - container.add_node('Gemm', [name_a, name_b, name_c], buffer_tensor_name, op_version=op_version, **attrs) - apply_reshape(scope, buffer_tensor_name, operator.outputs[0].full_name, container, - desired_shape=[-1, int(params.outputChannels), 1, 1]) + buffer_tensor_name = scope.get_unique_variable_name( + operator.full_name + "_buffer" + ) + container.add_node( + "Gemm", + [name_a, name_b, name_c], + buffer_tensor_name, + op_version=op_version, + **attrs + ) + apply_reshape( + scope, + buffer_tensor_name, + operator.outputs[0].full_name, + container, + desired_shape=[-1, int(params.outputChannels), 1, 1], + ) else: # Input shape is [N, C], so we don't need to change Gemm's output shape. - container.add_node('Gemm', [name_a, name_b, name_c], operator.outputs[0].full_name, - op_version=op_version, **attrs) + container.add_node( + "Gemm", + [name_a, name_b, name_c], + operator.outputs[0].full_name, + op_version=op_version, + **attrs + ) -register_converter('innerProduct', convert_inner_product) + +register_converter("innerProduct", convert_inner_product) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/L2Normalize.py b/onnxmltools/convert/coreml/operator_converters/neural_network/L2Normalize.py index 5aef9535..573711da 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/L2Normalize.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/L2Normalize.py @@ -4,9 +4,19 @@ def convert_l2_normalization(scope, operator, container): - # The first dimension is batch size, so the normalization is done along the 2nd axis (indexed by 1). - attrs = {'name': operator.full_name, 'axis': 1, 'p': 2} # Caffe normalization happens per image in one batch - container.add_node('LpNormalization', operator.input_full_names, operator.output_full_names, **attrs) + # The first dimension is batch size, so the + # normalization is done along the 2nd axis (indexed by 1). + attrs = { + "name": operator.full_name, + "axis": 1, + "p": 2, + } # Caffe normalization happens per image in one batch + container.add_node( + "LpNormalization", + operator.input_full_names, + operator.output_full_names, + **attrs + ) -register_converter('l2normalize', convert_l2_normalization) +register_converter("l2normalize", convert_l2_normalization) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/LRN.py b/onnxmltools/convert/coreml/operator_converters/neural_network/LRN.py index 054d6876..5240f843 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/LRN.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/LRN.py @@ -4,14 +4,16 @@ def convert_lrn(scope, operator, container): - op_type = 'LRN' + op_type = "LRN" params = operator.raw_operator.lrn - attrs = {'name': scope.get_unique_operator_name(op_type)} - attrs['size'] = params.localSize - attrs['alpha'] = float(params.alpha) - attrs['beta'] = float(params.beta) - attrs['bias'] = float(params.k) - container.add_node(op_type, operator.input_full_names, operator.output_full_names, **attrs) + attrs = {"name": scope.get_unique_operator_name(op_type)} + attrs["size"] = params.localSize + attrs["alpha"] = float(params.alpha) + attrs["beta"] = float(params.beta) + attrs["bias"] = float(params.k) + container.add_node( + op_type, operator.input_full_names, operator.output_full_names, **attrs + ) -register_converter('lrn', convert_lrn) +register_converter("lrn", convert_lrn) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/LSTM.py b/onnxmltools/convert/coreml/operator_converters/neural_network/LSTM.py index 88b1702d..839d55f3 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/LSTM.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/LSTM.py @@ -8,14 +8,21 @@ def convert_unidirectional_lstm(scope, operator, container): - # The LSTM inputs are feature vector, X, initial hidden state, h_init, and initial cell state, c_init. - # In CorML, their shapes respectively are [S, C_in], [1, C_out], and [1, C_out], where C_in is input feature - # length, # C_out is output dimension, and S is sequence length. Note that S-axis is also known as time axis. - # In ONNX, those shapes become [S, N, C_in] (X), [D, N, C_out] (h_init), and [D, N, C_out]. To simulate - # CoreML LSTM under ONNX, we need some extra operators in addition to LSTM itself. + # The LSTM inputs are feature vector, X, initial hidden state, + # h_init, and initial cell state, c_init. + # In CorML, their shapes respectively are [S, C_in], [1, C_out], + # and [1, C_out], where C_in is input feature + # length, # C_out is output dimension, and S is sequence length. + # Note that S-axis is also known as time axis. + # In ONNX, those shapes become [S, N, C_in] (X), [D, N, C_out] + # (h_init), and [D, N, C_out]. To simulate + # CoreML LSTM under ONNX, we need some extra operators + # in addition to LSTM itself. # - # Note that N=1 and D=1 are always true in ONNX if we are considering LSTM in CoreML because there is no - # batch size in CoreML spec and CoreML LSTM is always uni-directional. + # Note that N=1 and D=1 are always true in ONNX + # if we are considering LSTM in CoreML because there is no + # batch size in CoreML spec and CoreML LSTM + # is always uni-directional. # # Below we provide a visualization of our conversion for CoreML LSTM. # @@ -24,7 +31,8 @@ def convert_unidirectional_lstm(scope, operator, container): # X: input features of CoreML LSTM # h_init: initial LSTM hidden state in CoreML # c_init: initial LSTM cell state in CoreML - # Y: CoreML LSTM's output. It can be [S, C_out] (if sequence_output is on) or [1, C_out] (if sequence_output is off) + # Y: CoreML LSTM's output. It can be [S, C_out] + # (if sequence_output is on) or [1, C_out] (if sequence_output is off) # Y_h: CoreML LSTM's last hidden state # Y_c: CoreML LSTM's last cell state # @@ -131,35 +139,56 @@ def convert_unidirectional_lstm(scope, operator, container): hidden_size = params.outputVectorSize # Initialize materials needed to create ONNX LSTM - lstm_op_name = scope.get_unique_operator_name('LSTM') - lstm_attrs = {'name': lstm_op_name} + lstm_op_name = scope.get_unique_operator_name("LSTM") + lstm_attrs = {"name": lstm_op_name} lstm_inputs = [] lstm_outputs = [] # Reshape input feature vector in CoreML format into ONNX format - lstm_x_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_X_reshape') - apply_reshape(scope, operator.inputs[0].full_name, lstm_x_reshape_name, container, - desired_shape=[-1, 1, input_size]) + lstm_x_reshape_name = scope.get_unique_variable_name(lstm_op_name + "_X_reshape") + apply_reshape( + scope, + operator.inputs[0].full_name, + lstm_x_reshape_name, + container, + desired_shape=[-1, 1, input_size], + ) lstm_inputs.append(lstm_x_reshape_name) # Allocate LSTM's weight matrices and add them into ONNX LSTM's input list - matrices_w = np.concatenate([lstm_weights.inputGateWeightMatrix.floatValue, - lstm_weights.outputGateWeightMatrix.floatValue, - lstm_weights.forgetGateWeightMatrix.floatValue, - lstm_weights.blockInputWeightMatrix.floatValue]) - matrices_w_name = scope.get_unique_variable_name(lstm_op_name + '_W') - container.add_initializer(matrices_w_name, onnx_proto.TensorProto.FLOAT, - [1, 4 * hidden_size, input_size], matrices_w) + matrices_w = np.concatenate( + [ + lstm_weights.inputGateWeightMatrix.floatValue, + lstm_weights.outputGateWeightMatrix.floatValue, + lstm_weights.forgetGateWeightMatrix.floatValue, + lstm_weights.blockInputWeightMatrix.floatValue, + ] + ) + matrices_w_name = scope.get_unique_variable_name(lstm_op_name + "_W") + container.add_initializer( + matrices_w_name, + onnx_proto.TensorProto.FLOAT, + [1, 4 * hidden_size, input_size], + matrices_w, + ) lstm_inputs.append(matrices_w_name) # Allocate LSTM's recursion weight matrices and add them into ONNX LSTM's input list - matrices_r = np.concatenate([lstm_weights.inputGateRecursionMatrix.floatValue, - lstm_weights.outputGateRecursionMatrix.floatValue, - lstm_weights.forgetGateRecursionMatrix.floatValue, - lstm_weights.blockInputRecursionMatrix.floatValue]) - matrices_r_name = scope.get_unique_variable_name(lstm_op_name + '_R') - container.add_initializer(matrices_r_name, onnx_proto.TensorProto.FLOAT, - [1, 4 * hidden_size, hidden_size], matrices_r) + matrices_r = np.concatenate( + [ + lstm_weights.inputGateRecursionMatrix.floatValue, + lstm_weights.outputGateRecursionMatrix.floatValue, + lstm_weights.forgetGateRecursionMatrix.floatValue, + lstm_weights.blockInputRecursionMatrix.floatValue, + ] + ) + matrices_r_name = scope.get_unique_variable_name(lstm_op_name + "_R") + container.add_initializer( + matrices_r_name, + onnx_proto.TensorProto.FLOAT, + [1, 4 * hidden_size, hidden_size], + matrices_r, + ) lstm_inputs.append(matrices_r_name) # Handle bias vectors @@ -174,116 +203,180 @@ def convert_unidirectional_lstm(scope, operator, container): # added 1 into lstm_weights.forgetGateBiasVector.floatValue. pass if lstm_params.hasBiasVectors or lstm_params.forgetBias: - vectors_b_name = scope.get_unique_variable_name(lstm_op_name + '_B') - container.add_initializer(vectors_b_name, onnx_proto.TensorProto.FLOAT, - [1, 8 * hidden_size], vectors_b.flatten()) + vectors_b_name = scope.get_unique_variable_name(lstm_op_name + "_B") + container.add_initializer( + vectors_b_name, + onnx_proto.TensorProto.FLOAT, + [1, 8 * hidden_size], + vectors_b.flatten(), + ) lstm_inputs.append(vectors_b_name) else: - lstm_inputs.append('') + lstm_inputs.append("") # Converting CoreML LSTM doesn't need sequence length - lstm_inputs.append('') + lstm_inputs.append("") # Provide ONNX LSTM the initial hidden state when necessary if len(operator.inputs) > 1: - # Assign a Reshape to adjust CoreML hidden state's shape [1, C]/[1, C, 1, 1] into its ONNX counterpart [1, 1, C] - lstm_h_init_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_h_init_reshape') - apply_reshape(scope, operator.inputs[1].full_name, lstm_h_init_reshape_name, container, - desired_shape=[1, 1, hidden_size]) + # Assign a Reshape to adjust CoreML hidden state's shape + # [1, C]/[1, C, 1, 1] into its ONNX counterpart [1, 1, C] + lstm_h_init_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_h_init_reshape" + ) + apply_reshape( + scope, + operator.inputs[1].full_name, + lstm_h_init_reshape_name, + container, + desired_shape=[1, 1, hidden_size], + ) lstm_inputs.append(lstm_h_init_reshape_name) # Add a zero initializer to initial hidden state so that this variable becomes optional - container.add_initializer(operator.inputs[1].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[1].type.shape, - np.zeros(shape=operator.inputs[1].type.shape).flatten()) + container.add_initializer( + operator.inputs[1].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[1].type.shape, + np.zeros(shape=operator.inputs[1].type.shape).flatten(), + ) else: - lstm_inputs.append('') + lstm_inputs.append("") # Provide ONNX LSTM the initial cell state when necessary if len(operator.inputs) > 2: - lstm_c_init_reshape_name = scope.get_unique_variable_name(lstm_op_name + '_c_init_reshape') - apply_reshape(scope, operator.inputs[2].full_name, lstm_c_init_reshape_name, container, - desired_shape=[1, 1, hidden_size]) + lstm_c_init_reshape_name = scope.get_unique_variable_name( + lstm_op_name + "_c_init_reshape" + ) + apply_reshape( + scope, + operator.inputs[2].full_name, + lstm_c_init_reshape_name, + container, + desired_shape=[1, 1, hidden_size], + ) lstm_inputs.append(lstm_c_init_reshape_name) # Add a zero initializer to initial cell state so that this variable becomes optional - container.add_initializer(operator.inputs[2].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[2].type.shape, - np.zeros(shape=operator.inputs[2].type.shape).flatten()) + container.add_initializer( + operator.inputs[2].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[2].type.shape, + np.zeros(shape=operator.inputs[2].type.shape).flatten(), + ) else: - lstm_inputs.append('') + lstm_inputs.append("") # Add peephole vector when presenting if lstm_params.hasPeepholeVectors: - vectors_p = np.concatenate([lstm_weights.inputGatePeepholeVector.floatValue, - lstm_weights.outputGatePeepholeVector.floatValue, - lstm_weights.forgetGatePeepholeVector.floatValue]) - vectors_p_name = scope.get_unique_variable_name(lstm_op_name + '_P') - container.add_initializer(vectors_p_name, onnx_proto.TensorProto.FLOAT, - [1, 3 * hidden_size], vectors_p) + vectors_p = np.concatenate( + [ + lstm_weights.inputGatePeepholeVector.floatValue, + lstm_weights.outputGatePeepholeVector.floatValue, + lstm_weights.forgetGatePeepholeVector.floatValue, + ] + ) + vectors_p_name = scope.get_unique_variable_name(lstm_op_name + "_P") + container.add_initializer( + vectors_p_name, + onnx_proto.TensorProto.FLOAT, + [1, 3 * hidden_size], + vectors_p, + ) lstm_inputs.append(vectors_p_name) else: - lstm_inputs.append('') + lstm_inputs.append("") - # Parse activation functions' information and add them into ONNX LSTM's attribute dictionary + # Parse activation functions' information and add + # them into ONNX LSTM's attribute dictionary activation_types = [] alphas = [] betas = [] for activation in params.activations: activation_type, alpha, beta = extract_rnn_activation_info(activation) - activation_types.append(activation_type.encode('utf-8')) + activation_types.append(activation_type.encode("utf-8")) if alpha is not None: alphas.append(alpha) if beta is not None: betas.append(beta) - lstm_attrs['activations'] = activation_types + lstm_attrs["activations"] = activation_types if alphas: - lstm_attrs['activation_alpha'] = alphas + lstm_attrs["activation_alpha"] = alphas if betas: - lstm_attrs['activation_beta'] = betas + lstm_attrs["activation_beta"] = betas # Set up other attributes - lstm_attrs['direction'] = 'reverse' if params.reverseInput else 'forward' - lstm_attrs['hidden_size'] = hidden_size - lstm_attrs['clip'] = float(lstm_params.cellClipThreshold) - lstm_attrs['input_forget'] = lstm_params.coupledInputAndForgetGate + lstm_attrs["direction"] = "reverse" if params.reverseInput else "forward" + lstm_attrs["hidden_size"] = hidden_size + lstm_attrs["clip"] = float(lstm_params.cellClipThreshold) + lstm_attrs["input_forget"] = lstm_params.coupledInputAndForgetGate # Set up version-dependent attributes if container.target_opset < 7: - lstm_attrs['output_sequence'] = lstm_params.sequenceOutput + lstm_attrs["output_sequence"] = lstm_params.sequenceOutput op_version = 1 else: op_version = 7 # Create the main LSTM operator - lstm_y_name = scope.get_unique_variable_name(lstm_op_name + '_Y') - lstm_y_h_name = scope.get_unique_variable_name(lstm_op_name + '_Y_h') - lstm_c_name = scope.get_unique_variable_name(lstm_op_name + '_Y_c') + lstm_y_name = scope.get_unique_variable_name(lstm_op_name + "_Y") + lstm_y_h_name = scope.get_unique_variable_name(lstm_op_name + "_Y_h") + lstm_c_name = scope.get_unique_variable_name(lstm_op_name + "_Y_c") lstm_outputs.extend([lstm_y_name, lstm_y_h_name, lstm_c_name]) - container.add_node('LSTM', lstm_inputs, lstm_outputs, op_version=op_version, **lstm_attrs) + container.add_node( + "LSTM", lstm_inputs, lstm_outputs, op_version=op_version, **lstm_attrs + ) # Handle the first output of LSTM if lstm_params.sequenceOutput: # Handle the first output of LSTM - apply_reshape(scope, lstm_y_name, operator.outputs[0].full_name, container, desired_shape=[-1, hidden_size]) + apply_reshape( + scope, + lstm_y_name, + operator.outputs[0].full_name, + container, + desired_shape=[-1, hidden_size], + ) # Handle the second output of LSTM if len(operator.outputs) > 1: - apply_reshape(scope, lstm_y_h_name, operator.outputs[1].full_name, container, - desired_shape=[1, hidden_size]) + apply_reshape( + scope, + lstm_y_h_name, + operator.outputs[1].full_name, + container, + desired_shape=[1, hidden_size], + ) else: - # Here we ingore ONNX LSTM's first output because it's useless and use the second output of ONNX LSTM to produce + # Here we ingore ONNX LSTM's first output because + # it's useless and use the second output of ONNX LSTM to produce # the first output of CoreML LSTM - apply_reshape(scope, lstm_y_h_name, operator.outputs[0].full_name, container, desired_shape=[1, hidden_size]) + apply_reshape( + scope, + lstm_y_h_name, + operator.outputs[0].full_name, + container, + desired_shape=[1, hidden_size], + ) # Create the second LSTM output from the first output if len(operator.outputs) > 1: - container.add_node('Identity', operator.outputs[0].full_name, operator.outputs[1].full_name, - name=scope.get_unique_operator_name('Identity')) + container.add_node( + "Identity", + operator.outputs[0].full_name, + operator.outputs[1].full_name, + name=scope.get_unique_operator_name("Identity"), + ) # Handle the cell state output of LSTM if len(operator.outputs) > 2: - apply_reshape(scope, lstm_c_name, operator.outputs[2].full_name, container, desired_shape=[1, hidden_size]) + apply_reshape( + scope, + lstm_c_name, + operator.outputs[2].full_name, + container, + desired_shape=[1, hidden_size], + ) -register_converter('uniDirectionalLSTM', convert_unidirectional_lstm) +register_converter("uniDirectionalLSTM", convert_unidirectional_lstm) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstant.py b/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstant.py index f7ab7e66..78c75d2d 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstant.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstant.py @@ -5,13 +5,24 @@ from ....common._registration import register_converter from ....common._apply_operation import apply_constant + def convert_load_constant(scope, operator, container): params = operator.raw_operator.loadConstant - constant_name = scope.get_unique_variable_name('constant') - constant = helper.make_tensor(constant_name, onnx_proto.TensorProto.FLOAT, - params.shape, params.data.floatValue) + constant_name = scope.get_unique_variable_name("constant") + constant = helper.make_tensor( + constant_name, + onnx_proto.TensorProto.FLOAT, + params.shape, + params.data.floatValue, + ) + + apply_constant( + scope, + operator.output_full_names, + container, + operator_name=operator.full_name, + value=constant, + ) - apply_constant(scope, operator.output_full_names, container, - operator_name=operator.full_name, value=constant) -register_converter('loadConstant', convert_load_constant) +register_converter("loadConstant", convert_load_constant) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstantND.py b/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstantND.py index 65dee927..c76f9426 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstantND.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/LoadConstantND.py @@ -5,13 +5,24 @@ from ....common._registration import register_converter from ....common._apply_operation import apply_constant + def convert_load_constant_nd(scope, operator, container): params = operator.raw_operator.loadConstantND - constant_name = scope.get_unique_variable_name('constant') - constant = helper.make_tensor(constant_name, onnx_proto.TensorProto.FLOAT, - params.shape, params.data.floatValue) + constant_name = scope.get_unique_variable_name("constant") + constant = helper.make_tensor( + constant_name, + onnx_proto.TensorProto.FLOAT, + params.shape, + params.data.floatValue, + ) + + apply_constant( + scope, + operator.output_full_names, + container, + operator_name=operator.full_name, + value=constant, + ) - apply_constant(scope, operator.output_full_names, container, - operator_name=operator.full_name, value=constant) -register_converter('loadConstantND', convert_load_constant_nd) +register_converter("loadConstantND", convert_load_constant_nd) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Max.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Max.py index 3a91d6f2..e297a2cb 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Max.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Max.py @@ -5,8 +5,13 @@ def convert_max(scope, operator, container): - apply_max(scope, operator.input_full_names, operator.output_full_names, - container, operator_name=operator.full_name) + apply_max( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) -register_converter('max', convert_max) +register_converter("max", convert_max) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/MeanImage.py b/onnxmltools/convert/coreml/operator_converters/neural_network/MeanImage.py index 4a50c2ff..0ebd3a0b 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/MeanImage.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/MeanImage.py @@ -6,17 +6,29 @@ def convert_preprocessing_mean_image(scope, operator, container): - mean_tensor_name = scope.get_unique_variable_name(operator.full_name + '_mean') + mean_tensor_name = scope.get_unique_variable_name(operator.full_name + "_mean") - # We assume that the first input's shape is [N, C, H, W] so that the mean image's shape, [C, H, W], can + # We assume that the first input's shape is [N, C, H, W] + # so that the mean image's shape, [C, H, W], can # be inferred from the first input's shape. - container.add_initializer(mean_tensor_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[0].type.shape[1:], operator.raw_operator.meanImage) + container.add_initializer( + mean_tensor_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[0].type.shape[1:], + operator.raw_operator.meanImage, + ) - # We assume that the first input variable's shape is [N, C, H, W] while the mean image's shape is [C, H, W]. Thus, + # We assume that the first input variable's shape is [N, C, H, W] + # while the mean image's shape is [C, H, W]. Thus, # broadcasting should be enabled starting with axis=1. - apply_sub(scope, [operator.inputs[0].full_name, mean_tensor_name], operator.output_full_names, container, - axis=1, broadcast=1) + apply_sub( + scope, + [operator.inputs[0].full_name, mean_tensor_name], + operator.output_full_names, + container, + axis=1, + broadcast=1, + ) -register_converter('meanImagePreprocessor', convert_preprocessing_mean_image) +register_converter("meanImagePreprocessor", convert_preprocessing_mean_image) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/MeanVarianceNorm.py b/onnxmltools/convert/coreml/operator_converters/neural_network/MeanVarianceNorm.py index a80e5210..0a23bf98 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/MeanVarianceNorm.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/MeanVarianceNorm.py @@ -4,13 +4,15 @@ def convert_mean_variance_normalization(scope, operator, container): - op_type = 'MeanVarianceNormalization' + op_type = "MeanVarianceNormalization" op_name = scope.get_unique_operator_name(op_type) params = operator.raw_operator.mvn - attrs = {'name': op_name} - attrs['across_channels'] = params.acrossChannels - attrs['normalize_variance'] = params.normalizeVariance - container.add_node(op_type, operator.input_full_names, operator.output_full_names, **attrs) + attrs = {"name": op_name} + attrs["across_channels"] = params.acrossChannels + attrs["normalize_variance"] = params.normalizeVariance + container.add_node( + op_type, operator.input_full_names, operator.output_full_names, **attrs + ) -register_converter('mvn', convert_mean_variance_normalization) +register_converter("mvn", convert_mean_variance_normalization) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Min.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Min.py index 0380c023..b2a0c200 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Min.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Min.py @@ -5,7 +5,13 @@ def convert_min(scope, operator, container): - apply_min(scope, operator.input_full_names, operator.output_full_names, container, operator.full_name) + apply_min( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator.full_name, + ) -register_converter('min', convert_min) +register_converter("min", convert_min) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Multiply.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Multiply.py index cb3a7048..f0d30431 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Multiply.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Multiply.py @@ -8,12 +8,23 @@ def convert_multiply(scope, operator, container): if len(operator.input_full_names) == 1: # Multiply the input tensor by a scalar, named alpha. - scaler_name = scope.get_unique_variable_name(operator.full_name + '_B') - container.add_initializer(scaler_name, onnx_proto.TensorProto.FLOAT, [], [operator.raw_operator.multiply.alpha]) + scaler_name = scope.get_unique_variable_name(operator.full_name + "_B") + container.add_initializer( + scaler_name, + onnx_proto.TensorProto.FLOAT, + [], + [operator.raw_operator.multiply.alpha], + ) inputs = [operator.inputs[0].full_name, scaler_name] broadcast = 1 - apply_mul(scope, inputs, operator.output_full_names, container, operator_name=operator.full_name, - broadcast=broadcast) + apply_mul( + scope, + inputs, + operator.output_full_names, + container, + operator_name=operator.full_name, + broadcast=broadcast, + ) else: inputs = operator.input_full_names @@ -26,28 +37,49 @@ def convert_multiply(scope, operator, container): else: broadcast = 0 - apply_mul(scope, [left_tensor, right_tensor], operator.outputs[0].full_name, container, - operator_name=operator.full_name, broadcast=broadcast) + apply_mul( + scope, + [left_tensor, right_tensor], + operator.outputs[0].full_name, + container, + operator_name=operator.full_name, + broadcast=broadcast, + ) else: # In this case we calculate the multiplication of multiple input tensors # Sum up the first two inputs left_tensor = operator.inputs[0].full_name right_tensor = operator.inputs[1].full_name - intermediate_tensor_name = scope.get_unique_variable_name('buffer_tensor') - apply_mul(scope, [left_tensor, right_tensor], intermediate_tensor_name, container, - operator_name=operator.full_name, broadcast=1) + intermediate_tensor_name = scope.get_unique_variable_name("buffer_tensor") + apply_mul( + scope, + [left_tensor, right_tensor], + intermediate_tensor_name, + container, + operator_name=operator.full_name, + broadcast=1, + ) - # Accumulate other inputs onto intermediate tensors. Note that we may use the original operator's output as + # Accumulate other inputs onto intermediate tensors. + # Note that we may use the original operator's output as # the last intermediate tensor. for i in range(2, len(inputs)): left_tensor = intermediate_tensor_name right_tensor = inputs[i].full_name if i != len(inputs) - 1: - intermediate_tensor_name = scope.get_unique_variable_name('buffer_tensor') + intermediate_tensor_name = scope.get_unique_variable_name( + "buffer_tensor" + ) else: intermediate_tensor_name = operator.outputs[0].full_name - apply_mul(scope, [left_tensor, right_tensor], intermediate_tensor_name, container, broadcast=1) + apply_mul( + scope, + [left_tensor, right_tensor], + intermediate_tensor_name, + container, + broadcast=1, + ) -register_converter('multiply', convert_multiply) +register_converter("multiply", convert_multiply) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Pad.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Pad.py index a8c0c7ae..a16087df 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Pad.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Pad.py @@ -7,17 +7,21 @@ def convert_padding(scope, operator, container): params = operator.raw_operator.padding - pad_table = {'constant': 'constant', 'reflection': 'reflect', 'replication': 'edge'} - pad_type = params.WhichOneof('PaddingType') + pad_table = {"constant": "constant", "reflection": "reflect", "replication": "edge"} + pad_type = params.WhichOneof("PaddingType") if pad_type not in pad_table: - raise ValueError('Unsupported padding mode: {}'.format(pad_type)) + raise ValueError("Unsupported padding mode: {}".format(pad_type)) mode = pad_table[pad_type] - # CoreML only pads for their H- and W-axes. Here we assume the shape of the tensor to be padded + # CoreML only pads for their H- and W-axes. + # Here we assume the shape of the tensor to be padded # is [N, C, H, W], so we have 8 padding amounts - # pads = [N_begin_index, C_begin_index, H_begin_index, W_begin_index, - # N_end_index, C_end_index, H_end_index, W_end_index] - # Because only H- and W-axes are padded in CoreML, we leave padding amounts of N- and C-axes zeros. + # pads = [N_begin_index, C_begin_index, + # H_begin_index, W_begin_index, + # N_end_index, C_end_index, + # H_end_index, W_end_index] + # Because only H- and W-axes are padded in CoreML, + # we leave padding amounts of N- and C-axes zeros. pads = [0, 0, 0, 0, 0, 0, 0, 0] if len(params.paddingAmounts.borderAmounts) > 0: # Set H_begin_index @@ -29,13 +33,21 @@ def convert_padding(scope, operator, container): # Set W_end_index pads[7] = params.paddingAmounts.borderAmounts[1].endEdgeSize - if pad_type == 'constant': + if pad_type == "constant": value = params.constant.value else: value = None - apply_pad(scope, operator.input_full_names, operator.output_full_names, container, operator_name=operator.full_name, - mode=mode, pads=pads, value=value) + apply_pad( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator_name=operator.full_name, + mode=mode, + pads=pads, + value=value, + ) -register_converter('padding', convert_padding) +register_converter("padding", convert_padding) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Permute.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Permute.py index 10c1b5e9..ea8abc3f 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Permute.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Permute.py @@ -5,8 +5,14 @@ def convert_permute(scope, operator, container): - apply_transpose(scope, operator.input_full_names, operator.output_full_names, container, - operator_name=operator.full_name, perm=operator.raw_operator.permute.axis) + apply_transpose( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator_name=operator.full_name, + perm=operator.raw_operator.permute.axis, + ) -register_converter('permute', convert_permute) +register_converter("permute", convert_permute) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Pool.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Pool.py index 9d1e1727..890b60e1 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Pool.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Pool.py @@ -7,8 +7,9 @@ def calculate_legacy_pad_amount(H_in, pad_h, k_h, s_h): - ''' - This function calculate padding amount along H-axis. It can be applied to other axes. It should be only used with + """ + This function calculate padding amount along H-axis. + It can be applied to other axes. It should be only used with pooling conversion. :param H_in: input dimension along H-axis @@ -16,24 +17,29 @@ def calculate_legacy_pad_amount(H_in, pad_h, k_h, s_h): :param k_h: kernel's H-axis dimension :param s_h: stride along H-axis :return: (top_padding_amount, bottom_padding_amount) - ''' + """ # Calculate a common variable H_temp = H_in + 2 * pad_h - k_h # Pooling output shape under CoerML IncludeLastPixel padding mode H_include_last_pad_out = math.ceil(H_temp / s_h) + 1 # Pooling output shape under valid padding mode H_valid_pad_out = math.floor(H_temp / s_h) + 1 - # Amount of values padded at top boundary. For max pooling, the padded value should be "-inf." + # Amount of values padded at top boundary. + # For max pooling, the padded value should be "-inf." # For average pooling, we should pad zeros. pad_t = pad_h - # Amount of values padded at bottom boundary (add extra pixels so that H_include_last_pad_out = floor( (H_adjusted_out - k_h) / stride) + 1) + # Amount of values padded at bottom boundary + # (add extra pixels so that + # H_include_last_pad_out = floor( (H_adjusted_out - k_h) / stride) + 1) if H_include_last_pad_out > H_valid_pad_out: pad_b = pad_h + (s_h - H_temp % s_h) else: pad_b = pad_h - # Intermediate result with pad_t values at top and pad_b valules at bottom of the original input + # Intermediate result with pad_t values at top and pad_b + # valules at bottom of the original input H_adjusted_out = H_in + pad_t + pad_b - # Adjust padded result if the original pooling wants to cut off the last output pixel. + # Adjust padded result if the original pooling wants to cut + # off the last output pixel. if (H_include_last_pad_out - 1) * s_h >= H_in + pad_h: if H_adjusted_out % s_h == 0: H_adjusted_out -= s_h @@ -42,11 +48,26 @@ def calculate_legacy_pad_amount(H_in, pad_h, k_h, s_h): return (pad_t, H_adjusted_out - pad_t - H_in) -def create_legacy_pad(scope, input_name, output_name, H_in, W_in, k_h, k_w, - s_h, s_w, p_h, p_w, padded_value, container): - ''' - This function adds one Pad operator into its last argument, which is a Container object. By feeding the output of - the created Pad operator into Pool operator under valid padding mode, we can achieve the same functionality of +def create_legacy_pad( + scope, + input_name, + output_name, + H_in, + W_in, + k_h, + k_w, + s_h, + s_w, + p_h, + p_w, + padded_value, + container, +): + """ + This function adds one Pad operator into its last argument, + which is a Container object. By feeding the output of + the created Pad operator into Pool operator under valid + padding mode, we can achieve the same functionality of CoreML' pooling under IncludeLastPixel padding mode. :param scope: @@ -62,53 +83,66 @@ def create_legacy_pad(scope, input_name, output_name, H_in, W_in, k_h, k_w, :param p_w: padding amount at the beginning and the end of W-axis :param padded_value: value used to fill padded area :param container: Container object - ''' + """ # Add a Pad operator to pre-process 4-D tensor pad_t, pad_b = calculate_legacy_pad_amount(H_in, p_h, k_h, s_h) pad_l, pad_r = calculate_legacy_pad_amount(W_in, p_w, k_w, s_w) - # CoreML pooling operator pads only their H- and W-axes. Here we assume the shape of the tensor to be padded + # CoreML pooling operator pads only their H- and W-axes. + # Here we assume the shape of the tensor to be padded # is [N, C, H, W], so we have 8 padding amounts # pads = [N_begin_index, C_begin_index, H_begin_index, W_begin_index, # N_end_index, C_end_index, H_end_index, W_end_index] - # Because only H- and W-axes are padded in CoreML, we leave padding amounts of N- and C-axes zeros. + # Because only H- and W-axes are padded in CoreML, + # we leave padding amounts of N- and C-axes zeros. pads = [0, 0, pad_t, pad_l, 0, 0, pad_b, pad_r] apply_pad(scope, input_name, output_name, container, pads=pads, value=padded_value) -# The conversion of pooling has several possible outcomes. Let's first define some symbols and then discuss their +# The conversion of pooling has several possible outcomes. +# Let's first define some symbols and then discuss their # ONNX computational graphs case-by-case. # # Symbols: -# X: input 4-D tensor. It should have a shape [N, C, H, W] following CoreML's pooling definition. -# Y: output tensor identical to CorML's pooling. Its shapes depends on the pooling type applied. +# X: input 4-D tensor. It should have a shape [N, C, H, W] +# following CoreML's pooling definition. +# Y: output tensor identical to CorML's pooling. Its shapes +# depends on the pooling type applied. # # Case 1: global pooling # X ---> ONNX Global Pooling ---> Y -# In this case, ONNX's pooling implementation should directly match CoreML's implementation, so it's just a naive +# In this case, ONNX's pooling implementation should directly +# match CoreML's implementation, so it's just a naive # translation. # # Case 2: local max/L2 pooling with same or valid padding # # X ---> ONNX Local Max/L2 Pooling ---> Y # -# In this case, ONNX's pooling implementation should directly match CoreML's implementation, so it's just a naive +# In this case, ONNX's pooling implementation should directly +# match CoreML's implementation, so it's just a naive # translation. # # Case 3: local max/L2 pooling under CoreML's IncludeLastPixel padding # # X ---> Pad --> X' ---> ONNX Local Max/L2 Pooling ---> Y # -# CoreML's IncludeLastPixel padding mode is not supported in ONNX's pooling. We combine a Pad -# operator and a pooling to simulate CoreML's behavior. In this case, the Pad takes all padding-related -# parameters from CoreML's pooling while ONNX's pooling is working under valid padding mode. +# CoreML's IncludeLastPixel padding mode is not supported in ONNX's +# pooling. We combine a Pad +# operator and a pooling to simulate CoreML's behavior. +# In this case, the Pad takes all padding-related +# parameters from CoreML's pooling while ONNX's pooling is +# working under valid padding mode. # -# Case 4: local average pooling with same or valid padding. exclude_pad_area is on. +# Case 4: local average pooling with same or valid padding. +# exclude_pad_area is on. # # X ---> ONNX Local Average Pooling ---> Y # -# Current ONNX pooling operator just follows Caffe2, so the padded area is naturally excluded when calculating the -# numerator and denumerator of the pixel average covered by the kernel. That is, the translation from CoreML to ONNX +# Current ONNX pooling operator just follows Caffe2, so the padded area +# is naturally excluded when calculating the +# numerator and denumerator of the pixel average covered by the kernel. +# That is, the translation from CoreML to ONNX # is trivial. # # Case 5: local average pooling with same or valid padding. exclude_pad_area is off. @@ -118,11 +152,14 @@ def create_legacy_pad(scope, input_name, output_name, H_in, W_in, k_h, k_w, # | | # '---> Scaler ---> Z ---> ONNX L1-norm Pooling ---> Z' ---' # -# The Scaler has "alpha=0" and its "beta" is a constant. If "beta=1", the output of the L1-norm pooling, Z', is the -# effective kernel size applied at each pixel when padded area is excluded. Here we use "beta=1/kerenel_size" so +# The Scaler has "alpha=0" and its "beta" is a constant. If "beta=1", +# the output of the L1-norm pooling, Z', is the +# effective kernel size applied at each pixel when padded area is +# excluded. Here we use "beta=1/kerenel_size" so # that one value in Z' stands for # (the kernel size without padded area) / (the kernel size with padded area) -# at a pixel. The output Y' is computed with exclude_pad_area=on, so the element-wise multiplication of Y' and Z' +# at a pixel. The output Y' is computed with exclude_pad_area=on, +# so the element-wise multiplication of Y' and Z' # is Y. # # Case 6: local average pooling with IncludeLastPixel padding. exclude_pad_area is on. @@ -132,18 +169,24 @@ def create_legacy_pad(scope, input_name, output_name, H_in, W_in, k_h, k_w, # | | # '---> Scaler ---> Z ---> Pad ---> Z' ---> ONNX L1-norm Pooling ---> Z'' # -# The Scaler has "alpha=0" and its "beta" is a constant. If "beta=1", the output of the L1-norm pooling, Z'', is -# the effective kernel size applied at each pixel when padded area is excluded (since Pad fills the -# padded area with zeros so those padded pixels are not counted by the L1-norm pooling). Here we use +# The Scaler has "alpha=0" and its "beta" is a constant. +# If "beta=1", the output of the L1-norm pooling, Z'', is +# the effective kernel size applied at each pixel when padded +# area is excluded (since Pad fills the +# padded area with zeros so those padded pixels are not counted +# by the L1-norm pooling). Here we use # "beta=1/kerenel_size" so that one value in Z' stands for # (the kernel size without padded area) / (the kernel size with padded area) -# at a pixel. The output Y' is computed as if exclude_pad_area=on, so the element-wise division of Y' and Z'' is Y. +# at a pixel. The output Y' is computed as if exclude_pad_area=on, +# so the element-wise division of Y' and Z'' is Y. # -# Case 7: local average pooling with IncludeLastPixel padding. exclude_pad_area is off. +# Case 7: local average pooling with IncludeLastPixel padding. +# exclude_pad_area is off. # # X ---> Pad --> X' ---> ONNX Local Average Pooling ---> Y # -# Since Pad operators add zeros to X's margin and the local pooling here is working under valid padding, it's +# Since Pad operators add zeros to X's margin and the local +# pooling here is working under valid padding, it's # equivalent to the situation of exclude_pad_area=off. def convert_pooling(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import PoolingLayerParams as Params @@ -155,19 +198,22 @@ def convert_pooling(scope, operator, container): # The output of Pool outputs = [variable.full_name for variable in operator.outputs] - # Handle global pooling mode. This case if much simpler than the conversion of local pooling. - attrs = {'name': operator.full_name} + # Handle global pooling mode. This case if much simpler + # than the conversion of local pooling. + attrs = {"name": operator.full_name} if params.globalPooling: - pooling_table = {params.MAX: 'GlobalMaxPool', - Params.AVERAGE: 'GlobalAveragePool', - Params.L2: 'GlobalLpPool'} + pooling_table = { + params.MAX: "GlobalMaxPool", + Params.AVERAGE: "GlobalAveragePool", + Params.L2: "GlobalLpPool", + } if params.type not in pooling_table: - raise ValueError('Unsupported pooling type: {}'.format(params.type)) + raise ValueError("Unsupported pooling type: {}".format(params.type)) op_type = pooling_table[params.type] if params.type == Params.L2: - attrs['p'] = 2 + attrs["p"] = 2 container.add_node(op_type, inputs, outputs, op_version=2, **attrs) else: container.add_node(op_type, inputs, outputs, **attrs) @@ -175,41 +221,41 @@ def convert_pooling(scope, operator, container): # From here to the end of this function, we will handle local pooling mode if params.type == Params.MAX: - op_type = 'MaxPool' + op_type = "MaxPool" if container.target_opset < 8: op_version = 1 elif container.target_opset < 10: op_version = 8 else: op_version = 10 - attrs['ceil_mode'] = 0 + attrs["ceil_mode"] = 0 elif params.type == Params.AVERAGE: - op_type = 'AveragePool' + op_type = "AveragePool" if container.target_opset < 7: op_version = 1 elif container.target_opset < 10: op_version = 7 else: op_version = 10 - attrs['ceil_mode'] = 0 + attrs["ceil_mode"] = 0 elif params.type == Params.L2: - op_type = 'LpPool' - attrs['p'] = 2 + op_type = "LpPool" + attrs["p"] = 2 op_version = 2 else: - raise ValueError('Unsupported pooling type: {}'.format(params.type)) + raise ValueError("Unsupported pooling type: {}".format(params.type)) # CoreML default v.s. non-default parameters kernel_shape = [3, 3] if len(params.kernelSize) <= 0 else params.kernelSize strides = [1, 1] if len(params.stride) <= 0 else params.stride - attrs['kernel_shape'] = kernel_shape - attrs['strides'] = strides + attrs["kernel_shape"] = kernel_shape + attrs["strides"] = strides # Set up padding attributes pads = None auto_pad = None - pad_type = params.WhichOneof('PoolingPaddingType') - if pad_type == 'valid': + pad_type = params.WhichOneof("PoolingPaddingType") + if pad_type == "valid": if len(params.valid.paddingAmounts.borderAmounts) > 0: pads = [0, 0, 0, 0] pads[0] = params.valid.paddingAmounts.borderAmounts[0].startEdgeSize @@ -219,78 +265,127 @@ def convert_pooling(scope, operator, container): # If padding amounts are all zero, there should be no padding list. if all(pad == 0 for pad in pads): pads = None - auto_pad = 'VALID' + auto_pad = "VALID" else: - auto_pad = 'VALID' - elif pad_type == 'same': + auto_pad = "VALID" + elif pad_type == "same": if params.same.asymmetryMode == SamePadding.BOTTOM_RIGHT_HEAVY: - # In CoreML, BOTTOM_RIGHT_HEAVY means that the extra pixel (when not dividable by 2) will be added to the + # In CoreML, BOTTOM_RIGHT_HEAVY means that the extra pixel + # (when not dividable by 2) will be added to the # end of an axis. This behavior matches ONNX's SAME_UPPER mode. - # Reference: https://apple.github.io/coremltools/coremlspecification/sections/NeuralNetwork.html#samepadding - # https://github.com/onnx/onnx/blob/rel-1.2.1/docs/Operators.md#AveragePool - auto_pad = 'SAME_UPPER' + # Reference: + # https://apple.github.io/coremltools/ + # coremlspecification/sections/NeuralNetwork.html#samepadding + # https://github.com/onnx/onnx/blob/rel-1.2.1/docs/Operators.md#AveragePool + auto_pad = "SAME_UPPER" elif params.same.asymmetryMode == SamePadding.TOP_LEFT_HEAVY: - # In CoreML, TOP_LEFT_HEAVY means that the extra pixel (when not dividable by 2) will be added to the + # In CoreML, TOP_LEFT_HEAVY means that the extra + # pixel (when not dividable by 2) will be added to the # beginning of an axis. This behavior matches ONNX's SAME_LOWER mode. - # Reference: https://apple.github.io/coremltools/coremlspecification/sections/NeuralNetwork.html#samepadding - # https://github.com/onnx/onnx/blob/rel-1.2.1/docs/Operators.md#AveragePool - auto_pad = 'SAME_LOWER' + # Reference: + # https://apple.github.io/coremltools/ + # coremlspecification/sections/NeuralNetwork.html#samepadding + # https://github.com/onnx/onnx/blob/rel-1.2.1/docs/Operators.md#AveragePool + auto_pad = "SAME_LOWER" else: - raise ValueError('Unknown asymmetric mode: {}'.format(params.same.asymmetryMode)) - elif pad_type == 'includeLastPixel': + raise ValueError( + "Unknown asymmetric mode: {}".format(params.same.asymmetryMode) + ) + elif pad_type == "includeLastPixel": # Here we use a Pad operator to mimic the behavior of this CoreML padding. - auto_pad = 'VALID' + auto_pad = "VALID" H = operator.inputs[0].type.shape[2] W = operator.inputs[0].type.shape[3] pad_h = params.includeLastPixel.paddingAmounts[0] pad_w = params.includeLastPixel.paddingAmounts[1] - legacy_padded_tensor_name = scope.get_unique_variable_name('legacy_padded_tensor') - padded_value = 0. if params.type != Params.MAX else 1 + np.finfo(np.float32).min + legacy_padded_tensor_name = scope.get_unique_variable_name( + "legacy_padded_tensor" + ) + padded_value = ( + 0.0 if params.type != Params.MAX else 1 + np.finfo(np.float32).min + ) # Create a sub-graph of cases 3, 6, 7: X ---> Pad ---> X' - create_legacy_pad(scope, inputs[0], legacy_padded_tensor_name, H, W, kernel_shape[0], kernel_shape[1], - strides[0], strides[1], pad_h, pad_w, padded_value, container) - # Set the first input name to the output of Pad so that the following Pool operator won't access the + create_legacy_pad( + scope, + inputs[0], + legacy_padded_tensor_name, + H, + W, + kernel_shape[0], + kernel_shape[1], + strides[0], + strides[1], + pad_h, + pad_w, + padded_value, + container, + ) + # Set the first input name to the output of Pad so that + # the following Pool operator won't access the # original input. inputs[0] = legacy_padded_tensor_name else: - raise ValueError('Unsupported padding mode: {}'.format(pad_type)) + raise ValueError("Unsupported padding mode: {}".format(pad_type)) if pads is not None: - attrs['pads'] = pads + attrs["pads"] = pads if auto_pad is not None: - attrs['auto_pad'] = auto_pad + attrs["auto_pad"] = auto_pad - if (params.type == Params.AVERAGE and params.avgPoolExcludePadding and pad_type == 'includeLastPixel') or \ - (params.type == Params.AVERAGE and not params.avgPoolExcludePadding and pad_type != 'includeLastPixel'): + if ( + params.type == Params.AVERAGE + and params.avgPoolExcludePadding + and pad_type == "includeLastPixel" + ) or ( + params.type == Params.AVERAGE + and not params.avgPoolExcludePadding + and pad_type != "includeLastPixel" + ): # Case 5 & 6. See comment above. # X --> Affine --> Z X_name = operator.inputs[0].full_name Y_name = operator.outputs[0].full_name - Z_name = scope.get_unique_variable_name('Z') - apply_affine(scope, X_name, Z_name, container, alpha=0., beta=1. / (kernel_shape[0] * kernel_shape[1])) + Z_name = scope.get_unique_variable_name("Z") + apply_affine( + scope, + X_name, + Z_name, + container, + alpha=0.0, + beta=1.0 / (kernel_shape[0] * kernel_shape[1]), + ) - Z_prime_name = scope.get_unique_variable_name('Z_prime') - Y_prime_name = scope.get_unique_variable_name('Y_prime') + Z_prime_name = scope.get_unique_variable_name("Z_prime") + Y_prime_name = scope.get_unique_variable_name("Y_prime") - if pad_type != 'includeLastPixel': + if pad_type != "includeLastPixel": # Create the major Pool operator. # Associated sub-graph of case 5: X ---> Pool ---> Y' container.add_node(op_type, inputs, Y_prime_name, **attrs) # Create operators to calculate correction coefficients # Associated sub-graph of case 5: Z ---> L1Pool ---> Z' - lp_pool_attrs = {'name': scope.get_unique_operator_name('LpPool'), 'kernel_shape': kernel_shape, - 'strides': strides, 'p': 1} + lp_pool_attrs = { + "name": scope.get_unique_operator_name("LpPool"), + "kernel_shape": kernel_shape, + "strides": strides, + "p": 1, + } if pads is not None: - lp_pool_attrs['pads'] = pads + lp_pool_attrs["pads"] = pads if auto_pad is not None: - lp_pool_attrs['auto_pad'] = auto_pad - container.add_node('LpPool', Z_name, Z_prime_name, op_version=2, **lp_pool_attrs) + lp_pool_attrs["auto_pad"] = auto_pad + container.add_node( + "LpPool", Z_name, Z_prime_name, op_version=2, **lp_pool_attrs + ) - # Element-wisely apply adjustment coefficients and create the expected CoreML output + # Element-wisely apply adjustment coefficients and + # create the expected CoreML output # Associated sub-graph of case 5: Y', Z' ---> Mul ---> Y - apply_mul(scope, [Y_prime_name, Z_prime_name], Y_name, container, broadcast=0) + apply_mul( + scope, [Y_prime_name, Z_prime_name], Y_name, container, broadcast=0 + ) else: # Create the major Pool operator # Associated sub-graph of case 6: X' ---> Pool ---> Y' @@ -298,26 +393,55 @@ def convert_pooling(scope, operator, container): # Create operators to correct Pool's output Y_name = operator.outputs[0].full_name - Z_prime_prime_name = scope.get_unique_variable_name('Z_prime_prime') + Z_prime_prime_name = scope.get_unique_variable_name("Z_prime_prime") # Pad the constant tensor. # Associated sub-graph of case 6: Z ---> Pad ---> Z' - create_legacy_pad(scope, Z_name, Z_prime_name, operator.inputs[0].type.shape[2], - operator.inputs[0].type.shape[3], kernel_shape[0], kernel_shape[1], - strides[0], strides[1], params.includeLastPixel.paddingAmounts[0], - params.includeLastPixel.paddingAmounts[1], 0., container) + create_legacy_pad( + scope, + Z_name, + Z_prime_name, + operator.inputs[0].type.shape[2], + operator.inputs[0].type.shape[3], + kernel_shape[0], + kernel_shape[1], + strides[0], + strides[1], + params.includeLastPixel.paddingAmounts[0], + params.includeLastPixel.paddingAmounts[1], + 0.0, + container, + ) # Associated sub-graph of case 6: Z' ---> L1Pool ---> Z'' - lp_pool_attrs = {'name': scope.get_unique_operator_name('LpPool'), 'kernel_shape': kernel_shape, - 'strides': strides, 'p': 1, 'auto_pad': 'VALID'} - container.add_node('LpPool', Z_prime_name, Z_prime_prime_name, op_version=2, **lp_pool_attrs) + lp_pool_attrs = { + "name": scope.get_unique_operator_name("LpPool"), + "kernel_shape": kernel_shape, + "strides": strides, + "p": 1, + "auto_pad": "VALID", + } + container.add_node( + "LpPool", + Z_prime_name, + Z_prime_prime_name, + op_version=2, + **lp_pool_attrs + ) - # Element-wisely apply adjustment coefficients and create the expected CoreML output + # Element-wisely apply adjustment coefficients + # and create the expected CoreML output # Associated sub-graph of case 6: Y', Z'' ---> Div ---> Y - apply_div(scope, [Y_prime_name, Z_prime_prime_name], Y_name, container, broadcast=0) + apply_div( + scope, + [Y_prime_name, Z_prime_prime_name], + Y_name, + container, + broadcast=0, + ) else: # Create the major Pool operator container.add_node(op_type, inputs, outputs, op_version=op_version, **attrs) -register_converter('pooling', convert_pooling) +register_converter("pooling", convert_pooling) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Reduce.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Reduce.py index 58052d9f..9e8a44f5 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Reduce.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Reduce.py @@ -5,20 +5,40 @@ def convert_reduce(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import ReduceLayerParams as Params - reduce_mode_table = {Params.SUM: 'ReduceSum', Params.AVG: 'ReduceMean', Params.PROD: 'ReduceProd', - Params.LOGSUM: 'ReduceLogSum', Params.SUMSQUARE: 'ReduceSumSquare', - Params.L1: 'ReduceL1', Params.L2: 'ReduceL2', Params.MAX: 'ReduceMax', - Params.MIN: 'ReduceMin', Params.ARGMAX: 'ArgMax'} + + reduce_mode_table = { + Params.SUM: "ReduceSum", + Params.AVG: "ReduceMean", + Params.PROD: "ReduceProd", + Params.LOGSUM: "ReduceLogSum", + Params.SUMSQUARE: "ReduceSumSquare", + Params.L1: "ReduceL1", + Params.L2: "ReduceL2", + Params.MAX: "ReduceMax", + Params.MIN: "ReduceMin", + Params.ARGMAX: "ArgMax", + } params = operator.raw_operator.reduce reduce_mode = reduce_mode_table[params.mode] reduce_name = scope.get_unique_operator_name(reduce_mode) - # CoreML's reduce operator is used to process tensors with shape [C, H, W]. Notice that [C, H, W] in CoreML - # corresponds to [N, C, H, W] in ONNX because ONNX explicitly get the batch axis. If a CoreML reduce is working - # on CoreML's C-axis, the corresponding ONNX axis's index would be 1 (for the 2nd axis in [N, C, H, W]-system). - reduce_axis_table = {Params.CHW: [1, 2, 3], Params.HW: [2, 3], Params.C: [1], Params.H: [2], Params.W: [3]} + # CoreML's reduce operator is used to process tensors + # with shape [C, H, W]. Notice that [C, H, W] in CoreML + # corresponds to [N, C, H, W] in ONNX because ONNX explicitly + # get the batch axis. If a CoreML reduce is working + # on CoreML's C-axis, the corresponding ONNX axis's + # index would be 1 (for the 2nd axis in [N, C, H, W]-system). + reduce_axis_table = { + Params.CHW: [1, 2, 3], + Params.HW: [2, 3], + Params.C: [1], + Params.H: [2], + Params.W: [3], + } reduce_axis = reduce_axis_table[params.axis] - attrs = {'name': reduce_name, 'axes': reduce_axis} - container.add_node(reduce_mode, operator.input_full_names, operator.output_full_names, **attrs) + attrs = {"name": reduce_name, "axes": reduce_axis} + container.add_node( + reduce_mode, operator.input_full_names, operator.output_full_names, **attrs + ) -register_converter('reduce', convert_reduce) +register_converter("reduce", convert_reduce) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/ReorganizeData.py b/onnxmltools/convert/coreml/operator_converters/neural_network/ReorganizeData.py index 7a042922..34e2f5dc 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/ReorganizeData.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/ReorganizeData.py @@ -5,17 +5,20 @@ def convert_reorganize_data(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import ReorganizeDataLayerParams as Params + params = operator.raw_operator.reorganizeData if params.mode == Params.DEPTH_TO_SPACE: - op_type = 'DepthToSpace' + op_type = "DepthToSpace" elif params.mode == Params.SPACE_TO_DEPTH: - op_type = 'SpaceToDepth' + op_type = "SpaceToDepth" else: - raise ValueError('Unsupported reorganization mode {0}'.format(params.mode)) + raise ValueError("Unsupported reorganization mode {0}".format(params.mode)) op_name = scope.get_unique_operator_name(op_type) - attrs = {'name': op_name, 'blocksize': params.blockSize} - container.add_node(op_type, operator.input_full_names, operator.output_full_names, **attrs) + attrs = {"name": op_name, "blocksize": params.blockSize} + container.add_node( + op_type, operator.input_full_names, operator.output_full_names, **attrs + ) -register_converter('reorganizeData', convert_reorganize_data) +register_converter("reorganizeData", convert_reorganize_data) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Reshape.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Reshape.py index a271493c..601bd655 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Reshape.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Reshape.py @@ -10,13 +10,21 @@ def convert_reshape(scope, operator, container): params = operator.raw_operator.reshape if params.mode == Params.CHANNEL_LAST: - intra_variable_name = scope.get_unique_variable_name(operator.inputs[0].full_name + '_transpose') - apply_transpose(scope, operator.inputs[0].full_name, intra_variable_name, container, perm=[0, 2, 3, 1]) + intra_variable_name = scope.get_unique_variable_name( + operator.inputs[0].full_name + "_transpose" + ) + apply_transpose( + scope, + operator.inputs[0].full_name, + intra_variable_name, + container, + perm=[0, 2, 3, 1], + ) else: intra_variable_name = operator.inputs[0].full_name N = operator.inputs[0].type.shape[0] - if N == 'None': + if N == "None": N = -1 if len(params.targetShape) == 4: output_shape = [int(d) for d in params.targetShape] @@ -24,11 +32,20 @@ def convert_reshape(scope, operator, container): elif len(params.targetShape) == 3: output_shape = [N] + [int(d) for d in params.targetShape] else: - raise ValueError('The targeted shape of Reshape (name: %s) must be 3-element or 4-element array but got %s'\ - % (operator.full_name, params.targetShape)) - - apply_reshape(scope=scope, input_name=intra_variable_name, output_name=operator.outputs[0].full_name, - container=container, operator_name=operator.full_name, desired_shape=output_shape) - - -register_converter('reshape', convert_reshape) + raise ValueError( + "The targeted shape of Reshape (name: %s) " + "must be 3-element or 4-element array but got %s" + % (operator.full_name, params.targetShape) + ) + + apply_reshape( + scope=scope, + input_name=intra_variable_name, + output_name=operator.outputs[0].full_name, + container=container, + operator_name=operator.full_name, + desired_shape=output_shape, + ) + + +register_converter("reshape", convert_reshape) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/ReshapeStatic.py b/onnxmltools/convert/coreml/operator_converters/neural_network/ReshapeStatic.py index 047597be..e3ff7e6b 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/ReshapeStatic.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/ReshapeStatic.py @@ -5,7 +5,7 @@ def convert_reshape_static(scope, operator, container): - from coremltools.proto.NeuralNetwork_pb2 import ReshapeLayerParams as Params + pass params = operator.raw_operator.reshapeStatic @@ -13,7 +13,7 @@ def convert_reshape_static(scope, operator, container): intra_variable_name = operator.inputs[0].full_name N = operator.inputs[0].type.shape[0] - if N == 'None': + if N == "None": N = -1 if len(params.targetShape) == 4: output_shape = [int(d) for d in params.targetShape] @@ -23,11 +23,20 @@ def convert_reshape_static(scope, operator, container): elif len(params.targetShape) == 2: output_shape = [N] + [int(d) for d in params.targetShape] else: - raise ValueError('The targeted shape of Reshape (name: %s) must be 3-element or 4-element array but got %s'\ - % (operator.full_name, params.targetShape)) - - apply_reshape(scope=scope, input_name=intra_variable_name, output_name=operator.outputs[0].full_name, - container=container, operator_name=operator.full_name, desired_shape=output_shape) - - -register_converter('reshapeStatic', convert_reshape_static) + raise ValueError( + "The targeted shape of Reshape (name: %s) " + "must be 3-element or 4-element array but got %s" + % (operator.full_name, params.targetShape) + ) + + apply_reshape( + scope=scope, + input_name=intra_variable_name, + output_name=operator.outputs[0].full_name, + container=container, + operator_name=operator.full_name, + desired_shape=output_shape, + ) + + +register_converter("reshapeStatic", convert_reshape_static) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Scale.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Scale.py index 109a60bd..8070f8af 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Scale.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Scale.py @@ -6,12 +6,18 @@ def deduce_broadcast_axis_and_shape(target_opset, shape): - # This function is used to calculate the first axis aligned with the scalar and the scalar's ONNX shape for reduce- - # like operators. Assuming input variable is always a 4-D tensor, we provide a few of examples. If scalar's shape - # is [1, 2, 3] and input shape is [5, 2, 3, 8], the aligned axis is the [2] (indexed by 1 because indexes are 0-based) - # in [5, 2, 3, 8], and the desired scalar shape in ONNX is [2, 3] # (the [1] in [1, 2, 3] is redundant and can cause - # errors in ONNX's boardcasting). If the scaler's shape is [1], no matter what shape the input is, we leave the axis - # "None" because ONNX operator may automatically handle it. After ONNX-1.2, we adopt Numpy-style broadcasting rule. + # This function is used to calculate the first axis aligned + # with the scalar and the scalar's ONNX shape for reduce- + # like operators. Assuming input variable is always a 4-D tensor, + # we provide a few of examples. If scalar's shape + # is [1, 2, 3] and input shape is [5, 2, 3, 8], the aligned axis + # is the [2] (indexed by 1 because indexes are 0-based) + # in [5, 2, 3, 8], and the desired scalar shape in ONNX is [2, 3] + # (the [1] in [1, 2, 3] is redundant and can cause + # errors in ONNX's boardcasting). If the scaler's shape is [1], + # no matter what shape the input is, we leave the axis + # "None" because ONNX operator may automatically handle it. + # After ONNX-1.2, we adopt Numpy-style broadcasting rule. if target_opset < 7: # Input shape is [N, C, H, W] @@ -48,7 +54,8 @@ def deduce_broadcast_axis_and_shape(target_opset, shape): def convert_scale(scope, operator, container): - # In CoreML's ScaleLayer, the input is first scaled by their "scale" attribute and then a "bias" can be added. + # In CoreML's ScaleLayer, the input is first scaled by their + # "scale" attribute and then a "bias" can be added. # Symbols: # a: scale attribute in CoreML's ScaleLayer # b: bias attribute in CoreML's ScaleLayer @@ -56,40 +63,74 @@ def convert_scale(scope, operator, container): # y: output # The math formulation of ScaleLayer should be # y = a * x + b - # Therefore, our strategy of composing ScaleLayer is to have one multiplication followed by an addition. + # Therefore, our strategy of composing ScaleLayer is to + # have one multiplication followed by an addition. params = operator.raw_operator.scale - op1_type = 'Mul' - scale_axis, scale_shape = deduce_broadcast_axis_and_shape(container.target_opset, params.shapeScale) - scale_name = scope.get_unique_variable_name(op1_type + '_B') - container.add_initializer(scale_name, onnx_proto.TensorProto.FLOAT, scale_shape, params.scale.floatValue) + op1_type = "Mul" + scale_axis, scale_shape = deduce_broadcast_axis_and_shape( + container.target_opset, params.shapeScale + ) + scale_name = scope.get_unique_variable_name(op1_type + "_B") + container.add_initializer( + scale_name, onnx_proto.TensorProto.FLOAT, scale_shape, params.scale.floatValue + ) # CoreML is at most 3-D, so we always turn broadcasting on. scale_broadcast = 1 if not params.hasBias: - # Create a element-wise multiplication and use it to scale the input. The first input is the variable we want + # Create a element-wise multiplication and use it to scale + # the input. The first input is the variable we want # to scale while the second input is their multipliers. - apply_mul(scope, [operator.inputs[0].full_name, scale_name], operator.output_full_names, container, - operator_name=operator.full_name, axis=scale_axis, broadcast=scale_broadcast) + apply_mul( + scope, + [operator.inputs[0].full_name, scale_name], + operator.output_full_names, + container, + operator_name=operator.full_name, + axis=scale_axis, + broadcast=scale_broadcast, + ) else: # Declare a temporal variable to store the scaled input - intra_variable_name = scope.get_unique_variable_name(operator.inputs[0].full_name + '_scaled') - # Create a element-wise multiplication and use it to scale the input and save the result to a temporal variable - apply_mul(scope, [operator.inputs[0].full_name, scale_name], intra_variable_name, container, - operator_name=operator.full_name, axis=scale_axis, broadcast=scale_broadcast) + intra_variable_name = scope.get_unique_variable_name( + operator.inputs[0].full_name + "_scaled" + ) + # Create a element-wise multiplication and use it to + # scale the input and save the result to a temporal variable + apply_mul( + scope, + [operator.inputs[0].full_name, scale_name], + intra_variable_name, + container, + operator_name=operator.full_name, + axis=scale_axis, + broadcast=scale_broadcast, + ) # Prepare materials to build an Add operator for adding bias - bias_axis, bias_shape = deduce_broadcast_axis_and_shape(container.target_opset, params.shapeBias) + bias_axis, bias_shape = deduce_broadcast_axis_and_shape( + container.target_opset, params.shapeBias + ) # CoreML is at most 3-D, so we always turn broadcasting on. bias_broadcast = 1 - bias_name = scope.get_unique_variable_name(operator.full_name + '_B') - container.add_initializer(bias_name, onnx_proto.TensorProto.FLOAT, bias_shape, params.bias.floatValue) - # As bias exists, we add the bias into the output of the multiplication and then use the output of addition + bias_name = scope.get_unique_variable_name(operator.full_name + "_B") + container.add_initializer( + bias_name, onnx_proto.TensorProto.FLOAT, bias_shape, params.bias.floatValue + ) + # As bias exists, we add the bias into the output of + # the multiplication and then use the output of addition # as the final output of this conversion. - apply_add(scope, [intra_variable_name, bias_name], operator.output_full_names, container, - axis=bias_axis, broadcast=bias_broadcast) + apply_add( + scope, + [intra_variable_name, bias_name], + operator.output_full_names, + container, + axis=bias_axis, + broadcast=bias_broadcast, + ) -register_converter('scale', convert_scale) +register_converter("scale", convert_scale) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/SequenceRepeat.py b/onnxmltools/convert/coreml/operator_converters/neural_network/SequenceRepeat.py index 1ca56490..70344a2d 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/SequenceRepeat.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/SequenceRepeat.py @@ -5,7 +5,9 @@ def convert_sequence_repeat(scope, operator, container): - repeat_count = operator.raw_operator.sequenceRepeat.nRepetitions # number of copies along N-axis in CoreML + repeat_count = ( + operator.raw_operator.sequenceRepeat.nRepetitions + ) # number of copies along N-axis in CoreML if len(operator.inputs[0].type.shape) == 4: # Number of copies along [N, C, H, W] @@ -14,8 +16,14 @@ def convert_sequence_repeat(scope, operator, container): # Number of copies along [N, C] repeats = [int(repeat_count), 1] - apply_tile(scope, operator.input_full_names[0], operator.output_full_names[0], container, - operator_name=operator.full_name, repeats=repeats) + apply_tile( + scope, + operator.input_full_names[0], + operator.output_full_names[0], + container, + operator_name=operator.full_name, + repeats=repeats, + ) -register_converter('sequenceRepeat', convert_sequence_repeat) +register_converter("sequenceRepeat", convert_sequence_repeat) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/SimpleRNN.py b/onnxmltools/convert/coreml/operator_converters/neural_network/SimpleRNN.py index 2f1633a1..a8eb4fde 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/SimpleRNN.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/SimpleRNN.py @@ -7,58 +7,64 @@ def extract_rnn_activation_info(activation): - activation_type = activation.WhichOneof('NonlinearityType') + activation_type = activation.WhichOneof("NonlinearityType") alpha = None beta = None - activation_map = {'linear': 'Affine', - 'ReLU': 'Relu', - 'leakyReLU': 'LeakyRelu', - 'thresholdedReLU': 'ThresholdedRelu', - 'PReLU': 'PRelu', - 'tanh': 'Tanh', - 'scaledTanh': 'ScaledTanh', - 'sigmoid': 'Sigmoid', - 'sigmoidHard': 'HardSigmoid', - 'ELU': 'Elu', - 'softsign': 'Softsign', - 'softplus': 'Softplus', - 'parametricSoftplus': 'ParametricSoftplus'} + activation_map = { + "linear": "Affine", + "ReLU": "Relu", + "leakyReLU": "LeakyRelu", + "thresholdedReLU": "ThresholdedRelu", + "PReLU": "PRelu", + "tanh": "Tanh", + "scaledTanh": "ScaledTanh", + "sigmoid": "Sigmoid", + "sigmoidHard": "HardSigmoid", + "ELU": "Elu", + "softsign": "Softsign", + "softplus": "Softplus", + "parametricSoftplus": "ParametricSoftplus", + } if activation_type not in activation_map: - raise ValueError('Unsupported activation function: {}'.format(activation_type)) + raise ValueError("Unsupported activation function: {}".format(activation_type)) # Notice that if we see a default vaulue (i.e., 0 for float), we may replace it with # the real default parameter for the specified activation function if necessary. - if activation_type == 'leakyReLU': + if activation_type == "leakyReLU": alpha = activation.leakyReLU.alpha if alpha == 0: alpha = 0.3 - elif activation_type == 'PReLU': - raise RuntimeError('Unsupported activation function: {}'.format(activation_type)) - elif activation_type == 'ELU': + elif activation_type == "PReLU": + raise RuntimeError( + "Unsupported activation function: {}".format(activation_type) + ) + elif activation_type == "ELU": alpha = activation.ELU.alpha - elif activation_type == 'thresholdedReLU': + elif activation_type == "thresholdedReLU": alpha = activation.thresholdedReLU.alpha if alpha == 0: alpha = 1.0 - elif activation_type == 'scaledTanh': + elif activation_type == "scaledTanh": alpha = activation.scaledTanh.alpha beta = activation.scaledTanh.beta - elif activation_type == 'linear': + elif activation_type == "linear": alpha = activation.linear.alpha beta = activation.linear.beta if alpha == 0: alpha = 1.0 - elif activation_type == 'sigmoidHard': + elif activation_type == "sigmoidHard": alpha = activation.sigmoidHard.alpha beta = activation.sigmoidHard.beta if alpha == 0: alpha = 0.2 if beta == 0: beta = 0.5 - elif activation_type == 'parametricSoftplus': - raise RuntimeError('Unsupported activation function: {}'.format(activation_type)) + elif activation_type == "parametricSoftplus": + raise RuntimeError( + "Unsupported activation function: {}".format(activation_type) + ) return activation_map[activation_type], alpha, beta @@ -125,98 +131,158 @@ def convert_simple_rnn(scope, operator, container): hidden_size = params.outputVectorSize X_name = operator.inputs[0].full_name - X_reshape_name = scope.get_unique_variable_name('X') - apply_reshape(scope, X_name, X_reshape_name, container, desired_shape=[-1, 1, input_size]) + X_reshape_name = scope.get_unique_variable_name("X") + apply_reshape( + scope, X_name, X_reshape_name, container, desired_shape=[-1, 1, input_size] + ) - rnn_op_name = scope.get_unique_operator_name('RNN') - rnn_attrs = {'name': rnn_op_name} + rnn_op_name = scope.get_unique_operator_name("RNN") + rnn_attrs = {"name": rnn_op_name} rnn_inputs = [X_reshape_name] # Load RNN's weight matrix and add it into RNN's input list - rnn_w_name = scope.get_unique_variable_name(rnn_op_name + '_W') - container.add_initializer(rnn_w_name, onnx_proto.TensorProto.FLOAT, - [1, hidden_size, input_size], params.weightMatrix.floatValue) + rnn_w_name = scope.get_unique_variable_name(rnn_op_name + "_W") + container.add_initializer( + rnn_w_name, + onnx_proto.TensorProto.FLOAT, + [1, hidden_size, input_size], + params.weightMatrix.floatValue, + ) rnn_inputs.append(rnn_w_name) # Load RNN's recursion matrix and add it into RNN's input list - rnn_r_name = scope.get_unique_variable_name(rnn_op_name + '_R') - container.add_initializer(rnn_r_name, onnx_proto.TensorProto.FLOAT, - [1, hidden_size, hidden_size], params.recursionMatrix.floatValue) + rnn_r_name = scope.get_unique_variable_name(rnn_op_name + "_R") + container.add_initializer( + rnn_r_name, + onnx_proto.TensorProto.FLOAT, + [1, hidden_size, hidden_size], + params.recursionMatrix.floatValue, + ) rnn_inputs.append(rnn_r_name) if params.hasBiasVector: # Load RNN's bias vector and add it into RNN's input list - rnn_b_name = scope.get_unique_variable_name(rnn_op_name + '_B') - rnn_b_content = np.concatenate([params.biasVector.floatValue, np.zeros(hidden_size)]).flatten() - container.add_initializer(rnn_b_name, onnx_proto.TensorProto.FLOAT, [1, 2 * hidden_size], rnn_b_content) + rnn_b_name = scope.get_unique_variable_name(rnn_op_name + "_B") + rnn_b_content = np.concatenate( + [params.biasVector.floatValue, np.zeros(hidden_size)] + ).flatten() + container.add_initializer( + rnn_b_name, + onnx_proto.TensorProto.FLOAT, + [1, 2 * hidden_size], + rnn_b_content, + ) rnn_inputs.append(rnn_b_name) else: - # Input names are position-sensitive, so for optional but missing inputs, we need to provide an empty string. - rnn_inputs.append('') + # Input names are position-sensitive, so for optional + # but missing inputs, we need to provide an empty string. + rnn_inputs.append("") - # The input, sequence_lens, in ONNX is alwasy optional for this conversion, so here is always an empty string. - rnn_inputs.append('') + # The input, sequence_lens, in ONNX is alwasy optional for + # this conversion, so here is always an empty string. + rnn_inputs.append("") - # If initial hidden state is provided, we add it into RNN's input list after adjusting its shape. + # If initial hidden state is provided, we add it into RNN's + # input list after adjusting its shape. if len(operator.inputs) == 2: - rnn_h_init_reshape_name = scope.get_unique_variable_name(rnn_op_name + '_h_init') - apply_reshape(scope, operator.inputs[1].full_name, rnn_h_init_reshape_name, container, - desired_shape=[1, 1, hidden_size]) + rnn_h_init_reshape_name = scope.get_unique_variable_name( + rnn_op_name + "_h_init" + ) + apply_reshape( + scope, + operator.inputs[1].full_name, + rnn_h_init_reshape_name, + container, + desired_shape=[1, 1, hidden_size], + ) rnn_inputs.append(rnn_h_init_reshape_name) - # Add a zero initializer to initial hidden state so that this variable becomes optional - container.add_initializer(operator.inputs[1].full_name, onnx_proto.TensorProto.FLOAT, - operator.inputs[1].type.shape, - np.zeros(shape=operator.inputs[1].type.shape).flatten()) + # Add a zero initializer to initial hidden state so that + # this variable becomes optional + container.add_initializer( + operator.inputs[1].full_name, + onnx_proto.TensorProto.FLOAT, + operator.inputs[1].type.shape, + np.zeros(shape=operator.inputs[1].type.shape).flatten(), + ) else: - # Input names are position-sensitive, so for optional but missing inputs, we need to provide an empty string. - rnn_inputs.append('') + # Input names are position-sensitive, so for optional + # but missing inputs, we need to provide an empty string. + rnn_inputs.append("") # Add RNN's information of activation function activation, alpha, beta = extract_rnn_activation_info(params.activation) - rnn_attrs['activations'] = [activation.encode('utf-8')] + rnn_attrs["activations"] = [activation.encode("utf-8")] if alpha is not None: - rnn_attrs['activation_alpha'] = [alpha] + rnn_attrs["activation_alpha"] = [alpha] if beta is not None: - rnn_attrs['activation_beta'] = [beta] + rnn_attrs["activation_beta"] = [beta] # Set up other attributes - rnn_attrs['direction'] = 'reverse' if params.reverseInput else 'forward' - rnn_attrs['hidden_size'] = hidden_size + rnn_attrs["direction"] = "reverse" if params.reverseInput else "forward" + rnn_attrs["hidden_size"] = hidden_size # Set up version-dependent attributes if container.target_opset < 7: - rnn_attrs['output_sequence'] = params.sequenceOutput + rnn_attrs["output_sequence"] = params.sequenceOutput op_version = 1 else: op_version = 7 - # We use the collected information to build ONNX's RNN. ONNX RNN's outputs will be saved onto two intermediate - # tensors and we will adjust them subsequently to mimic Keras output format. - rnn_y_name = scope.get_unique_variable_name(rnn_op_name + '_Y') - rnn_h_name = scope.get_unique_variable_name(rnn_op_name + '_Y_h') - container.add_node('RNN', rnn_inputs, [rnn_y_name, rnn_h_name], op_version=op_version, **rnn_attrs) + # We use the collected information to build ONNX's RNN. + # ONNX RNN's outputs will be saved onto two intermediate + # tensors and we will adjust them subsequently to mimic + # Keras output format. + rnn_y_name = scope.get_unique_variable_name(rnn_op_name + "_Y") + rnn_h_name = scope.get_unique_variable_name(rnn_op_name + "_Y_h") + container.add_node( + "RNN", rnn_inputs, [rnn_y_name, rnn_h_name], op_version=op_version, **rnn_attrs + ) # Set up outputs' of RNN if params.sequenceOutput: # Connect ONNX's output and CoreML's output via a reshape operator - apply_reshape(scope, rnn_y_name, operator.outputs[0].full_name, container, desired_shape=[-1, hidden_size]) + apply_reshape( + scope, + rnn_y_name, + operator.outputs[0].full_name, + container, + desired_shape=[-1, hidden_size], + ) # Handel the second RNN output (aka last hidden state), which is optional. if len(operator.outputs) == 2: # Connect ONNX's output and CoreML's output via a reshape operator - apply_reshape(scope, rnn_h_name, operator.outputs[1].full_name, container, desired_shape=[1, hidden_size]) + apply_reshape( + scope, + rnn_h_name, + operator.outputs[1].full_name, + container, + desired_shape=[1, hidden_size], + ) else: - # According to CoreML, its two outputs are always identical, so we just need to compute one of them and produce - # the other one using an identity operator. Note that the first ONNX RNN output is undefined in this case. + # According to CoreML, its two outputs are always identical, so + # we just need to compute one of them and produce + # the other one using an identity operator. Note that the + # first ONNX RNN output is undefined in this case. # Reshape last hidden state's ONNX format to its CoreML format - apply_reshape(scope, rnn_h_name, operator.outputs[0].full_name, container, desired_shape=[1, hidden_size]) + apply_reshape( + scope, + rnn_h_name, + operator.outputs[0].full_name, + container, + desired_shape=[1, hidden_size], + ) if len(operator.outputs) == 2: # Copy the first output to the second output - container.add_node('Identity', operator.outputs[0].full_name, operator.outputs[1].full_name, - name=scope.get_unique_operator_name('Identity')) + container.add_node( + "Identity", + operator.outputs[0].full_name, + operator.outputs[1].full_name, + name=scope.get_unique_operator_name("Identity"), + ) -register_converter('simpleRecurrent', convert_simple_rnn) +register_converter("simpleRecurrent", convert_simple_rnn) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Slice.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Slice.py index 8e3d8bcd..04be6e6b 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Slice.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Slice.py @@ -7,12 +7,13 @@ def convert_slice(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import SliceLayerParams as Params - op_type = 'Slice' + op_type = "Slice" op_name = scope.get_unique_operator_name(op_type) - attrs = {'name': op_name} + attrs = {"name": op_name} params = operator.raw_operator.slice - # Set up slice range of C-, H-, and W-axes. Notice that only one of them will be actually sliced. + # Set up slice range of C-, H-, and W-axes. Notice that + # only one of them will be actually sliced. axis_map = {Params.CHANNEL_AXIS: 0, Params.HEIGHT_AXIS: 1, Params.WIDTH_AXIS: 2} starts = [0, 0, 0] ends = [-1, -1, -1] @@ -20,34 +21,48 @@ def convert_slice(scope, operator, container): ends[axis_map[params.axis]] = params.endIndex if params.stride != 1: - raise ValueError('Stride must be 1 but got %s' % params.stride) + raise ValueError("Stride must be 1 but got %s" % params.stride) if container.target_opset < 10: - # The input shape should be [N, C, H, W] in ONNX. Because CoreML only slices one of C-, H-, or W-axes, the - # "axes" attribute in ONNX is [1, 2, 3]. Note that for the axes not really sliced, their starting and ending + # The input shape should be [N, C, H, W] in ONNX. + # Because CoreML only slices one of C-, H-, or W-axes, the + # "axes" attribute in ONNX is [1, 2, 3]. Note that for the + # axes not really sliced, their starting and ending # indexes are 0 and -1, respectively. - attrs['axes'] = [1, 2, 3] - attrs['starts'] = starts - attrs['ends'] = ends + attrs["axes"] = [1, 2, 3] + attrs["starts"] = starts + attrs["ends"] = ends op_version = 1 - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_version=op_version, **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_version=op_version, + **attrs + ) else: - starts_name = scope.get_unique_variable_name(operator.full_name + '_starts') - container.add_initializer(starts_name, onnx_proto.TensorProto.INT64, - [3], starts) - ends_name = scope.get_unique_variable_name(operator.full_name + '_ends') - container.add_initializer(ends_name, onnx_proto.TensorProto.INT64, - [3], ends) - axes_name = scope.get_unique_variable_name(operator.full_name + '_axes') - container.add_initializer(axes_name, onnx_proto.TensorProto.INT64, - [3], [1, 2, 3]) + starts_name = scope.get_unique_variable_name(operator.full_name + "_starts") + container.add_initializer( + starts_name, onnx_proto.TensorProto.INT64, [3], starts + ) + ends_name = scope.get_unique_variable_name(operator.full_name + "_ends") + container.add_initializer(ends_name, onnx_proto.TensorProto.INT64, [3], ends) + axes_name = scope.get_unique_variable_name(operator.full_name + "_axes") + container.add_initializer( + axes_name, onnx_proto.TensorProto.INT64, [3], [1, 2, 3] + ) if container.target_opset < 11: op_version = 10 else: op_version = 11 - container.add_node(op_type, [operator.input_full_names[0], starts_name, ends_name, axes_name], operator.output_full_names, op_version=op_version, - **attrs) + container.add_node( + op_type, + [operator.input_full_names[0], starts_name, ends_name, axes_name], + operator.output_full_names, + op_version=op_version, + **attrs + ) -register_converter('slice', convert_slice) +register_converter("slice", convert_slice) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Softmax.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Softmax.py index 0f743ed8..8c9987f6 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Softmax.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Softmax.py @@ -4,11 +4,11 @@ def convert_softmax(scope, operator, container): - op_type = 'Softmax' + op_type = "Softmax" inputs = [variable.full_name for variable in operator.inputs] outputs = [variable.full_name for variable in operator.outputs] - attrs = {'name': operator.full_name} + attrs = {"name": operator.full_name} container.add_node(op_type, inputs, outputs, **attrs) -register_converter('softmax', convert_softmax) +register_converter("softmax", convert_softmax) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Split.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Split.py index 2e6bafc3..122bf084 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Split.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Split.py @@ -5,11 +5,19 @@ def convert_split(scope, operator, container): - # ONNX Split may evenly divide the input along the specified axis if "split" attribute is not specified. - # Also, CoreML always evenly split the input along C-axis. Consequently, we only need to specify the axis + # ONNX Split may evenly divide the input along the specified + # axis if "split" attribute is not specified. + # Also, CoreML always evenly split the input along C-axis. + # Consequently, we only need to specify the axis # and make sure the number of outputs in ONNX matches that in CoreML. - apply_split(scope, operator.input_full_names, operator.output_full_names, container, - operator_name=operator.full_name, axis=1) + apply_split( + scope, + operator.input_full_names, + operator.output_full_names, + container, + operator_name=operator.full_name, + axis=1, + ) -register_converter('split', convert_split) +register_converter("split", convert_split) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/UnaryFunction.py b/onnxmltools/convert/coreml/operator_converters/neural_network/UnaryFunction.py index 5b63e9af..261d2dd2 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/UnaryFunction.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/UnaryFunction.py @@ -1,6 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -from ....common._apply_operation import * +import onnx +from ....common._apply_operation import ( + apply_abs, + apply_log, + apply_exp, + apply_reciprocal, + apply_sqrt, + apply_clip, + apply_pow, + apply_affine, +) from ....common._registration import register_converter @@ -8,41 +18,89 @@ def convert_unary(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import UnaryFunctionLayerParams as Params params = operator.raw_operator.unary - preprocessor_name = scope.get_unique_operator_name('Affine') - preprocessed_variable_name = scope.get_unique_variable_name(preprocessor_name + '_output') - apply_affine(scope, operator.input_full_names[0], preprocessed_variable_name, container, - operator_name=preprocessor_name, alpha=params.scale, beta=params.shift) + preprocessor_name = scope.get_unique_operator_name("Affine") + preprocessed_variable_name = scope.get_unique_variable_name( + preprocessor_name + "_output" + ) + apply_affine( + scope, + operator.input_full_names[0], + preprocessed_variable_name, + container, + operator_name=preprocessor_name, + alpha=params.scale, + beta=params.shift, + ) if params.type == Params.RSQRT: - sqrt_tensor_name = scope.get_unique_variable_name(operator.full_name + '_intra_tensor') + sqrt_tensor_name = scope.get_unique_variable_name( + operator.full_name + "_intra_tensor" + ) apply_sqrt(scope, preprocessed_variable_name, sqrt_tensor_name, container) apply_reciprocal(scope, sqrt_tensor_name, operator.output_full_names, container) elif params.type == Params.POWER: - exp_name = scope.get_unique_variable_name('Y') - container.add_initializer(exp_name, onnx_proto.TensorProto.FLOAT, [], [params.alpha]) + exp_name = scope.get_unique_variable_name("Y") + container.add_initializer(exp_name, onnx.TensorProto.FLOAT, [], [params.alpha]) - apply_pow(scope, [preprocessed_variable_name, exp_name], operator.output_full_names, container, - operator_name=operator.full_name, broadcast=1) + apply_pow( + scope, + [preprocessed_variable_name, exp_name], + operator.output_full_names, + container, + operator_name=operator.full_name, + broadcast=1, + ) elif params.type == Params.THRESHOLD: - apply_clip(scope, preprocessed_variable_name, operator.output_full_names, container, - operator_name=operator.full_name, min=params.alpha) + apply_clip( + scope, + preprocessed_variable_name, + operator.output_full_names, + container, + operator_name=operator.full_name, + min=params.alpha, + ) elif params.type == Params.SQRT: - apply_sqrt(scope, preprocessed_variable_name, operator.output_full_names, container, - operator_name=operator.full_name) + apply_sqrt( + scope, + preprocessed_variable_name, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) elif params.type == Params.INVERSE: - apply_reciprocal(scope, preprocessed_variable_name, operator.output_full_names, container, - operator_name=operator.full_name) + apply_reciprocal( + scope, + preprocessed_variable_name, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) elif params.type == Params.EXP: - apply_exp(scope, preprocessed_variable_name, operator.output_full_names, container, - operator_name=operator.full_name) + apply_exp( + scope, + preprocessed_variable_name, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) elif params.type == Params.LOG: - apply_log(scope, preprocessed_variable_name, operator.output_full_names, container, - operator_name=operator.full_name) + apply_log( + scope, + preprocessed_variable_name, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) elif params.type == Params.ABS: - apply_abs(scope, preprocessed_variable_name, operator.output_full_names, container, - operator_name=operator.full_name) + apply_abs( + scope, + preprocessed_variable_name, + operator.output_full_names, + container, + operator_name=operator.full_name, + ) else: - raise ValueError('Unsupported unary function :{}'.format(params.type)) + raise ValueError("Unsupported unary function :{}".format(params.type)) -register_converter('unary', convert_unary) +register_converter("unary", convert_unary) diff --git a/onnxmltools/convert/coreml/operator_converters/neural_network/Upsample.py b/onnxmltools/convert/coreml/operator_converters/neural_network/Upsample.py index bb8852bd..54f7ad50 100644 --- a/onnxmltools/convert/coreml/operator_converters/neural_network/Upsample.py +++ b/onnxmltools/convert/coreml/operator_converters/neural_network/Upsample.py @@ -6,20 +6,28 @@ def convert_upsample(scope, operator, container): from coremltools.proto.NeuralNetwork_pb2 import UpsampleLayerParams as Params + params = operator.raw_operator.upsample if params.mode == Params.NN: - mode = 'NEAREST' + mode = "NEAREST" elif params.mode == Params.BILINEAR: - mode = 'BILINEAR' + mode = "BILINEAR" else: - raise ValueError('Unsupported interpolation model in up-sampling') + raise ValueError("Unsupported interpolation model in up-sampling") width_scale = float(params.scalingFactor[1]) height_scale = float(params.scalingFactor[0]) - apply_upsample(scope, operator.input_full_names[0], operator.output_full_names, container, operator_name=None, - mode=mode, scales=[1, 1, height_scale, width_scale]) + apply_upsample( + scope, + operator.input_full_names[0], + operator.output_full_names, + container, + operator_name=None, + mode=mode, + scales=[1, 1, height_scale, width_scale], + ) -register_converter('upsample', convert_upsample) +register_converter("upsample", convert_upsample) diff --git a/onnxmltools/convert/coreml/shape_calculators/ArrayFeatureExtractor.py b/onnxmltools/convert/coreml/shape_calculators/ArrayFeatureExtractor.py index 56662b90..a179dff8 100644 --- a/onnxmltools/convert/coreml/shape_calculators/ArrayFeatureExtractor.py +++ b/onnxmltools/convert/coreml/shape_calculators/ArrayFeatureExtractor.py @@ -7,17 +7,21 @@ def calculate_array_feature_extractor_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N, C'] C' is the number of extracted features. - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType, StringTensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType, StringTensorType] + ) N = operator.inputs[0].type.shape[0] - extracted_feature_number = len(operator.raw_operator.arrayFeatureExtractor.extractIndex) + extracted_feature_number = len( + operator.raw_operator.arrayFeatureExtractor.extractIndex + ) # Save doc_string before over-writing by us doc_string = operator.outputs[0].type.doc_string @@ -27,4 +31,6 @@ def calculate_array_feature_extractor_output_shapes(operator): operator.outputs[0].type.doc_string = doc_string -register_shape_calculator('arrayFeatureExtractor', calculate_array_feature_extractor_output_shapes) +register_shape_calculator( + "arrayFeatureExtractor", calculate_array_feature_extractor_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/Classifier.py b/onnxmltools/convert/coreml/shape_calculators/Classifier.py index 75dbdae6..2be0efe2 100644 --- a/onnxmltools/convert/coreml/shape_calculators/Classifier.py +++ b/onnxmltools/convert/coreml/shape_calculators/Classifier.py @@ -1,13 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 from ...common._registration import register_shape_calculator -from ...common.data_types import FloatTensorType, Int64TensorType, FloatType, Int64Type, DictionaryType, StringType, \ - SequenceType, StringTensorType +from ...common.data_types import ( + FloatTensorType, + Int64TensorType, + FloatType, + Int64Type, + DictionaryType, + SequenceType, + StringTensorType, +) from ...common.utils import check_input_and_output_numbers, check_input_and_output_types def calculate_traditional_classifier_output_shapes(operator): - ''' + """ For classifiers, allowed input/output patterns are 1. [N, C_1], ..., [N, C_n] ---> [N], Sequence of Map 2. [N, C_1], ..., [N, C_n] ---> [N] @@ -15,24 +22,35 @@ def calculate_traditional_classifier_output_shapes(operator): For regressors, allowed input/output patterns are 1. [N, C_1], ..., [N, C_n] ---> [N, 1] - Core ML classifiers and regressors support multiple input feature tensors, so we need to concatenate them before - feeding them into their ONNX counterparts. Note that the N must be 1 as long as ZipMap only produces dictionary. - ''' - check_input_and_output_numbers(operator, input_count_range=[1, None], output_count_range=[1, 2]) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType, FloatType, Int64Type]) + Core ML classifiers and regressors support multiple + input feature tensors, so we need to concatenate them before + feeding them into their ONNX counterparts. + Note that the N must be 1 as long as ZipMap only produces dictionary. + """ + check_input_and_output_numbers( + operator, input_count_range=[1, None], output_count_range=[1, 2] + ) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType, Int64TensorType, FloatType, Int64Type], + ) if any(len(variable.type.shape) != 2 for variable in operator.inputs): - raise RuntimeError('Input(s) must be [N, C]-tensor(s)') + raise RuntimeError("Input(s) must be [N, C]-tensor(s)") - model_type = operator.raw_operator.WhichOneof('Type') - if model_type == 'treeEnsembleClassifier': - class_label_type = operator.raw_operator.treeEnsembleClassifier.WhichOneof('ClassLabels') - elif model_type == 'glmClassifier': - class_label_type = operator.raw_operator.glmClassifier.WhichOneof('ClassLabels') - elif model_type == 'supportVectorClassifier': - class_label_type = operator.raw_operator.supportVectorClassifier.WhichOneof('ClassLabels') + model_type = operator.raw_operator.WhichOneof("Type") + if model_type == "treeEnsembleClassifier": + class_label_type = operator.raw_operator.treeEnsembleClassifier.WhichOneof( + "ClassLabels" + ) + elif model_type == "glmClassifier": + class_label_type = operator.raw_operator.glmClassifier.WhichOneof("ClassLabels") + elif model_type == "supportVectorClassifier": + class_label_type = operator.raw_operator.supportVectorClassifier.WhichOneof( + "ClassLabels" + ) else: - raise ValueError('%s has no class label' % model_type) + raise ValueError("%s has no class label" % model_type) N = operator.inputs[0].type.shape[0] if operator.target_opset < 7: @@ -40,28 +58,48 @@ def calculate_traditional_classifier_output_shapes(operator): else: output_shape = [N] - if class_label_type == 'stringClassLabels': - operator.outputs[0].type = StringTensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + if class_label_type == "stringClassLabels": + operator.outputs[0].type = StringTensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) if len(operator.outputs) == 2: if operator.target_opset < 7: - operator.outputs[1].type = DictionaryType(StringTensorType([1]), FloatTensorType([1]), - doc_string=operator.outputs[1].type.doc_string) + operator.outputs[1].type = DictionaryType( + StringTensorType([1]), + FloatTensorType([1]), + doc_string=operator.outputs[1].type.doc_string, + ) else: - operator.outputs[1].type = SequenceType(DictionaryType(StringTensorType([]), FloatTensorType([])), - doc_string=operator.outputs[1].type.doc_string) - elif class_label_type == 'int64ClassLabels': - operator.outputs[0].type = Int64TensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + operator.outputs[1].type = SequenceType( + DictionaryType(StringTensorType([]), FloatTensorType([])), + doc_string=operator.outputs[1].type.doc_string, + ) + elif class_label_type == "int64ClassLabels": + operator.outputs[0].type = Int64TensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) if len(operator.outputs) == 2: if operator.target_opset < 7: - operator.outputs[1].type = DictionaryType(Int64TensorType([1]), FloatTensorType([1]), - doc_string=operator.outputs[1].type.doc_string) + operator.outputs[1].type = DictionaryType( + Int64TensorType([1]), + FloatTensorType([1]), + doc_string=operator.outputs[1].type.doc_string, + ) else: - operator.outputs[1].type = SequenceType(DictionaryType(Int64TensorType([]), FloatTensorType([])), - doc_string=operator.outputs[1].type.doc_string) + operator.outputs[1].type = SequenceType( + DictionaryType(Int64TensorType([]), FloatTensorType([])), + doc_string=operator.outputs[1].type.doc_string, + ) else: - raise ValueError('Traditional classifier must include label information') + raise ValueError("Traditional classifier must include label information") -register_shape_calculator('glmClassifier', calculate_traditional_classifier_output_shapes) -register_shape_calculator('supportVectorClassifier', calculate_traditional_classifier_output_shapes) -register_shape_calculator('treeEnsembleClassifier', calculate_traditional_classifier_output_shapes) +register_shape_calculator( + "glmClassifier", calculate_traditional_classifier_output_shapes +) +register_shape_calculator( + "supportVectorClassifier", calculate_traditional_classifier_output_shapes +) +register_shape_calculator( + "treeEnsembleClassifier", calculate_traditional_classifier_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/DictVectorizer.py b/onnxmltools/convert/coreml/shape_calculators/DictVectorizer.py index 68d744cd..2900c818 100644 --- a/onnxmltools/convert/coreml/shape_calculators/DictVectorizer.py +++ b/onnxmltools/convert/coreml/shape_calculators/DictVectorizer.py @@ -6,32 +6,41 @@ def calculate_dictionary_vectorizer_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. Map ---> [1, C] C is the number of all allowed keys in the input dictionary. - ''' + """ # We assume all dictionaries' value types are float. It seems be reasonable to CoreML's # model input, but the existence of other map types leads to some concerns. check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - # Two types are allowed. One is DictionaryType and the other one, SequenceType, means a sequence of dictionaries. - check_input_and_output_types(operator, good_input_types=[DictionaryType, SequenceType]) + # Two types are allowed. One is DictionaryType and + # the other one, SequenceType, means a sequence of dictionaries. + check_input_and_output_types( + operator, good_input_types=[DictionaryType, SequenceType] + ) params = operator.raw_operator.dictVectorizer string_key_vector = params.stringToIndex.vector int64_key_vector = params.int64ToIndex.vector if len(string_key_vector) > 0 and len(int64_key_vector) > 0: - raise RuntimeError('Only one key type can present at the same time') + raise RuntimeError("Only one key type can present at the same time") doc_string = operator.outputs[0].type.doc_string if len(string_key_vector) > 0: - operator.outputs[0].type = FloatTensorType([1, len(string_key_vector)], doc_string=doc_string) + operator.outputs[0].type = FloatTensorType( + [1, len(string_key_vector)], doc_string=doc_string + ) elif len(int64_key_vector) > 0: - operator.outputs[1].type.shape = FloatTensorType([1, len(int64_key_vector)], doc_string=doc_string) + operator.outputs[1].type.shape = FloatTensorType( + [1, len(int64_key_vector)], doc_string=doc_string + ) else: - raise ValueError('Key vector cannot be empty') + raise ValueError("Key vector cannot be empty") -register_shape_calculator('dictVectorizer', calculate_dictionary_vectorizer_output_shapes) +register_shape_calculator( + "dictVectorizer", calculate_dictionary_vectorizer_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/FeatureVectorizer.py b/onnxmltools/convert/coreml/shape_calculators/FeatureVectorizer.py index ae34e3c5..f5396625 100644 --- a/onnxmltools/convert/coreml/shape_calculators/FeatureVectorizer.py +++ b/onnxmltools/convert/coreml/shape_calculators/FeatureVectorizer.py @@ -6,40 +6,58 @@ def calculate_feature_vectorizer_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C_1], ..., [N, C_n] ---> [N, C_1 + ... + C_n] - Feature vectorizer concatenates all input tensors along the C-axis, so the output dimension along C-axis is simply + Feature vectorizer concatenates all input tensors + along the C-axis, so the output dimension along C-axis is simply a sum of all input features. - ''' - check_input_and_output_numbers(operator, input_count_range=[1, None], output_count_range=1) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType, FloatType, Int64Type]) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, None], output_count_range=1 + ) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType, Int64TensorType, FloatType, Int64Type], + ) if any(len(variable.type.shape) != 2 for variable in operator.inputs): - raise RuntimeError('Input(s) must be 2-D tensor(s)') + raise RuntimeError("Input(s) must be 2-D tensor(s)") # Find the first batch size which is not unknown - N = 'None' + N = "None" for variable in operator.inputs: - if variable.type.shape[0] != 'None': + if variable.type.shape[0] != "None": N = variable.type.shape[0] break for variable in operator.inputs: - if variable.type.shape[0] not in ['None', N]: - raise RuntimeError('The batch dimensions should be the same to all input tensors.') - - C = sum(info.inputDimensions for info in operator.raw_operator.featureVectorizer.inputList) - - # Currently, we only expect numerical inputs. If both of integers and floats exist, we may convert integers into - # floats before concatenating them. Thus, the output type is integer-like only if all inputs are integer-like. + if variable.type.shape[0] not in ["None", N]: + raise RuntimeError( + "The batch dimensions should be the same to all input tensors." + ) + + C = sum( + info.inputDimensions + for info in operator.raw_operator.featureVectorizer.inputList + ) + + # Currently, we only expect numerical inputs. If both of + # integers and floats exist, we may convert integers into + # floats before concatenating them. Thus, the output type + # is integer-like only if all inputs are integer-like. doc_string = operator.outputs[0].type.doc_string - if all(isinstance(variable.type, (Int64TensorType, Int64Type)) for variable in operator.inputs): + if all( + isinstance(variable.type, (Int64TensorType, Int64Type)) + for variable in operator.inputs + ): operator.outputs[0].type = Int64TensorType([N, C], doc_string=doc_string) elif isinstance(operator.inputs[0].type, (FloatTensorType, FloatType)): operator.outputs[0].type = FloatTensorType([N, C], doc_string=doc_string) else: - raise ValueError('Unsupported input type: %s' % type(operator.inputs[0].type)) + raise ValueError("Unsupported input type: %s" % type(operator.inputs[0].type)) -register_shape_calculator('featureVectorizer', calculate_feature_vectorizer_output_shapes) +register_shape_calculator( + "featureVectorizer", calculate_feature_vectorizer_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/Identity.py b/onnxmltools/convert/coreml/shape_calculators/Identity.py index c6fbead4..03ee1659 100644 --- a/onnxmltools/convert/coreml/shape_calculators/Identity.py +++ b/onnxmltools/convert/coreml/shape_calculators/Identity.py @@ -16,8 +16,7 @@ def calculate_identity_output_shapes(operator): output.type.doc_string = doc_string -register_shape_calculator('identity', calculate_identity_output_shapes) -register_shape_calculator('imputer', calculate_identity_output_shapes) -register_shape_calculator('scaler', calculate_identity_output_shapes) -register_shape_calculator('normalizer', calculate_identity_output_shapes) - +register_shape_calculator("identity", calculate_identity_output_shapes) +register_shape_calculator("imputer", calculate_identity_output_shapes) +register_shape_calculator("scaler", calculate_identity_output_shapes) +register_shape_calculator("normalizer", calculate_identity_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/OneHotEncoder.py b/onnxmltools/convert/coreml/shape_calculators/OneHotEncoder.py index 7e10746a..bd3535e5 100644 --- a/onnxmltools/convert/coreml/shape_calculators/OneHotEncoder.py +++ b/onnxmltools/convert/coreml/shape_calculators/OneHotEncoder.py @@ -6,16 +6,16 @@ def calculate_one_hot_encoder_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, 1] ---> [N, C'] C' is the total number of categorical values. - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) if operator.inputs[0].type.shape[1] != 1 or len(operator.inputs[0].type.shape) > 2: - raise RuntimeError('Input must be [N, 1]-tensor') + raise RuntimeError("Input must be [N, 1]-tensor") int_categories = operator.raw_operator.oneHotEncoder.int64Categories.vector str_categories = operator.raw_operator.oneHotEncoder.stringCategories.vector @@ -23,13 +23,15 @@ def calculate_one_hot_encoder_output_shapes(operator): N = operator.inputs[0].type.shape[0] if len(int_categories) > 0: - operator.outputs[0].type = FloatTensorType([N, len(int_categories)], - doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = FloatTensorType( + [N, len(int_categories)], doc_string=operator.outputs[0].type.doc_string + ) elif len(str_categories) > 0 and type(operator.inputs[0].type) == StringTensorType: - operator.outputs[0].type = FloatTensorType([N, len(str_categories)], - doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = FloatTensorType( + [N, len(str_categories)], doc_string=operator.outputs[0].type.doc_string + ) else: - raise ValueError('Categorical indexes are missing') + raise ValueError("Categorical indexes are missing") -register_shape_calculator('oneHotEncoder', calculate_one_hot_encoder_output_shapes) +register_shape_calculator("oneHotEncoder", calculate_one_hot_encoder_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/Regressor.py b/onnxmltools/convert/coreml/shape_calculators/Regressor.py index 2a2bdf36..4ecd0400 100644 --- a/onnxmltools/convert/coreml/shape_calculators/Regressor.py +++ b/onnxmltools/convert/coreml/shape_calculators/Regressor.py @@ -6,32 +6,45 @@ def calculate_traditional_regressor_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N, C'] - The number C' is the length of prediction vector. It can be a scalar (C'=1) or a vector (C'>1) - ''' - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType, FloatType, Int64Type]) + The number C' is the length of prediction vector. + It can be a scalar (C'=1) or a vector (C'>1) + """ + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType, Int64TensorType, FloatType, Int64Type], + ) if any(len(variable.type.shape) != 2 for variable in operator.inputs): - raise RuntimeError('Input(s) must be 2-D tensor(s)') + raise RuntimeError("Input(s) must be 2-D tensor(s)") - model_type = operator.raw_operator.WhichOneof('Type') - if model_type == 'glmRegressor': + model_type = operator.raw_operator.WhichOneof("Type") + if model_type == "glmRegressor": glm = operator.raw_operator.glmRegressor C = len(glm.weights) - elif model_type == 'treeEnsembleRegressor': + elif model_type == "treeEnsembleRegressor": tree = operator.raw_operator.treeEnsembleRegressor.treeEnsemble C = len(tree.basePredictionValue) - elif model_type == 'supportVectorRegressor': + elif model_type == "supportVectorRegressor": C = 1 else: - raise ValueError('Model should be one of linear model, tree-based model, and support vector machine') + raise ValueError( + "Model should be one of linear model, tree-based model, and support vector machine" + ) N = operator.inputs[0].type.shape[0] - operator.outputs[0].type = FloatTensorType([N, C], doc_string=operator.outputs[0].type.doc_string) - -register_shape_calculator('glmRegressor', calculate_traditional_regressor_output_shapes) -register_shape_calculator('supportVectorRegressor', calculate_traditional_regressor_output_shapes) -register_shape_calculator('treeEnsembleRegressor', calculate_traditional_regressor_output_shapes) + operator.outputs[0].type = FloatTensorType( + [N, C], doc_string=operator.outputs[0].type.doc_string + ) + + +register_shape_calculator("glmRegressor", calculate_traditional_regressor_output_shapes) +register_shape_calculator( + "supportVectorRegressor", calculate_traditional_regressor_output_shapes +) +register_shape_calculator( + "treeEnsembleRegressor", calculate_traditional_regressor_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/TensorToLabel.py b/onnxmltools/convert/coreml/shape_calculators/TensorToLabel.py index d563e1db..ef3c3efc 100644 --- a/onnxmltools/convert/coreml/shape_calculators/TensorToLabel.py +++ b/onnxmltools/convert/coreml/shape_calculators/TensorToLabel.py @@ -1,17 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 from ...common._registration import register_shape_calculator -from ...common.data_types import FloatTensorType, Int64TensorType, Int64Type, StringTensorType, StringType +from ...common.data_types import ( + FloatTensorType, + Int64TensorType, + Int64Type, + StringTensorType, + StringType, +) from ...common.utils import check_input_and_output_numbers, check_input_and_output_types def calculte_tensor_to_label_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N, 1] - Note that N must be 1 currently because TensorToProbability doesn't support batch size larger than 1. - ''' + Note that N must be 1 currently because TensorToProbability + doesn't support batch size larger than 1. + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -22,11 +29,15 @@ def calculte_tensor_to_label_output_shapes(operator): output_shape = [N, 1] if type(operator.outputs[0].type) in [Int64Type, Int64TensorType]: - operator.outputs[0].type = Int64TensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = Int64TensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) elif type(operator.outputs[0].type) in [StringType, StringTensorType]: - operator.outputs[0].type = StringTensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = StringTensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) else: - raise ValueError('Unsupported label type') + raise ValueError("Unsupported label type") -register_shape_calculator('tensorToLabel', calculte_tensor_to_label_output_shapes) +register_shape_calculator("tensorToLabel", calculte_tensor_to_label_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/TensorToProbabilityMap.py b/onnxmltools/convert/coreml/shape_calculators/TensorToProbabilityMap.py index 2b0f08df..69f457d1 100644 --- a/onnxmltools/convert/coreml/shape_calculators/TensorToProbabilityMap.py +++ b/onnxmltools/convert/coreml/shape_calculators/TensorToProbabilityMap.py @@ -1,12 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 from ...common._registration import register_shape_calculator -from ...common.data_types import DictionaryType, FloatTensorType, SequenceType, StringTensorType, Int64TensorType +from ...common.data_types import ( + DictionaryType, + FloatTensorType, + SequenceType, + StringTensorType, + Int64TensorType, +) from ...common.utils import check_input_and_output_numbers, check_input_and_output_types def calculate_tensor_to_probability_map_output_shapes(operator): - ''' + """ Allowed input/output patterns are ONNX < 1.2 1. [1, C] ---> ---> A map @@ -15,34 +21,46 @@ def calculate_tensor_to_probability_map_output_shapes(operator): 1. [N, C] ---> ---> A sequence of maps 2. [N, C_1, ..., C_n] ---> A sequence of maps - Note that N must be 1 currently if you're using ONNX<1.2 because old ZipMap doesn't produce a seqneuce of map If the - input is not [N, C], it will be reshaped into [N, C_1 x C_2, x ... x C_n] before being fed into ONNX ZipMap. - ''' + Note that N must be 1 currently if you're using + ONNX<1.2 because old ZipMap doesn't produce a seqneuce of map If the + input is not [N, C], it will be reshaped into + [N, C_1 x C_2, x ... x C_n] before being fed into ONNX ZipMap. + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) - model_type = operator.raw_operator.WhichOneof('Type') - if model_type == 'neuralNetworkClassifier': - class_label_type = operator.raw_operator.neuralNetworkClassifier.WhichOneof('ClassLabels') + model_type = operator.raw_operator.WhichOneof("Type") + if model_type == "neuralNetworkClassifier": + class_label_type = operator.raw_operator.neuralNetworkClassifier.WhichOneof( + "ClassLabels" + ) else: - raise TypeError('%s has no class label' % model_type) + raise TypeError("%s has no class label" % model_type) N = operator.inputs[0].type.shape[0] doc_string = operator.outputs[0].type.doc_string - if class_label_type == 'stringClassLabels': + if class_label_type == "stringClassLabels": if operator.target_opset < 7: - operator.outputs[0].type = DictionaryType(StringTensorType([1]), FloatTensorType([1]), doc_string) + operator.outputs[0].type = DictionaryType( + StringTensorType([1]), FloatTensorType([1]), doc_string + ) else: - operator.outputs[0].type = \ - SequenceType(DictionaryType(StringTensorType([]), FloatTensorType([])), N, doc_string) - elif class_label_type == 'int64ClassLabels': + operator.outputs[0].type = SequenceType( + DictionaryType(StringTensorType([]), FloatTensorType([])), N, doc_string + ) + elif class_label_type == "int64ClassLabels": if operator.target_opset < 7: - operator.outputs[0].type = DictionaryType(Int64TensorType([1]), FloatTensorType([1]), doc_string) + operator.outputs[0].type = DictionaryType( + Int64TensorType([1]), FloatTensorType([1]), doc_string + ) else: - operator.outputs[0].type = \ - SequenceType(DictionaryType(Int64TensorType([]), FloatTensorType([])), N, doc_string) + operator.outputs[0].type = SequenceType( + DictionaryType(Int64TensorType([]), FloatTensorType([])), N, doc_string + ) else: - raise ValueError('Unsupported label type') + raise ValueError("Unsupported label type") -register_shape_calculator('tensorToProbabilityMap', calculate_tensor_to_probability_map_output_shapes) +register_shape_calculator( + "tensorToProbabilityMap", calculate_tensor_to_probability_map_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/BatchNorm.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/BatchNorm.py index 844ce3aa..54e972c4 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/BatchNorm.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/BatchNorm.py @@ -3,25 +3,28 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_batch_normalization_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N, C] 2. [N, C, H, W] ---> [N, C, H, W] This operator just uses the operator input shape as its output shape. - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) input_shape = operator.inputs[0].type.shape if len(input_shape) not in [2, 4]: - raise RuntimeError('Input must be a 2-D or a 4-D tensor') + raise RuntimeError("Input must be a 2-D or a 4-D tensor") operator.outputs[0].type.shape = copy.deepcopy(operator.inputs[0].type.shape) -register_shape_calculator('batchnorm', calculate_batch_normalization_output_shapes) +register_shape_calculator("batchnorm", calculate_batch_normalization_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/BidirectionalLSTM.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/BidirectionalLSTM.py index 55718197..5bbdca36 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/BidirectionalLSTM.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/BidirectionalLSTM.py @@ -2,38 +2,53 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_bidirectional_lstm_output_shapes(operator): - ''' + """ See bidirectional LSTM's conversion function for its output shapes. - ''' - check_input_and_output_numbers(operator, input_count_range=[1, 5], output_count_range=[1, 5]) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, 5], output_count_range=[1, 5] + ) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) input_shape = operator.inputs[0].type.shape # LSTM accepts [N, C] and [N, C, 1, 1] inputs if len(input_shape) not in [2, 4]: - raise RuntimeError('Input must be a 2-D or 4-D tensor') + raise RuntimeError("Input must be a 2-D or 4-D tensor") params = operator.raw_operator.biDirectionalLSTM # The following line is more accurate but it may break some tests - # output_shape = ['None', params.outputVectorSize] if params.params.sequenceOutput else [1, 2 *params.outputVectorSize] - output_shape = ['None', 2 * params.outputVectorSize] + # output_shape = ['None', params.outputVectorSize] + # if params.params.sequenceOutput else [1, 2 *params.outputVectorSize] + output_shape = ["None", 2 * params.outputVectorSize] state_shape = [1, params.outputVectorSize] - # TODO: Changing input shapes of an operator is dangerous, this should be move to Topology's _fix_shapes function + # TODO: Changing input shapes of an operator is dangerous, + # this should be move to Topology's _fix_shapes function if len(operator.inputs) > 1: - Y_h_in = operator.inputs[1] # The forward initial hidden state of a single sequence + Y_h_in = operator.inputs[ + 1 + ] # The forward initial hidden state of a single sequence Y_h_in.type.shape = state_shape - Y_h_rev_in = operator.inputs[3] # The backward initial hidden state of a single sequence + Y_h_rev_in = operator.inputs[ + 3 + ] # The backward initial hidden state of a single sequence Y_h_rev_in.type.shape = state_shape if len(operator.inputs) > 2: - Y_c_in = operator.inputs[2] # The forward initial cell state of a single sequence + Y_c_in = operator.inputs[ + 2 + ] # The forward initial cell state of a single sequence Y_c_in.type.shape = state_shape - Y_c_rev_in = operator.inputs[4] # The backward initial cell state of a single sequence + Y_c_rev_in = operator.inputs[ + 4 + ] # The backward initial cell state of a single sequence Y_c_rev_in.type.shape = state_shape operator.outputs[0].type.shape = output_shape @@ -45,4 +60,6 @@ def calculate_bidirectional_lstm_output_shapes(operator): operator.outputs[4].type.shape = state_shape -register_shape_calculator('biDirectionalLSTM', calculate_bidirectional_lstm_output_shapes) +register_shape_calculator( + "biDirectionalLSTM", calculate_bidirectional_lstm_output_shapes +) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Concat.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Concat.py index a4ebcae8..377c07d4 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Concat.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Concat.py @@ -4,27 +4,39 @@ from ....common.utils import check_input_and_output_numbers from ....common._registration import register_shape_calculator + def calculate_concat_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N_1, C, H, W], ..., [N_n, C, H, W] ---> [N_1 + ... + N_n, C, H, W] 2. [N, C_1, H, W], ..., [N, C_n, H, W] ---> [N, C_1 + ... + C_n, H, W] - ''' - check_input_and_output_numbers(operator, input_count_range=[1, None], output_count_range=[1, 1]) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, None], output_count_range=[1, 1] + ) output_shape = copy.deepcopy(operator.inputs[0].type.shape) dims = [] for variable in operator.inputs: - if variable.type.shape[0] != 'None' and variable.type.shape[0] != output_shape[0]: - raise RuntimeError('Only dimensions along C-axis can be different') - if variable.type.shape[2] != 'None' and variable.type.shape[2] != output_shape[2]: - raise RuntimeError('Only dimensions along C-axis can be different') - if variable.type.shape[3] != 'None' and variable.type.shape[3] != output_shape[3]: - raise RuntimeError('Only dimensions along C-axis can be different') + if ( + variable.type.shape[0] != "None" + and variable.type.shape[0] != output_shape[0] + ): + raise RuntimeError("Only dimensions along C-axis can be different") + if ( + variable.type.shape[2] != "None" + and variable.type.shape[2] != output_shape[2] + ): + raise RuntimeError("Only dimensions along C-axis can be different") + if ( + variable.type.shape[3] != "None" + and variable.type.shape[3] != output_shape[3] + ): + raise RuntimeError("Only dimensions along C-axis can be different") dims.append(variable.type.shape[1]) - output_shape[1] = 'None' if 'None' in dims else sum(dims) + output_shape[1] = "None" if "None" in dims else sum(dims) operator.outputs[0].type.shape = output_shape -register_shape_calculator('concat', calculate_concat_output_shapes) +register_shape_calculator("concat", calculate_concat_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Convolution.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Convolution.py index 749e10da..76e1888f 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Convolution.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Convolution.py @@ -7,56 +7,88 @@ def calculate_convolution_and_pooling_1D_output_shape( - input_size, kernel_size, kernel_dilation, stride, pad_mode, pad_head, pad_tail, output_size=0): + input_size, + kernel_size, + kernel_dilation, + stride, + pad_mode, + pad_head, + pad_tail, + output_size=0, +): if not isinstance(input_size, numbers.Integral): - return 'None' + return "None" if output_size > 0: return int(output_size) # Must use output_size = 1 for global pooling - effective_kernel_size = 1 + kernel_dilation * (kernel_size - 1) # For pooling, we always have dilation = 1. - if pad_mode == 'valid': - return int(math.floor((input_size + pad_head + pad_tail - effective_kernel_size) / stride) + 1) - elif pad_mode == 'same': + effective_kernel_size = 1 + kernel_dilation * ( + kernel_size - 1 + ) # For pooling, we always have dilation = 1. + if pad_mode == "valid": + return int( + math.floor( + (input_size + pad_head + pad_tail - effective_kernel_size) / stride + ) + + 1 + ) + elif pad_mode == "same": return int(math.ceil(input_size / stride)) - elif pad_mode == 'includeLastPixel': + elif pad_mode == "includeLastPixel": if pad_head != pad_tail: - raise ValueError('Padding amounts at the beginning and the end of an axis must be the same') + raise ValueError( + "Padding amounts at the beginning and the end of an axis must be the same" + ) effective_input_size = input_size + pad_head + pad_tail - effective_kernel_size out_size = math.ceil(effective_input_size / stride) + 1 if (out_size - 1) * stride >= input_size + pad_head: out_size -= 1 return out_size else: - raise ValueError('Unknown padding mode: %s' % pad_mode) + raise ValueError("Unknown padding mode: %s" % pad_mode) def calculate_convolution_transpose_1D_output_shape( - input_size, kernel_size, kernel_dilation, stride, pad_mode, pad_head, pad_tail, output_size=0): + input_size, + kernel_size, + kernel_dilation, + stride, + pad_mode, + pad_head, + pad_tail, + output_size=0, +): if not isinstance(input_size, numbers.Integral): - return 'None' + return "None" if output_size > 0: return output_size effective_kernel_size = 1 + kernel_dilation * (kernel_size - 1) - if pad_mode == 'valid': - return int((input_size - 1) * stride - pad_head - pad_tail + effective_kernel_size) - elif pad_mode == 'same': + if pad_mode == "valid": + return int( + (input_size - 1) * stride - pad_head - pad_tail + effective_kernel_size + ) + elif pad_mode == "same": return int(input_size * stride) else: - raise ValueError('Unknown padding mode: %s' % pad_mode) + raise ValueError("Unknown padding mode: %s" % pad_mode) def calculate_convolution_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C, H', W'] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) params = operator.raw_operator.convolution input_shape = operator.inputs[0].type.shape - operator.outputs[0].type.shape = [0, 0, 0, 0] # Initialize output shape. It will be modified below. + operator.outputs[0].type.shape = [ + 0, + 0, + 0, + 0, + ] # Initialize output shape. It will be modified below. output_shape = operator.outputs[0].type.shape # Adjust N-axis @@ -78,13 +110,14 @@ def calculate_convolution_output_shapes(operator): specified_output_shape = [0, 0] # Only used with convolution transpose if params.isDeconvolution and len(params.outputShape) > 0: specified_output_shape = list(int(i) for i in params.outputShape) - pad_mode = params.WhichOneof('ConvolutionPaddingType') - if pad_mode == 'valid' and len(params.valid.paddingAmounts.borderAmounts) > 0: + pad_mode = params.WhichOneof("ConvolutionPaddingType") + if pad_mode == "valid" and len(params.valid.paddingAmounts.borderAmounts) > 0: pad_amounts = params.valid.paddingAmounts.borderAmounts pad_heads = [pad_amounts[0].startEdgeSize, pad_amounts[1].startEdgeSize] pad_tails = [pad_amounts[0].endEdgeSize, pad_amounts[1].endEdgeSize] else: - # Padding amounts are useless for same padding and valid padding uses [0, 0] by default. + # Padding amounts are useless for same + # padding and valid padding uses [0, 0] by default. pad_heads = [0, 0] pad_tails = [0, 0] @@ -92,12 +125,25 @@ def calculate_convolution_output_shapes(operator): for i in range(2): if params.isDeconvolution: output_shape[i + 2] = calculate_convolution_transpose_1D_output_shape( - input_shape[i + 2], kernel_shape[i], dilations[i], strides[i], - pad_mode, pad_heads[i], pad_tails[i], specified_output_shape[i]) + input_shape[i + 2], + kernel_shape[i], + dilations[i], + strides[i], + pad_mode, + pad_heads[i], + pad_tails[i], + specified_output_shape[i], + ) else: output_shape[i + 2] = calculate_convolution_and_pooling_1D_output_shape( - input_shape[i + 2], kernel_shape[i], dilations[i], strides[i], - pad_mode, pad_heads[i], pad_tails[i]) + input_shape[i + 2], + kernel_shape[i], + dilations[i], + strides[i], + pad_mode, + pad_heads[i], + pad_tails[i], + ) -register_shape_calculator('convolution', calculate_convolution_output_shapes) +register_shape_calculator("convolution", calculate_convolution_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Crop.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Crop.py index a23f23a9..d9664cf5 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Crop.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Crop.py @@ -3,16 +3,21 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_crop_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C, H', W'] 2. [N, C, H, W], shape-ref [N', C', H', W'] ---> [N, C, H', W'] - ''' - check_input_and_output_numbers(operator, input_count_range=[1, 2], output_count_range=1) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, 2], output_count_range=1 + ) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) output_shape = copy.deepcopy(operator.inputs[0].type.shape) @@ -28,9 +33,9 @@ def calculate_crop_output_shapes(operator): output_shape[2] = operator.inputs[1].type.shape[2] output_shape[3] = operator.inputs[1].type.shape[3] else: - raise RuntimeError('Too many inputs for Crop operator') + raise RuntimeError("Too many inputs for Crop operator") operator.outputs[0].type.shape = output_shape -register_shape_calculator('crop', calculate_crop_output_shapes) +register_shape_calculator("crop", calculate_crop_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Dot.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Dot.py index 4d264d00..81691852 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Dot.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Dot.py @@ -3,20 +3,23 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_dot_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C], [N, C] ---> [N, 1] 2. [N, C, 1, 1], [N, C, 1, 1] ---> [N, 1, 1, 1] - ''' + """ check_input_and_output_numbers(operator, input_count_range=2, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) if operator.inputs[0].type.shape != operator.inputs[1].type.shape: - raise RuntimeError('Input shapes must be identical') + raise RuntimeError("Input shapes must be identical") # Assume that inputs are [N, C]- or [N, C, 1, 1]-tensors output_shape = copy.deepcopy(operator.inputs[0].type.shape) @@ -24,4 +27,4 @@ def calculate_dot_output_shapes(operator): operator.outputs[0].type.shape = output_shape -register_shape_calculator('dot', calculate_dot_output_shapes) +register_shape_calculator("dot", calculate_dot_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Embed.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Embed.py index d7fd096b..459333e8 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Embed.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Embed.py @@ -2,24 +2,31 @@ from ....common._registration import register_shape_calculator from ....common.data_types import Int64Type, Int64TensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_embedding_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, 1] ---> [N, C] 2. [N, 1, 1, 1] ---> [N, C, 1, 1] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[Int64Type, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[Int64Type, Int64TensorType] + ) output = operator.outputs[0] input_shape = operator.inputs[0].type.shape - if input_shape[1] != 1 or (len(input_shape) > 2 and (input_shape[2] != 1 or input_shape[3] != 1)): - raise RuntimeError('If input is a 4-D tensor, its shape must be [N, 1, 1, 1]') + if input_shape[1] != 1 or ( + len(input_shape) > 2 and (input_shape[2] != 1 or input_shape[3] != 1) + ): + raise RuntimeError("If input is a 4-D tensor, its shape must be [N, 1, 1, 1]") params = operator.raw_operator.embedding if len(input_shape) == 4: @@ -27,9 +34,9 @@ def calculate_embedding_output_shapes(operator): elif len(input_shape) == 2: output_shape = [input_shape[0], params.outputChannels] else: - raise RuntimeError('Input must be a 2-D or a 4-D tensor') + raise RuntimeError("Input must be a 2-D or a 4-D tensor") output.type.shape = output_shape -register_shape_calculator('embedding', calculate_embedding_output_shapes) +register_shape_calculator("embedding", calculate_embedding_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Flatten.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Flatten.py index bc587d7f..c75ca865 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Flatten.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Flatten.py @@ -2,15 +2,18 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_flatten_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N, C] 2. [N, C, H, W] ---> [N, C * H * W] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -18,21 +21,21 @@ def calculate_flatten_output_shapes(operator): output = operator.outputs[0] if len(input.type.shape) not in [2, 4]: - raise RuntimeError('Input must be 2-D or 4-D float tensor') + raise RuntimeError("Input must be 2-D or 4-D float tensor") input_shape = input.type.shape output_shape = [input_shape[0], 1] # Calculate the multiplication of C, H, and W. for i in input_shape[1:]: - if i != 'None': + if i != "None": output_shape[1] *= i else: # If any of C, H, W-dimensions is unknown, the flatten C-dimension is unknown - output_shape[1] = 'None' + output_shape[1] = "None" break output.type.shape = output_shape -register_shape_calculator('flatten', calculate_flatten_output_shapes) +register_shape_calculator("flatten", calculate_flatten_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/GRU.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/GRU.py index 2d570798..55dc9951 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/GRU.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/GRU.py @@ -2,34 +2,44 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_gru_output_shapes(operator): - ''' + """ See GRU's conversion function for its output shapes. - ''' - check_input_and_output_numbers(operator, input_count_range=[1, 2], output_count_range=[1, 2]) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, 2], output_count_range=[1, 2] + ) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) input_shape = operator.inputs[0].type.shape if len(input_shape) not in [2, 4]: - raise RuntimeError('Input must be a [N, C]- or [N, C, 1, 1]-tensor') + raise RuntimeError("Input must be a [N, C]- or [N, C, 1, 1]-tensor") - if operator.type == 'gru': + if operator.type == "gru": params = operator.raw_operator.gru - elif operator.type == 'simpleRecurrent': + elif operator.type == "simpleRecurrent": params = operator.raw_operator.simpleRecurrent else: - raise RuntimeError('Only GRU and SimpleRNN are supported') + raise RuntimeError("Only GRU and SimpleRNN are supported") # The following line is more accurate but it may break some tests - # output_shape = ['None', params.outputVectorSize] if params.params.sequenceOutput else [2, params.outputVectorSize] - output_shape = [input_shape[0] if params.sequenceOutput else 'None', params.outputVectorSize] # 'None' should be 1 + # output_shape = ['None', params.outputVectorSize] + # if params.params.sequenceOutput else [2, params.outputVectorSize] + output_shape = [ + input_shape[0] if params.sequenceOutput else "None", + params.outputVectorSize, + ] # 'None' should be 1 state_shape = [1, params.outputVectorSize] - # TODO: Changing input shapes of an operator is dangerous, this should be move to Topology's _fix_shapes function + # TODO: Changing input shapes of an operator is dangerous, + # this should be move to Topology's _fix_shapes function if len(operator.inputs) > 1: Y_h_in = operator.inputs[1] # The initial hidden state of a single sequence Y_h_in.type.shape = state_shape @@ -39,5 +49,5 @@ def calculate_gru_output_shapes(operator): operator.outputs[1].type.shape = state_shape -register_shape_calculator('gru', calculate_gru_output_shapes) -register_shape_calculator('simpleRecurrent', calculate_gru_output_shapes) +register_shape_calculator("gru", calculate_gru_output_shapes) +register_shape_calculator("simpleRecurrent", calculate_gru_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/IdentityFloat.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/IdentityFloat.py index 69f9f4d3..d83da217 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/IdentityFloat.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/IdentityFloat.py @@ -3,7 +3,10 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_identical_float_tensor_shapes(operator): @@ -14,20 +17,24 @@ def calculate_identical_float_tensor_shapes(operator): output = operator.outputs[0] doc_string = output.type.doc_string - output.type.shape = copy.deepcopy(input.type.shape) # Similar to identity but only accept floats + output.type.shape = copy.deepcopy( + input.type.shape + ) # Similar to identity but only accept floats output.type.doc_string = doc_string # Preprocessing layers in CoreML -register_shape_calculator('scalerPreprocessor', calculate_identical_float_tensor_shapes) -register_shape_calculator('meanImagePreprocessor', calculate_identical_float_tensor_shapes) +register_shape_calculator("scalerPreprocessor", calculate_identical_float_tensor_shapes) +register_shape_calculator( + "meanImagePreprocessor", calculate_identical_float_tensor_shapes +) # Standard neural network layers -register_shape_calculator('activation', calculate_identical_float_tensor_shapes) -register_shape_calculator('bias', calculate_identical_float_tensor_shapes) -register_shape_calculator('l2normalize', calculate_identical_float_tensor_shapes) -register_shape_calculator('lrn', calculate_identical_float_tensor_shapes) -register_shape_calculator('mvn', calculate_identical_float_tensor_shapes) -register_shape_calculator('scale', calculate_identical_float_tensor_shapes) -register_shape_calculator('softmax', calculate_identical_float_tensor_shapes) -register_shape_calculator('unary', calculate_identical_float_tensor_shapes) +register_shape_calculator("activation", calculate_identical_float_tensor_shapes) +register_shape_calculator("bias", calculate_identical_float_tensor_shapes) +register_shape_calculator("l2normalize", calculate_identical_float_tensor_shapes) +register_shape_calculator("lrn", calculate_identical_float_tensor_shapes) +register_shape_calculator("mvn", calculate_identical_float_tensor_shapes) +register_shape_calculator("scale", calculate_identical_float_tensor_shapes) +register_shape_calculator("softmax", calculate_identical_float_tensor_shapes) +register_shape_calculator("unary", calculate_identical_float_tensor_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/InnerProduct.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/InnerProduct.py index 834dfcca..3487764e 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/InnerProduct.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/InnerProduct.py @@ -2,14 +2,18 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) + def calculate_inner_product_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N, C'] 2. [N, C, 1, 1] ---> [N, C', 1, 1] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -18,20 +22,22 @@ def calculate_inner_product_output_shapes(operator): input_shape = input.type.shape if len(input_shape) == 4 and (input_shape[2] != 1 or input_shape[3] != 1): - raise RuntimeError('If input is a 4-D tensor, its shape must be [N, C, 1, 1]') + raise RuntimeError("If input is a 4-D tensor, its shape must be [N, C, 1, 1]") params = operator.raw_operator.innerProduct if input_shape[1] != params.inputChannels: - raise RuntimeError('Dimension mismatch along C-axis. Expected %s but got %s' % - (params.inputChannels, input_shape[1])) + raise RuntimeError( + "Dimension mismatch along C-axis. Expected %s but got %s" + % (params.inputChannels, input_shape[1]) + ) if len(input_shape) == 4: output.type.shape = [input_shape[0], params.outputChannels, 1, 1] elif len(input_shape) == 2: output.type.shape = [input_shape[0], params.outputChannels] else: - raise RuntimeError('Input must be a 2-D or a 4-D tensor') + raise RuntimeError("Input must be a 2-D or a 4-D tensor") -register_shape_calculator('innerProduct', calculate_inner_product_output_shapes) +register_shape_calculator("innerProduct", calculate_inner_product_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/LSTM.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/LSTM.py index 69451a2f..da2bb1b9 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/LSTM.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/LSTM.py @@ -2,29 +2,36 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_lstm_output_shapes(operator): - ''' + """ See LSTM's conversion function for its output shapes. - ''' - check_input_and_output_numbers(operator, input_count_range=[1, 3], output_count_range=[1, 3]) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, 3], output_count_range=[1, 3] + ) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) input_shape = operator.inputs[0].type.shape if len(input_shape) not in [2, 4]: - raise RuntimeError('Input must be a 2-D tensor') + raise RuntimeError("Input must be a 2-D tensor") params = operator.raw_operator.uniDirectionalLSTM # The following line is more accurate but it may break some tests - # output_shape = ['None', params.outputVectorSize] if params.params.sequenceOutput else [1, params.outputVectorSize] - output_shape = ['None', params.outputVectorSize] + # output_shape = ['None', params.outputVectorSize] + # if params.params.sequenceOutput else [1, params.outputVectorSize] + output_shape = ["None", params.outputVectorSize] state_shape = [1, params.outputVectorSize] - # TODO: Changing input shapes of an operator is dangerous, this should be move to Topology's _fix_shapes function + # TODO: Changing input shapes of an operator is dangerous, + # this should be move to Topology's _fix_shapes function if len(operator.inputs) > 1: Y_h_in = operator.inputs[1] # The initial hidden state of a single sequence Y_h_in.type.shape = state_shape @@ -39,4 +46,4 @@ def calculate_lstm_output_shapes(operator): operator.outputs[2].type.shape = state_shape -register_shape_calculator('uniDirectionalLSTM', calculate_lstm_output_shapes) +register_shape_calculator("uniDirectionalLSTM", calculate_lstm_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstant.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstant.py index 73177f6a..5fb6d103 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstant.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstant.py @@ -4,8 +4,11 @@ from ....common.data_types import TensorType, FloatTensorType from ....common.utils import check_input_and_output_numbers + def calculate_load_constant_output_shapes(operator): - check_input_and_output_numbers(operator, input_count_range=None, output_count_range=1) + check_input_and_output_numbers( + operator, input_count_range=None, output_count_range=1 + ) output = operator.outputs[0] @@ -18,9 +21,9 @@ def calculate_load_constant_output_shapes(operator): output.type = FloatTensorType(const_shape, doc_string=output.type.doc_string) else: if not isinstance(output.type, TensorType): - raise RuntimeError('Type conflict detected. Output must be a tensor.') + raise RuntimeError("Type conflict detected. Output must be a tensor.") # If output type exists, we just modify its shape. output.type.shape = const_shape -register_shape_calculator('loadConstant', calculate_load_constant_output_shapes) +register_shape_calculator("loadConstant", calculate_load_constant_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstantND.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstantND.py index 436d9dbc..cd6ddec4 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstantND.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/LoadConstantND.py @@ -4,8 +4,11 @@ from ....common.data_types import TensorType, FloatTensorType from ....common.utils import check_input_and_output_numbers + def calculate_load_constant_nd_output_shapes(operator): - check_input_and_output_numbers(operator, input_count_range=None, output_count_range=1) + check_input_and_output_numbers( + operator, input_count_range=None, output_count_range=1 + ) output = operator.outputs[0] @@ -18,9 +21,9 @@ def calculate_load_constant_nd_output_shapes(operator): output.type = FloatTensorType(const_shape, doc_string=output.type.doc_string) else: if not isinstance(output.type, TensorType): - raise RuntimeError('Type conflict detected. Output must be a tensor.') + raise RuntimeError("Type conflict detected. Output must be a tensor.") # If output type exists, we just modify its shape. output.type.shape = const_shape -register_shape_calculator('loadConstantND', calculate_load_constant_nd_output_shapes) +register_shape_calculator("loadConstantND", calculate_load_constant_nd_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Merge.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Merge.py index 43c0e5cd..0b2fbb18 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Merge.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Merge.py @@ -2,35 +2,44 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_merge_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C_1, H_1, W_1], ..., [N, C_n, H_n, W_n] ---> [N, max(C_1, ..., C_n), max(H_1, ..., H_n), max(W_1, ..., W_n)] If 'None' happens at any coordinate, that coordinate's final dimension would be 'None'. - ''' - check_input_and_output_numbers(operator, input_count_range=[1, None], output_count_range=1) + """ + check_input_and_output_numbers( + operator, input_count_range=[1, None], output_count_range=1 + ) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) # [TODO] Fix reduce-like shape inference. We now assume all inputs are 4-D. n_dims = max(len(variable.type.shape) for variable in operator.inputs) output_shape = [0] * n_dims for i in range(n_dims): - input_dims = [variable.type.shape[i] for variable in operator.inputs if len(variable.type.shape) > i] - if 'None' in input_dims: - output_shape[i] = 'None' + input_dims = [ + variable.type.shape[i] + for variable in operator.inputs + if len(variable.type.shape) > i + ] + if "None" in input_dims: + output_shape[i] = "None" else: output_shape[i] = max(input_dims) operator.outputs[0].type.shape = output_shape -register_shape_calculator('add', calculate_merge_output_shapes) -register_shape_calculator('average', calculate_merge_output_shapes) -register_shape_calculator('max', calculate_merge_output_shapes) -register_shape_calculator('min', calculate_merge_output_shapes) -register_shape_calculator('multiply', calculate_merge_output_shapes) +register_shape_calculator("add", calculate_merge_output_shapes) +register_shape_calculator("average", calculate_merge_output_shapes) +register_shape_calculator("max", calculate_merge_output_shapes) +register_shape_calculator("min", calculate_merge_output_shapes) +register_shape_calculator("multiply", calculate_merge_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Pad.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Pad.py index 92990dab..97cf0869 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Pad.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Pad.py @@ -3,14 +3,17 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_padding_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C, H', W'] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -26,4 +29,4 @@ def calculate_padding_output_shapes(operator): operator.outputs[0].type.shape = output_shape -register_shape_calculator('padding', calculate_padding_output_shapes) +register_shape_calculator("padding", calculate_padding_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Permute.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Permute.py index f3169228..6dc1dae7 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Permute.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Permute.py @@ -3,19 +3,25 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType, Int64TensorType, StringTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_permute_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N', C', H', W'] Note that here [N', C', H', W'] means all possible permutations of [N, C, H, W] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType, StringTensorType], - good_output_types=[FloatTensorType, Int64TensorType, StringTensorType]) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType, Int64TensorType, StringTensorType], + good_output_types=[FloatTensorType, Int64TensorType, StringTensorType], + ) input = operator.inputs[0] output = operator.outputs[0] @@ -25,4 +31,4 @@ def calculate_permute_output_shapes(operator): output.type.shape = [input_shape[a] for a in axes] -register_shape_calculator('permute', calculate_permute_output_shapes) +register_shape_calculator("permute", calculate_permute_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Pool.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Pool.py index 05e3ecc2..84a080ff 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Pool.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Pool.py @@ -2,14 +2,18 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) from .Convolution import calculate_convolution_and_pooling_1D_output_shape + def calculate_pooling_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C, H', W'] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -17,7 +21,7 @@ def calculate_pooling_output_shapes(operator): input_shape = operator.inputs[0].type.shape if len(input.type.shape) != 4: - raise RuntimeError('Input must be 4-D float tensor') + raise RuntimeError("Input must be 4-D float tensor") operator.outputs[0].type.shape = [0, 0, 0, 0] output_shape = operator.outputs[0].type.shape @@ -29,20 +33,29 @@ def calculate_pooling_output_shapes(operator): output_shape[1] = input_shape[1] params = operator.raw_operator.pooling - # Set up default and non-default parameters. Notice that they are only set for H- and W-axes. - dilations = [1, 1] # CoreML Pooling doesn't allow dilation, so we use [1, 1] which is equivalent to no dilation. + # Set up default and non-default parameters. Notice that + # they are only set for H- and W-axes. + # CoreML Pooling doesn't allow dilation, so we use [1, 1] + # which is equivalent to no dilation. + dilations = [ + 1, + 1, + ] kernel_shape = [3, 3] if len(params.kernelSize) > 0: kernel_shape = params.kernelSize strides = [1, 1] if len(params.stride) > 0: strides = params.stride - pad_mode = params.WhichOneof('PoolingPaddingType') - if pad_mode == 'valid' and len(params.valid.paddingAmounts.borderAmounts) > 0: + pad_mode = params.WhichOneof("PoolingPaddingType") + if pad_mode == "valid" and len(params.valid.paddingAmounts.borderAmounts) > 0: pad_amounts = params.valid.paddingAmounts.borderAmounts pad_heads = [pad_amounts[0].startEdgeSize, pad_amounts[1].startEdgeSize] pad_tails = [pad_amounts[0].endEdgeSize, pad_amounts[1].endEdgeSize] - elif pad_mode == 'includeLastPixel' and len(params.includeLastPixel.paddingAmounts) > 0: + elif ( + pad_mode == "includeLastPixel" + and len(params.includeLastPixel.paddingAmounts) > 0 + ): pad_amounts = params.includeLastPixel.paddingAmounts pad_heads = [pad_amounts[0], pad_amounts[1]] pad_tails = [pad_amounts[0], pad_amounts[1]] @@ -54,8 +67,15 @@ def calculate_pooling_output_shapes(operator): # Calculate output shape along H- and W-axes for i in range(2): output_shape[i + 2] = calculate_convolution_and_pooling_1D_output_shape( - input_shape[i + 2], kernel_shape[i], dilations[i], strides[i], - pad_mode, pad_heads[i], pad_tails[i], params.globalPooling) + input_shape[i + 2], + kernel_shape[i], + dilations[i], + strides[i], + pad_mode, + pad_heads[i], + pad_tails[i], + params.globalPooling, + ) -register_shape_calculator('pooling', calculate_pooling_output_shapes) +register_shape_calculator("pooling", calculate_pooling_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Reduce.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Reduce.py index 42fb76da..4df817bb 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Reduce.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Reduce.py @@ -3,18 +3,21 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_reduce_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, 1, H, W] 2. [N, C, H, W] ---> [N, C, 1, W] 3. [N, C, H, W] ---> [N, C, H, 1] 4. [N, C, H, W] ---> [N, C, 1, 1] 5. [N, C, H, W] ---> [N, 1, 1, 1] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -22,6 +25,7 @@ def calculate_reduce_output_shapes(operator): params = operator.raw_operator.reduce from coremltools.proto.NeuralNetwork_pb2 import ReduceLayerParams as Params + # Adjust C-axis if params.axis in [Params.CHW, Params.C]: output_shape[1] = 1 @@ -35,4 +39,4 @@ def calculate_reduce_output_shapes(operator): operator.outputs[0].type.shape = output_shape -register_shape_calculator('reduce', calculate_reduce_output_shapes) +register_shape_calculator("reduce", calculate_reduce_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/ReorganizeData.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/ReorganizeData.py index 513c1c62..61675eab 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/ReorganizeData.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/ReorganizeData.py @@ -3,17 +3,20 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_reorganize_data_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C * B * B , H / B, W / B] 2. [N, C, H, W] ---> [N, C / B / B , H * B, W * B] Note taht B is the block size specified in this operator. - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -25,21 +28,36 @@ def calculate_reorganize_data_output_shapes(operator): if params.mode == Params.DEPTH_TO_SPACE: if output_shape[1] % (params.blockSize * params.blockSize) != 0: - raise RuntimeError('Channel number must be divisible by the square of block size') - - output_shape = [output_shape[0], output_shape[1] / params.blockSize / params.blockSize, - output_shape[2] * params.blockSize, output_shape[3] * params.blockSize] + raise RuntimeError( + "Channel number must be divisible by the square of block size" + ) + + output_shape = [ + output_shape[0], + output_shape[1] / params.blockSize / params.blockSize, + output_shape[2] * params.blockSize, + output_shape[3] * params.blockSize, + ] elif params.mode == Params.SPACE_TO_DEPTH: - if output_shape[2] % params.blockSize != 0 or output_shape[3] % params.blockSize != 0: - raise RuntimeError('Height and weight must be divisible by block size') - - output_shape = [output_shape[0], output_shape[1] * params.blockSize * params.blockSize, - output_shape[2] / params.blockSize, output_shape[3] / params.blockSize] + if ( + output_shape[2] % params.blockSize != 0 + or output_shape[3] % params.blockSize != 0 + ): + raise RuntimeError("Height and weight must be divisible by block size") + + output_shape = [ + output_shape[0], + output_shape[1] * params.blockSize * params.blockSize, + output_shape[2] / params.blockSize, + output_shape[3] / params.blockSize, + ] else: - raise ValueError('Unsupport reorganization mode {0}'.format(params.mode)) + raise ValueError("Unsupport reorganization mode {0}".format(params.mode)) - operator.outputs[0].type = FloatTensorType([int(i) if i != 'None' else 'None' for i in output_shape], - doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = FloatTensorType( + [int(i) if i != "None" else "None" for i in output_shape], + doc_string=operator.outputs[0].type.doc_string, + ) -register_shape_calculator('reorganizeData', calculate_reorganize_data_output_shapes) +register_shape_calculator("reorganizeData", calculate_reorganize_data_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Reshape.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Reshape.py index fb2f4bbe..f4e46cde 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Reshape.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Reshape.py @@ -2,15 +2,19 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) + def calculate_reshape_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C', H', W'] Note that C*H*W should equal to C'*H'*W'. - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -24,4 +28,4 @@ def calculate_reshape_output_shapes(operator): operator.outputs[0].type.shape = output_shape -register_shape_calculator('reshape', calculate_reshape_output_shapes) +register_shape_calculator("reshape", calculate_reshape_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/ReshapeStatic.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/ReshapeStatic.py index 525e00c4..4160706f 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/ReshapeStatic.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/ReshapeStatic.py @@ -2,15 +2,19 @@ from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) + def calculate_reshape_static_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C', H', W'] Note that C*H*W should equal to C'*H'*W'. - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -24,4 +28,4 @@ def calculate_reshape_static_output_shapes(operator): operator.outputs[0].type.shape = output_shape -register_shape_calculator('reshapeStatic', calculate_reshape_static_output_shapes) +register_shape_calculator("reshapeStatic", calculate_reshape_static_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/SequenceRepeat.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/SequenceRepeat.py index bf44c568..ca2e6de5 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/SequenceRepeat.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/SequenceRepeat.py @@ -3,26 +3,34 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_sequence_repeat_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N', C, H, W] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) input_shape = operator.inputs[0].type.shape if len(input_shape) not in [2, 4]: - raise RuntimeError('Input shape of CoreML SequenceRepeat must be either 2-D or 4-D but got %s' % input_shape) + raise RuntimeError( + "Input shape of CoreML SequenceRepeat must be either 2-D or 4-D but got %s" + % input_shape + ) output_shape = copy.deepcopy(operator.inputs[0].type.shape) - if output_shape[0] != None: + if output_shape[0] is not None: output_shape[0] *= operator.raw_operator.sequenceRepeat.nRepetitions - operator.outputs[0].type = FloatTensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = FloatTensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) -register_shape_calculator('sequenceRepeat', calculate_sequence_repeat_output_shapes) +register_shape_calculator("sequenceRepeat", calculate_sequence_repeat_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Slice.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Slice.py index 84132e1d..e063102c 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Slice.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Slice.py @@ -3,16 +3,19 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_slice_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C', H, W] 2. [N, C, H, W] ---> [N, C, H', W] 3. [N, C, H, W] ---> [N, C, H, W'] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -21,14 +24,21 @@ def calculate_slice_output_shapes(operator): params = operator.raw_operator.slice from coremltools.proto.NeuralNetwork_pb2 import SliceLayerParams as Params + axis_map = {Params.CHANNEL_AXIS: 1, Params.HEIGHT_AXIS: 2, Params.WIDTH_AXIS: 3} if params.startIndex >= 0: - output_shape[axis_map[Params.CHANNEL_AXIS]] = params.endIndex - params.startIndex + output_shape[axis_map[Params.CHANNEL_AXIS]] = ( + params.endIndex - params.startIndex + ) else: - output_shape[axis_map[Params.CHANNEL_AXIS]] += 1 + params.endIndex - params.startIndex + output_shape[axis_map[Params.CHANNEL_AXIS]] += ( + 1 + params.endIndex - params.startIndex + ) - operator.outputs[0].type = FloatTensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = FloatTensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) -register_shape_calculator('slice', calculate_slice_output_shapes) +register_shape_calculator("slice", calculate_slice_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Split.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Split.py index 270667b4..c9a423aa 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Split.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Split.py @@ -3,23 +3,30 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_split_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C] ---> [N', C] 2. [N, C, H, W] ---> [N', C, H, W] - ''' - check_input_and_output_numbers(operator, input_count_range=1, output_count_range=[1, None]) + """ + check_input_and_output_numbers( + operator, input_count_range=1, output_count_range=[1, None] + ) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) output_shape = copy.deepcopy(operator.inputs[0].type.shape) divided = output_shape[1] / operator.raw_operator.split.nOutputs if divided != int(divided): - raise RuntimeError('Variable dimension along C-axis must be divisible by partition number') + raise RuntimeError( + "Variable dimension along C-axis must be divisible by partition number" + ) output_shape[1] = int(divided) @@ -27,4 +34,4 @@ def calculate_split_output_shapes(operator): operator.outputs[i].type.shape = copy.deepcopy(output_shape) -register_shape_calculator('split', calculate_split_output_shapes) +register_shape_calculator("split", calculate_split_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/Upsample.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/Upsample.py index 3527a746..b602d138 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/Upsample.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/Upsample.py @@ -3,14 +3,17 @@ import copy from ....common._registration import register_shape_calculator from ....common.data_types import FloatTensorType -from ....common.utils import check_input_and_output_numbers, check_input_and_output_types +from ....common.utils import ( + check_input_and_output_numbers, + check_input_and_output_types, +) def calculate_upsample_output_shapes(operator): - ''' + """ Allowed input/output patterns are 1. [N, C, H, W] ---> [N, C, H', W'] - ''' + """ check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) @@ -20,7 +23,9 @@ def calculate_upsample_output_shapes(operator): output_shape[2] *= scales[0] output_shape[3] *= scales[1] - operator.outputs[0].type = FloatTensorType(output_shape, doc_string=operator.outputs[0].type.doc_string) + operator.outputs[0].type = FloatTensorType( + output_shape, doc_string=operator.outputs[0].type.doc_string + ) -register_shape_calculator('upsample', calculate_upsample_output_shapes) +register_shape_calculator("upsample", calculate_upsample_output_shapes) diff --git a/onnxmltools/convert/coreml/shape_calculators/neural_network/__init__.py b/onnxmltools/convert/coreml/shape_calculators/neural_network/__init__.py index c5185d6d..246e14db 100644 --- a/onnxmltools/convert/coreml/shape_calculators/neural_network/__init__.py +++ b/onnxmltools/convert/coreml/shape_calculators/neural_network/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -# To register shape calculators for Core ML neural network operators, import associated modules here. +# To register shape calculators for Core ML neural +# network operators, import associated modules here. from . import BatchNorm from . import BidirectionalLSTM from . import Concat diff --git a/onnxmltools/convert/h2o/_parse.py b/onnxmltools/convert/h2o/_parse.py index dccfb720..d8224948 100644 --- a/onnxmltools/convert/h2o/_parse.py +++ b/onnxmltools/convert/h2o/_parse.py @@ -6,34 +6,45 @@ def _parse_h2o(scope, model, inputs): - ''' + """ :param scope: Scope object :param model: A h2o model data object :param inputs: A list of variables :return: A list of output variables which will be passed to next stage - ''' + """ this_operator = scope.declare_local_operator("H2OTreeMojo", model) this_operator.inputs = inputs if model["params"]["classifier"]: - label_variable = scope.declare_local_variable('label', FloatTensorType()) - probability_map_variable = scope.declare_local_variable('probabilities', FloatTensorType()) + label_variable = scope.declare_local_variable("label", FloatTensorType()) + probability_map_variable = scope.declare_local_variable( + "probabilities", FloatTensorType() + ) this_operator.outputs.append(label_variable) this_operator.outputs.append(probability_map_variable) else: - variable = scope.declare_local_variable('variable', FloatTensorType()) + variable = scope.declare_local_variable("variable", FloatTensorType()) this_operator.outputs.append(variable) return this_operator.outputs -def parse_h2o(model, initial_types=None, target_opset=None, - custom_conversion_functions=None, custom_shape_calculators=None): +def parse_h2o( + model, + initial_types=None, + target_opset=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): raw_model_container = H2OModelContainer(model) - topology = Topology(raw_model_container, default_batch_size='None', - initial_types=initial_types, target_opset=target_opset, - custom_conversion_functions=custom_conversion_functions, - custom_shape_calculators=custom_shape_calculators) - scope = topology.declare_scope('__root__') + topology = Topology( + raw_model_container, + default_batch_size="None", + initial_types=initial_types, + target_opset=target_opset, + custom_conversion_functions=custom_conversion_functions, + custom_shape_calculators=custom_shape_calculators, + ) + scope = topology.declare_scope("__root__") inputs = [] for var_name, initial_type in initial_types: diff --git a/onnxmltools/convert/h2o/convert.py b/onnxmltools/convert/h2o/convert.py index 3797f671..26d6292d 100644 --- a/onnxmltools/convert/h2o/convert.py +++ b/onnxmltools/convert/h2o/convert.py @@ -15,29 +15,45 @@ from . import operator_converters, shape_calculators # noqa -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): - ''' +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) + - 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 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 + :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: @@ -46,11 +62,11 @@ def convert(model, name=None, initial_types=None, doc_string='', target_opset=No >>> 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']))] + initial_types = [("input", FloatTensorType(shape=["None", "None"]))] if isinstance(model, str): model_path = model @@ -63,10 +79,20 @@ def convert(model, name=None, initial_types=None, doc_string='', target_opset=No 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"]) + "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 = 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) + onnx_model = convert_topology( + topology, name, doc_string, target_opset, targeted_onnx + ) return onnx_model diff --git a/onnxmltools/convert/h2o/operator_converters/h2o.py b/onnxmltools/convert/h2o/operator_converters/h2o.py index d9967730..0e3de060 100644 --- a/onnxmltools/convert/h2o/operator_converters/h2o.py +++ b/onnxmltools/convert/h2o/operator_converters/h2o.py @@ -3,9 +3,9 @@ from ...common._registration import register_converter _LINK_FUNCTION_TO_POST_TRANSFORM = { - 'identity': 'NONE', - 'logit': 'LOGISTIC', - 'ologit': 'LOGISTIC' + "identity": "NONE", + "logit": "LOGISTIC", + "ologit": "LOGISTIC", } @@ -13,7 +13,7 @@ def _get_post_transform(params): link_function = params["link_function"] family = params["family"] if family == "multinomial": - return 'SOFTMAX' + return "SOFTMAX" elif link_function not in _LINK_FUNCTION_TO_POST_TRANSFORM.keys(): raise ValueError("Link function %s not supported." % link_function) else: @@ -21,46 +21,61 @@ def _get_post_transform(params): def _get_default_tree_attribute_pairs(is_classifier, params): - attrs = { - 'post_transform': _get_post_transform(params) - } + attrs = {"post_transform": _get_post_transform(params)} nclasses = params["nclasses"] if is_classifier: predicted_classes = nclasses if nclasses > 2 else 1 - attrs['base_values'] = [params["base_score"] for _ in range(0, predicted_classes)] + attrs["base_values"] = [ + params["base_score"] for _ in range(0, predicted_classes) + ] else: - attrs['n_targets'] = 1 - attrs['base_values'] = [params["base_score"]] - for k in {'nodes_treeids', 'nodes_nodeids', - 'nodes_featureids', 'nodes_modes', 'nodes_values', - 'nodes_truenodeids', 'nodes_falsenodeids', 'nodes_missing_value_tracks_true'}: + attrs["n_targets"] = 1 + attrs["base_values"] = [params["base_score"]] + for k in { + "nodes_treeids", + "nodes_nodeids", + "nodes_featureids", + "nodes_modes", + "nodes_values", + "nodes_truenodeids", + "nodes_falsenodeids", + "nodes_missing_value_tracks_true", + }: attrs[k] = [] node_attr_prefix = _node_attr_prefix(is_classifier) - for k in {'_treeids', '_nodeids', '_ids', '_weights'}: + for k in {"_treeids", "_nodeids", "_ids", "_weights"}: attrs[node_attr_prefix + k] = [] return attrs def _add_node( - attr_pairs, is_classifier, tree_id, node_id, - feature_id, mode, value, true_child_id, false_child_id, weights, - missing + attr_pairs, + is_classifier, + tree_id, + node_id, + feature_id, + mode, + value, + true_child_id, + false_child_id, + weights, + missing, ): - attr_pairs['nodes_treeids'].append(tree_id) - attr_pairs['nodes_nodeids'].append(node_id) - attr_pairs['nodes_featureids'].append(feature_id) - attr_pairs['nodes_modes'].append(mode) - attr_pairs['nodes_values'].append(float(value)) - attr_pairs['nodes_truenodeids'].append(true_child_id) - attr_pairs['nodes_falsenodeids'].append(false_child_id) - attr_pairs['nodes_missing_value_tracks_true'].append(missing) - if mode == 'LEAF': + attr_pairs["nodes_treeids"].append(tree_id) + attr_pairs["nodes_nodeids"].append(node_id) + attr_pairs["nodes_featureids"].append(feature_id) + attr_pairs["nodes_modes"].append(mode) + attr_pairs["nodes_values"].append(float(value)) + attr_pairs["nodes_truenodeids"].append(true_child_id) + attr_pairs["nodes_falsenodeids"].append(false_child_id) + attr_pairs["nodes_missing_value_tracks_true"].append(missing) + if mode == "LEAF": node_attr_prefix = _node_attr_prefix(is_classifier) for i, w in enumerate(weights): - attr_pairs[node_attr_prefix + '_treeids'].append(tree_id) - attr_pairs[node_attr_prefix + '_nodeids'].append(node_id) - attr_pairs[node_attr_prefix + '_ids'].append(i) - attr_pairs[node_attr_prefix + '_weights'].append(float(w)) + attr_pairs[node_attr_prefix + "_treeids"].append(tree_id) + attr_pairs[node_attr_prefix + "_nodeids"].append(node_id) + attr_pairs[node_attr_prefix + "_ids"].append(i) + attr_pairs[node_attr_prefix + "_weights"].append(float(w)) def _node_attr_prefix(is_classifier): @@ -68,39 +83,41 @@ def _node_attr_prefix(is_classifier): def _fill_node_attributes(tree_id, node, attr_pairs, is_classifier): - if 'leftChild' in node: + if "leftChild" in node: if node["isCategorical"]: raise ValueError("categorical splits not supported, use one_hot_explicit") else: - operator = 'BRANCH_GTE' - value = node['splitValue'] + operator = "BRANCH_GTE" + value = node["splitValue"] _add_node( attr_pairs=attr_pairs, is_classifier=is_classifier, tree_id=tree_id, mode=operator, value=value, - node_id=node['id'], - feature_id=node['colId'], - true_child_id=node['rightChild']['id'], - false_child_id=node['leftChild']['id'], + node_id=node["id"], + feature_id=node["colId"], + true_child_id=node["rightChild"]["id"], + false_child_id=node["leftChild"]["id"], weights=None, missing=(0 if node["leftward"] else 1), ) _fill_node_attributes(tree_id, node["leftChild"], attr_pairs, is_classifier) _fill_node_attributes(tree_id, node["rightChild"], attr_pairs, is_classifier) else: # leaf - weights = [node['predValue']] + weights = [node["predValue"]] _add_node( attr_pairs=attr_pairs, is_classifier=is_classifier, tree_id=tree_id, - value=0., - node_id=node['id'], - feature_id=0, mode='LEAF', - true_child_id=0, false_child_id=0, + value=0.0, + node_id=node["id"], + feature_id=0, + mode="LEAF", + true_child_id=0, + false_child_id=0, weights=weights, - missing=False + missing=False, ) @@ -126,28 +143,41 @@ def convert_regression(scope, operator, container, params): fill_tree_attributes(model, attr_pairs, False) # add nodes - container.add_node('TreeEnsembleRegressor', operator.input_full_names, - operator.output_full_names, op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('TreeEnsembleRegressor'), **attr_pairs) + container.add_node( + "TreeEnsembleRegressor", + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("TreeEnsembleRegressor"), + **attr_pairs + ) def convert_classifier(scope, operator, container, params): if params["family"] == "multinomial" and params["nclasses"] == 2: - raise ValueError("Multinomial distribution with two classes not supported, use binomial distribution.") + raise ValueError( + "Multinomial distribution with two classes " + "not supported, use binomial distribution." + ) model = operator.raw_operator attr_pairs = _get_default_tree_attribute_pairs(True, params) fill_tree_attributes(model, attr_pairs, True) n_trees_in_group = params["n_trees_in_group"] - attr_pairs['class_ids'] = [v % n_trees_in_group for v in attr_pairs['class_treeids']] - attr_pairs['classlabels_strings'] = params["class_labels"] + attr_pairs["class_ids"] = [ + v % n_trees_in_group for v in attr_pairs["class_treeids"] + ] + attr_pairs["classlabels_strings"] = params["class_labels"] - container.add_node('TreeEnsembleClassifier', operator.input_full_names, - operator.output_full_names, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('TreeEnsembleClassifier'), - **attr_pairs) + container.add_node( + "TreeEnsembleClassifier", + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("TreeEnsembleClassifier"), + **attr_pairs + ) def convert_h2o(scope, operator, container): @@ -159,4 +189,4 @@ def convert_h2o(scope, operator, container): convert_regression(scope, operator, container, params) -register_converter('H2OTreeMojo', convert_h2o) +register_converter("H2OTreeMojo", convert_h2o) diff --git a/onnxmltools/convert/h2o/shape_calculators/h2otreemojo.py b/onnxmltools/convert/h2o/shape_calculators/h2otreemojo.py index 73063711..9d4adb9e 100644 --- a/onnxmltools/convert/h2o/shape_calculators/h2otreemojo.py +++ b/onnxmltools/convert/h2o/shape_calculators/h2otreemojo.py @@ -3,7 +3,7 @@ from ...common._registration import register_shape_calculator from ...common.shape_calculator import calculate_linear_regressor_output_shapes from ...common.utils import check_input_and_output_numbers, check_input_and_output_types -from ...common.data_types import (FloatTensorType, StringTensorType, Int64TensorType) +from ...common.data_types import FloatTensorType, StringTensorType, Int64TensorType def calculate_h2otree_output_shapes(operator): @@ -16,7 +16,9 @@ def calculate_h2otree_output_shapes(operator): def calculate_tree_classifier_output_shapes(operator, params): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=2) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) N = operator.inputs[0].type.shape[0] nclasses = params["nclasses"] operator.outputs[0].type = StringTensorType(shape=[N]) @@ -24,4 +26,4 @@ def calculate_tree_classifier_output_shapes(operator, params): operator.outputs[1].type = FloatTensorType([N, nclasses]) -register_shape_calculator('H2OTreeMojo', calculate_h2otree_output_shapes) +register_shape_calculator("H2OTreeMojo", calculate_h2otree_output_shapes) diff --git a/onnxmltools/convert/libsvm/__init__.py b/onnxmltools/convert/libsvm/__init__.py index 11db0855..017852fc 100644 --- a/onnxmltools/convert/libsvm/__init__.py +++ b/onnxmltools/convert/libsvm/__init__.py @@ -1,3 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 +from . import operator_converters, shape_calculators from .convert import convert diff --git a/onnxmltools/convert/libsvm/_parse.py b/onnxmltools/convert/libsvm/_parse.py index f2c1ad6f..9e135570 100644 --- a/onnxmltools/convert/libsvm/_parse.py +++ b/onnxmltools/convert/libsvm/_parse.py @@ -6,25 +6,27 @@ def _parse_libsvm_simple_model(scope, model, inputs): - ''' + """ This function handles all non-pipeline models. :param scope: Scope object :param model: A libsvm object (e.g., OneHotEncoder and LogisticRegression) :param inputs: A list of variables :return: A list of output variables which will be passed to next stage - ''' + """ if model.get_svm_type() in (0, 1): - label_variable = scope.declare_local_variable('label', FloatTensorType()) - probability_map_variable = scope.declare_local_variable('probabilities', FloatTensorType()) + label_variable = scope.declare_local_variable("label", FloatTensorType()) + probability_map_variable = scope.declare_local_variable( + "probabilities", FloatTensorType() + ) this_operator = scope.declare_local_operator("LibSvmSVC", model) this_operator.inputs = inputs this_operator.outputs.append(label_variable) this_operator.outputs.append(probability_map_variable) elif model.get_svm_type() in (4, 3): # We assume that all scikit-learn operator can only produce a single float tensor. - variable = scope.declare_local_variable('variable', FloatTensorType()) + variable = scope.declare_local_variable("variable", FloatTensorType()) this_operator = scope.declare_local_operator("LibSvmSVR", model) this_operator.inputs = inputs this_operator.outputs.append(variable) @@ -34,20 +36,26 @@ def _parse_libsvm_simple_model(scope, model, inputs): def _parse_libsvm(scope, model, inputs): - ''' - This is a delegate function. It doesn't nothing but invoke the correct parsing function according to the input + """ + This is a delegate function. It doesn't nothing but invoke + the correct parsing function according to the input model's type. + :param scope: Scope object :param model: A scikit-learn object (e.g., OneHotEncoder and LogisticRegression) :param inputs: A list of variables :return: The output variables produced by the input model - ''' + """ return _parse_libsvm_simple_model(scope, model, inputs) -def parse_libsvm(model, initial_types=None, target_opset=None, - custom_conversion_functions=None, - custom_shape_calculators=None): +def parse_libsvm( + model, + initial_types=None, + target_opset=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): # Put svmlib object into an abstract container so that our framework # can work seamlessly on models created # with different machine learning tools. @@ -55,16 +63,19 @@ def parse_libsvm(model, initial_types=None, target_opset=None, # Declare a computational graph. It will become a representation of # the input scikit-learn model after parsing. - topology = Topology(raw_model_container, default_batch_size='None', - initial_types=initial_types, - target_opset=target_opset, - custom_conversion_functions=custom_conversion_functions, - custom_shape_calculators=custom_shape_calculators) + topology = Topology( + raw_model_container, + default_batch_size="None", + initial_types=initial_types, + target_opset=target_opset, + custom_conversion_functions=custom_conversion_functions, + custom_shape_calculators=custom_shape_calculators, + ) # Declare an object to provide variables' and operators' naming mechanism. # In contrast to CoreML, one global scope # is enough for parsing scikit-learn models. - scope = topology.declare_scope('__root__') + scope = topology.declare_scope("__root__") # Declare input variables. They should be the inputs of the scikit-learn model # you want to convert into ONNX. diff --git a/onnxmltools/convert/libsvm/convert.py b/onnxmltools/convert/libsvm/convert.py index d1b5a819..7b79e8cc 100644 --- a/onnxmltools/convert/libsvm/convert.py +++ b/onnxmltools/convert/libsvm/convert.py @@ -7,40 +7,59 @@ from ._parse import parse_libsvm # Invoke the registration of all our converters and shape calculators -from . import shape_calculators -from . import operator_converters -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): +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, +): """ :param model: a libsvm model - :param initial_types: a python list. Each element is a tuple of a variable name and a type defined in data_types.py - :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 name: The name of the graph (type: GraphProto) + in the produced ONNX model (type: ModelProto) :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 scikit-learn 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 scikit-learn model """ if initial_types is None: - raise ValueError('Initial types are required. See usage of convert(...) in \ - onnxmltools.convert.libsvm.convert for details') + raise ValueError( + "Initial types are required. See usage of convert(...) in \ + onnxmltools.convert.libsvm.convert for details" + ) if name is None: name = str(uuid4().hex) target_opset = target_opset if target_opset else get_maximum_opset_supported() # Parse scikit-learn model as our internal data structure (i.e., Topology) - topology = parse_libsvm(model, initial_types, custom_conversion_functions, - custom_shape_calculators) + topology = parse_libsvm( + model, initial_types, custom_conversion_functions, custom_shape_calculators + ) # Infer variable shapes topology.compile() # Convert our Topology object into ONNX. The outcome is an ONNX model. - onnx_model = convert_topology(topology, name, doc_string, target_opset, targeted_onnx) + onnx_model = convert_topology( + topology, name, doc_string, target_opset, targeted_onnx + ) return onnx_model diff --git a/onnxmltools/convert/libsvm/operator_converters/SVMConverter.py b/onnxmltools/convert/libsvm/operator_converters/SVMConverter.py index 1aca71a0..0be18537 100644 --- a/onnxmltools/convert/libsvm/operator_converters/SVMConverter.py +++ b/onnxmltools/convert/libsvm/operator_converters/SVMConverter.py @@ -4,6 +4,7 @@ from ...common._registration import register_converter from ...common.utils import cast_list import numpy + try: from libsvm import svm except ImportError: @@ -15,41 +16,49 @@ class SVMConverter: """ Converts a SVM model trained with *svmlib*. """ + @staticmethod def validate(svm_node): try: - hasattr(svm_node, 'param') - hasattr(svm_node, 'SV') - hasattr(svm_node, 'nSV') - hasattr(svm_node, 'sv_coef') - hasattr(svm_node, 'l') - hasattr(svm_node.param, 'gamma') - hasattr(svm_node.param, 'coef0') - hasattr(svm_node.param, 'degree') - hasattr(svm_node.param, 'kernel_type') - hasattr(svm_node, 'rho') + hasattr(svm_node, "param") + hasattr(svm_node, "SV") + hasattr(svm_node, "nSV") + hasattr(svm_node, "sv_coef") + hasattr(svm_node, "l") + hasattr(svm_node.param, "gamma") + hasattr(svm_node.param, "coef0") + hasattr(svm_node.param, "degree") + hasattr(svm_node.param, "kernel_type") + hasattr(svm_node, "rho") except AttributeError as e: raise RuntimeError("Missing type from svm node:" + str(e)) - @staticmethod def get_sv(svm_node): labels = svm_node.get_labels() sv = svm_node.get_SV() if len(sv) == 0: - raise RuntimeError("No support vector machine. This usually happens with very small datasets or the training failed.") + raise RuntimeError( + "No support vector machine. This usually happens " + "with very small datasets or the training failed." + ) maxk = max(max(row.keys() for row in sv)) - mat = numpy.zeros((len(sv), maxk+1), dtype=numpy.float32) + mat = numpy.zeros((len(sv), maxk + 1), dtype=numpy.float32) for i, row in enumerate(sv): - for k,v in row.items(): + for k, v in row.items(): if k == -1: k = 0 try: mat[i, k] = v except IndexError: - raise RuntimeError("Issue with one dimension\nlabels={0}\n#sv={1}\nshape={2}\npos={3}x{4}-maxk={5}-svm.l={6}\nrow={7}".format(labels, nsv, mat.shape, i, k, maxk, svm_node.l, row)) + raise RuntimeError( + "Issue with one dimension\nlabels={0}\n#sv={1}\n" + "shape={2}\npos={3}x{4}-maxk={5}-svm.l={6}\nrow={7}".format( + labels, sv, mat.shape, i, k, maxk, svm_node.l, row + ) + ) # We do not consider the first row (class -1). mat = mat[:, 1:] @@ -66,18 +75,18 @@ def get_sv(svm_node): def convert(operator, scope, container, svm_node, inputs, model_name, nb_class): kt = svm_node.param.kernel_type if kt == svm.RBF: - kt = 'RBF' + kt = "RBF" elif kt == svm.SIGMOID: - kt = 'SIGMOID' + kt = "SIGMOID" elif kt == svm.POLY: - kt = 'POLY' + kt = "POLY" elif kt == svm.LINEAR: kt = "LINEAR" else: raise RuntimeError("Unexpected value for kernel: {0}".format(kt)) def copy_sv_coef(sv_coef): - nrc = svm_node.nr_class-1 + nrc = svm_node.nr_class - 1 res = numpy.zeros((svm_node.l, nrc), dtype=numpy.float64) for i in range(0, svm_node.l): for j in range(nrc): @@ -90,22 +99,35 @@ def copy_sv_coef(sv_coef): else: coef = numpy.array(svm_node.get_sv_coef()).ravel() - atts = dict(kernel_type=kt, - kernel_params=[float(_) for _ in [svm_node.param.gamma, svm_node.param.coef0, svm_node.param.degree]], - coefficients=list(coef.ravel())) + atts = dict( + kernel_type=kt, + kernel_params=[ + float(_) + for _ in [ + svm_node.param.gamma, + svm_node.param.coef0, + svm_node.param.degree, + ] + ], + coefficients=list(coef.ravel()), + ) + + return dict( + node="SVMConverter", + inputs=operator.input_full_names, + outputs=[o.full_name for o in operator.outputs], + op_domain="ai.onnx.ml", + attrs=atts, + ) - return dict(node='SVMConverter', inputs=operator.input_full_names, - outputs = [o.full_name for o in operator.outputs], - op_domain='ai.onnx.ml', attrs=atts) class SVCConverter(SVMConverter): - @staticmethod def validate(svm_node): SVMConverter.validate(svm_node) try: - hasattr(svm_node, 'probA') - hasattr(svm_node, 'probB') + hasattr(svm_node, "probA") + hasattr(svm_node, "probB") except AttributeError as e: raise RuntimeError("Missing type from svm node:" + str(e)) @@ -113,83 +135,99 @@ def validate(svm_node): def convert(operator, scope, container, svm_node, inputs): nbclass = len(svm_node.get_labels()) # See converter for sklearn. - nb = SVMConverter.convert(operator, scope, container, svm_node, inputs, "SVMClassifier", nbclass) - sign_rho = -1. + nb = SVMConverter.convert( + operator, scope, container, svm_node, inputs, "SVMClassifier", nbclass + ) + sign_rho = -1.0 st = svm_node.param.svm_type if svm_node.is_probability_model(): if st == svm.C_SVC or st == svm.NU_SVC: n_class = len(svm_node.get_labels()) - n = int(n_class*(n_class-1)/2) + n = int(n_class * (n_class - 1) / 2) probA = [svm_node.probA[i] for i in range(n)] probB = [svm_node.probB[i] for i in range(n)] nb["attrs"]["prob_a"] = probA nb["attrs"]["prob_b"] = probB - nb["attrs"]['rho'] = [svm_node.rho[i] * sign_rho for i in range(n)] + nb["attrs"]["rho"] = [svm_node.rho[i] * sign_rho for i in range(n)] else: - nb["attrs"]['rho'] = [svm_node.rho[0] * sign_rho] + nb["attrs"]["rho"] = [svm_node.rho[0] * sign_rho] elif st == svm.C_SVC or st == svm.NU_SVC: n_class = len(svm_node.get_labels()) - n = int(n_class*(n_class-1)/2) - nb["attrs"]['rho'] = [svm_node.rho[i] * sign_rho for i in range(n)] + n = int(n_class * (n_class - 1) / 2) + nb["attrs"]["rho"] = [svm_node.rho[i] * sign_rho for i in range(n)] else: - nb["attrs"]['rho'] = [svm_node.rho[0] * sign_rho] + nb["attrs"]["rho"] = [svm_node.rho[0] * sign_rho] class_labels = cast_list(int, svm_node.get_labels()) # Predictions are different when label are not sorted (multi-classification). class_labels.sort() - nb["attrs"]['classlabels_ints'] = class_labels - output_type = onnx_proto.TensorProto.INT64 + nb["attrs"]["classlabels_ints"] = class_labels + onnx_proto.TensorProto.INT64 - if len(nb['outputs']) != 2: - raise RuntimeError("The model outputs label and probabilities not {0}".format(nb['outputs'])) + if len(nb["outputs"]) != 2: + raise RuntimeError( + "The model outputs label and probabilities not {0}".format( + nb["outputs"] + ) + ) nbclass = len(svm_node.get_labels()) - nb["attrs"]['vectors_per_class'] = [svm_node.nSV[i] for i in range(nbclass)] - nb["attrs"]['post_transform'] = "NONE" - nb["attrs"]['support_vectors'] = SVCConverter.get_sv(svm_node) + nb["attrs"]["vectors_per_class"] = [svm_node.nSV[i] for i in range(nbclass)] + nb["attrs"]["post_transform"] = "NONE" + nb["attrs"]["support_vectors"] = SVCConverter.get_sv(svm_node) # Add a vec dictionizer to handle the map output - container.add_node('SVMClassifier', nb['inputs'], - nb['outputs'], op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('SVMClassifier'), - **nb['attrs']) + container.add_node( + "SVMClassifier", + nb["inputs"], + nb["outputs"], + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("SVMClassifier"), + **nb["attrs"] + ) class SVRConverter(SVMConverter): - @staticmethod def validate(svm_node): SVMConverter.validate(svm_node) try: - hasattr(svm_node, 'l') + hasattr(svm_node, "l") except AttributeError as e: raise RuntimeError("Missing type from svm node:" + str(e)) @staticmethod def convert(operator, scope, container, svm_node, inputs): - nb = SVMConverter.convert(operator, scope, container, svm_node, inputs, "SVMRegressor", 0) + nb = SVMConverter.convert( + operator, scope, container, svm_node, inputs, "SVMRegressor", 0 + ) - nb['attrs']["n_supports"] = svm_node.l - nb['attrs']['post_transform'] = "NONE" - nb['attrs']['rho'] = [-svm_node.rho[0]] - nb['attrs']['support_vectors'] = SVCConverter.get_sv(svm_node) + nb["attrs"]["n_supports"] = svm_node.l + nb["attrs"]["post_transform"] = "NONE" + nb["attrs"]["rho"] = [-svm_node.rho[0]] + nb["attrs"]["support_vectors"] = SVCConverter.get_sv(svm_node) - container.add_node('SVMRegressor', nb['inputs'], - nb['outputs'], op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('SVMRegressor'), - **nb['attrs']) + container.add_node( + "SVMRegressor", + nb["inputs"], + nb["outputs"], + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("SVMRegressor"), + **nb["attrs"] + ) class AnyLibSvmConverter: - @staticmethod def select(svm_node): if svm_node.param.svm_type in (svm.C_SVC, svm.NU_SVC): return SVCConverter if svm_node.param.svm_type in (svm.EPSILON_SVR, svm.NU_SVR): return SVRConverter - raise RuntimeError("svm_node type is unexpected '{0}'".format(svm_node.param.svm_type)) + raise RuntimeError( + "svm_node type is unexpected '{0}'".format(svm_node.param.svm_type) + ) @staticmethod def validate(svm_node): @@ -203,12 +241,9 @@ def convert(operator, scope, container, svm_node, inputs): def convert_libsvm(scope, operator, container): - inputs = operator.inputs model = operator.raw_operator converter = AnyLibSvmConverter - onnx_nodes = [] - outputs = None converter.validate(model) converter.convert(operator, scope, container, model, inputs) diff --git a/onnxmltools/convert/libsvm/shape_calculators/Classifier.py b/onnxmltools/convert/libsvm/shape_calculators/Classifier.py index b986e7d4..31483001 100644 --- a/onnxmltools/convert/libsvm/shape_calculators/Classifier.py +++ b/onnxmltools/convert/libsvm/shape_calculators/Classifier.py @@ -3,6 +3,7 @@ from ...common._registration import register_shape_calculator from ...common.data_types import FloatTensorType, Int64TensorType from ...common.utils import check_input_and_output_numbers + try: from libsvm.svm import C_SVC, NU_SVC except ImportError: @@ -13,11 +14,16 @@ def calculate_classifier_output_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=2) - N = (operator.inputs[0].type.shape[0] - if len(operator.inputs[0].type.shape) > 0 else None) + N = ( + operator.inputs[0].type.shape[0] + if len(operator.inputs[0].type.shape) > 0 + else None + ) if len(operator.outputs) != 2: - raise RuntimeError("Expect only two outputs not {0}".format(len(operator.outputs))) - svm_node = operator.raw_operator + raise RuntimeError( + "Expect only two outputs not {0}".format(len(operator.outputs)) + ) + svm_node = operator.raw_operator if svm_node.is_probability_model(): nc = svm_node.nr_class else: @@ -27,9 +33,9 @@ def calculate_classifier_output_shapes(operator): nc = svm_node.nr_class st = svm_node.param.svm_type if (st == C_SVC or st == NU_SVC) and nc > 2: - nc = (nc * (nc-1)) // 2 + nc = (nc * (nc - 1)) // 2 operator.outputs[0].type = Int64TensorType([N, 1]) operator.outputs[1].type = FloatTensorType([N, nc]) -register_shape_calculator('LibSvmSVC', calculate_classifier_output_shapes) +register_shape_calculator("LibSvmSVC", calculate_classifier_output_shapes) diff --git a/onnxmltools/convert/libsvm/shape_calculators/Regressor.py b/onnxmltools/convert/libsvm/shape_calculators/Regressor.py index 0d1069d1..e8b5d503 100644 --- a/onnxmltools/convert/libsvm/shape_calculators/Regressor.py +++ b/onnxmltools/convert/libsvm/shape_calculators/Regressor.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -from ...common._registration import register_shape_calculator, register_shape_calculator +from ...common._registration import register_shape_calculator from ...common.data_types import FloatTensorType from ...common.utils import check_input_and_output_numbers @@ -8,9 +8,12 @@ def calculate_regressor_output_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - N = (operator.inputs[0].type.shape[0] - if len(operator.inputs[0].type.shape) > 0 else None) + N = ( + operator.inputs[0].type.shape[0] + if len(operator.inputs[0].type.shape) > 0 + else None + ) operator.outputs[0].type = FloatTensorType([N, 1]) -register_shape_calculator('LibSvmSVR', calculate_regressor_output_shapes) +register_shape_calculator("LibSvmSVR", calculate_regressor_output_shapes) diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py index aee8869c..10b71996 100644 --- a/onnxmltools/convert/lightgbm/_parse.py +++ b/onnxmltools/convert/lightgbm/_parse.py @@ -4,84 +4,92 @@ from ..common._container import LightGbmModelContainer from ..common._topology import Topology -from ..common.data_types import (FloatTensorType, - SequenceType, DictionaryType, StringType, Int64Type) +from ..common.data_types import ( + FloatTensorType, + SequenceType, + DictionaryType, + StringType, + Int64Type, +) from lightgbm import LGBMClassifier, LGBMRegressor lightgbm_classifier_list = [LGBMClassifier] -# Associate scikit-learn types with our operator names. If two scikit-learn models share a single name, it means their +# Associate scikit-learn types with our operator names. +# If two scikit-learn models share a single name, it means their # are equivalent in terms of conversion. -lightgbm_operator_name_map = {LGBMClassifier: 'LgbmClassifier', - LGBMRegressor: 'LgbmRegressor'} +lightgbm_operator_name_map = { + LGBMClassifier: "LgbmClassifier", + LGBMRegressor: "LgbmRegressor", +} class WrappedBooster: - def __init__(self, booster): self.booster_ = booster self.n_features_ = self.booster_.feature_name() self.objective_ = self.get_objective() - if self.objective_.startswith('binary'): - self.operator_name = 'LgbmClassifier' + if self.objective_.startswith("binary"): + self.operator_name = "LgbmClassifier" self.classes_ = self._generate_classes(booster) - elif self.objective_.startswith('multiclass'): - self.operator_name = 'LgbmClassifier' + elif self.objective_.startswith("multiclass"): + self.operator_name = "LgbmClassifier" self.classes_ = self._generate_classes(booster) - elif self.objective_.startswith('regression'): - self.operator_name = 'LgbmRegressor' + elif self.objective_.startswith("regression"): + self.operator_name = "LgbmRegressor" else: raise NotImplementedError( - 'Unsupported LightGbm objective: %r.' % self.objective_) + "Unsupported LightGbm objective: %r." % self.objective_ + ) try: - average_output = self.booster_.attr('average_output') + average_output = self.booster_.attr("average_output") except AttributeError: average_output = self.booster_.params.get("average_output", None) if average_output: - self.boosting_type = 'rf' + self.boosting_type = "rf" else: # Other than random forest, other boosting types do not affect later conversion. # Here `gbdt` is chosen for no reason. - self.boosting_type = 'gbdt' + self.boosting_type = "gbdt" @staticmethod def _generate_classes(booster): if isinstance(booster, dict): - num_class = booster['num_class'] + num_class = booster["num_class"] else: try: - num_class = booster.attr('num_class') + num_class = booster.attr("num_class") except AttributeError: - num_class = booster.params.get('num_class', None) + num_class = booster.params.get("num_class", None) if num_class is None: dp = booster.dump_model(num_iteration=1) - num_class = dp['num_class'] + num_class = dp["num_class"] if num_class == 1: return numpy.asarray([0, 1]) return numpy.arange(num_class) def get_objective(self): "Returns the objective." - if hasattr(self, 'objective_') and self.objective_ is not None: + if hasattr(self, "objective_") and self.objective_ is not None: return self.objective_ try: - objective = self.booster_.attr('objective') + objective = self.booster_.attr("objective") except AttributeError: objective = self.booster_.params.get("objective", None) if objective is not None: return objective dp = self.booster_.dump_model(num_iteration=1) - return dp['objective'] + return dp["objective"] def _get_lightgbm_operator_name(model): - ''' + """ Get operator name of the input argument :param model: A lightgbm object. :return: A string which stands for the type of the input model in our conversion framework - ''' + """ if isinstance(model, WrappedBooster): return model.operator_name model_type = type(model) @@ -91,7 +99,7 @@ def _get_lightgbm_operator_name(model): def _parse_lightgbm_simple_model(scope, model, inputs, split=None): - ''' + """ This function handles all non-pipeline models. :param scope: Scope object @@ -100,73 +108,82 @@ def _parse_lightgbm_simple_model(scope, model, inputs, split=None): :param split: split TreeEnsembleRegressor into multiple node to reduce discrepancies :return: A list of output variables which will be passed to next stage - ''' + """ operator_name = _get_lightgbm_operator_name(model) this_operator = scope.declare_local_operator(operator_name, model) this_operator.split = split this_operator.inputs = inputs - if operator_name == 'LgbmClassifier': - # For classifiers, we may have two outputs, one for label and the other one for probabilities of all classes. - # Notice that their types here are not necessarily correct and they will be fixed in shape inference phase - label_variable = scope.declare_local_variable('lgbmlabel', FloatTensorType()) - probability_map_variable = scope.declare_local_variable('lgbmprobabilities', FloatTensorType()) + if operator_name == "LgbmClassifier": + # For classifiers, we may have two outputs, one for + # label and the other one for probabilities of all classes. + # Notice that their types here are not necessarily correct + # and they will be fixed in shape inference phase + label_variable = scope.declare_local_variable("lgbmlabel", FloatTensorType()) + probability_map_variable = scope.declare_local_variable( + "lgbmprobabilities", FloatTensorType() + ) this_operator.outputs.append(label_variable) this_operator.outputs.append(probability_map_variable) else: # We assume that all scikit-learn operator can only produce a single float tensor. - variable = scope.declare_local_variable('variable', FloatTensorType()) + variable = scope.declare_local_variable("variable", FloatTensorType()) this_operator.outputs.append(variable) return this_operator.outputs def _parse_sklearn_classifier(scope, model, inputs, zipmap=True): - probability_tensor = _parse_lightgbm_simple_model( - scope, model, inputs) - this_operator = scope.declare_local_operator('LgbmZipMap') + probability_tensor = _parse_lightgbm_simple_model(scope, model, inputs) + this_operator = scope.declare_local_operator("LgbmZipMap") this_operator.inputs = probability_tensor this_operator.zipmap = zipmap classes = model.classes_ label_type = Int64Type() - if (isinstance(model.classes_, list) and - isinstance(model.classes_[0], numpy.ndarray)): + if isinstance(model.classes_, list) and isinstance( + model.classes_[0], numpy.ndarray + ): # multi-label problem # this_operator.classlabels_int64s = list(range(0, len(classes))) raise NotImplementedError("multi-label is not supported") elif numpy.issubdtype(model.classes_.dtype, numpy.floating): classes = numpy.array(list(map(lambda x: int(x), classes))) if set(map(lambda x: float(x), classes)) != set(model.classes_): - raise RuntimeError("skl2onnx implicitly converts float class " - "labels into integers but at least one label " - "is not an integer. Class labels should " - "be integers or strings.") + raise RuntimeError( + "skl2onnx implicitly converts float class " + "labels into integers but at least one label " + "is not an integer. Class labels should " + "be integers or strings." + ) this_operator.classlabels_int64s = classes elif numpy.issubdtype(model.classes_.dtype, numpy.signedinteger): this_operator.classlabels_int64s = classes else: - classes = numpy.array([s.encode('utf-8') for s in classes]) + classes = numpy.array([s.encode("utf-8") for s in classes]) this_operator.classlabels_strings = classes label_type = StringType() - output_label = scope.declare_local_variable('label', label_type) + output_label = scope.declare_local_variable("label", label_type) if zipmap: output_probability = scope.declare_local_variable( - 'probabilities', - SequenceType(DictionaryType(label_type, FloatTensorType()))) + "probabilities", SequenceType(DictionaryType(label_type, FloatTensorType())) + ) else: output_probability = scope.declare_local_variable( - 'probabilities', FloatTensorType(shape=[None, len(classes)])) + "probabilities", FloatTensorType(shape=[None, len(classes)]) + ) this_operator.outputs.append(output_label) this_operator.outputs.append(output_probability) return this_operator.outputs def _parse_lightgbm(scope, model, inputs, zipmap=True, split=None): - ''' - This is a delegate function. It doesn't nothing but invoke the correct parsing function according to the input + """ + This is a delegate function. It doesn't nothing but + invoke the correct parsing function according to the input model's type. + :param scope: Scope object :param model: A lightgbm object :param inputs: A list of variables @@ -174,24 +191,33 @@ def _parse_lightgbm(scope, model, inputs, zipmap=True, split=None): :param split: split TreeEnsembleRegressor into multiple node to reduce discrepancies :return: The output variables produced by the input model - ''' + """ if isinstance(model, LGBMClassifier): return _parse_sklearn_classifier(scope, model, inputs, zipmap=zipmap) - if (isinstance(model, WrappedBooster) and - model.operator_name == 'LgbmClassifier'): + if isinstance(model, WrappedBooster) and model.operator_name == "LgbmClassifier": return _parse_sklearn_classifier(scope, model, inputs, zipmap=zipmap) return _parse_lightgbm_simple_model(scope, model, inputs, split=split) -def parse_lightgbm(model, initial_types=None, target_opset=None, - custom_conversion_functions=None, custom_shape_calculators=None, - zipmap=True, split=None): +def parse_lightgbm( + model, + initial_types=None, + target_opset=None, + custom_conversion_functions=None, + custom_shape_calculators=None, + zipmap=True, + split=None, +): raw_model_container = LightGbmModelContainer(model) - topology = Topology(raw_model_container, default_batch_size='None', - initial_types=initial_types, target_opset=target_opset, - custom_conversion_functions=custom_conversion_functions, - custom_shape_calculators=custom_shape_calculators) - scope = topology.declare_scope('__root__') + topology = Topology( + raw_model_container, + default_batch_size="None", + initial_types=initial_types, + target_opset=target_opset, + custom_conversion_functions=custom_conversion_functions, + custom_shape_calculators=custom_shape_calculators, + ) + scope = topology.declare_scope("__root__") inputs = [] for var_name, initial_type in initial_types: diff --git a/onnxmltools/convert/lightgbm/convert.py b/onnxmltools/convert/lightgbm/convert.py index 72409bb7..83799cbd 100644 --- a/onnxmltools/convert/lightgbm/convert.py +++ b/onnxmltools/convert/lightgbm/convert.py @@ -12,33 +12,53 @@ from . import operator_converters, shape_calculators # noqa -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, without_onnx_ml=False, zipmap=True, - split=None): - ''' +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, + without_onnx_ml=False, + zipmap=True, + split=None, +): + """ This function produces an equivalent ONNX model of the given lightgbm model. The supported lightgbm modules are listed below. - * `LGBMClassifiers `_ - * `LGBMRegressor `_ - * `Booster `_ + * `LGBMClassifiers + `_ + * `LGBMRegressor + `_ + * `Booster + `_ :param model: A LightGBM model - :param initial_types: a python list. Each element is a tuple of a variable name and a type defined in data_types.py - :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 name: The name of the graph (type: GraphProto) in the + produced ONNX model (type: ModelProto) :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 - :param without_onnx_ml: whether to generate a model composed by ONNX operators only, or to allow the converter + :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 + :param without_onnx_ml: whether to generate a model composed + by ONNX operators only, or to allow the converter :param zipmap: remove operator ZipMap from the ONNX graph :param split: this parameter is usefull to reduce the level of discrepancies for big regression forest (number of trees > 100). lightgbm does all the computation with double whereas ONNX is using floats. Instead of having one single node - TreeEnsembleRegressor, the converter splits it into multiple nodes TreeEnsembleRegressor, + TreeEnsembleRegressor, the converter splits it into + multiple nodes TreeEnsembleRegressor, casts the output in double and before additioning all the outputs. The final graph is slower but keeps the discrepancies constant (it is proportional to the number of trees in a node TreeEnsembleRegressor). @@ -47,30 +67,44 @@ def convert(model, name=None, initial_types=None, doc_string='', target_opset=No probabilities significantly reduces the discrepancies. to use ONNX-ML operators as well. :return: An ONNX model (type: ModelProto) which is equivalent to the input lightgbm model - ''' + """ if initial_types is None: - raise ValueError('Initial types are required. See usage of convert(...) in ' - 'onnxmltools.convert.lightgbm.convert for details') + raise ValueError( + "Initial types are required. See usage of convert(...) in " + "onnxmltools.convert.lightgbm.convert for details" + ) if without_onnx_ml and not hummingbird_installed(): raise RuntimeError( - 'Hummingbird is not installed. Please install hummingbird to use this feature: ' - 'pip install hummingbird-ml') + "Hummingbird is not installed. Please install hummingbird to use this feature: " + "pip install hummingbird-ml" + ) if isinstance(model, lightgbm.Booster): model = WrappedBooster(model) if name is None: name = str(uuid4().hex) target_opset = target_opset if target_opset else get_maximum_opset_supported() - topology = parse_lightgbm(model, initial_types, target_opset, custom_conversion_functions, - custom_shape_calculators, zipmap=zipmap, split=split) + topology = parse_lightgbm( + model, + initial_types, + target_opset, + custom_conversion_functions, + custom_shape_calculators, + zipmap=zipmap, + split=split, + ) topology.compile() - onnx_ml_model = convert_topology(topology, name, doc_string, target_opset, targeted_onnx) + onnx_ml_model = convert_topology( + topology, name, doc_string, target_opset, targeted_onnx + ) if without_onnx_ml: if zipmap: raise NotImplementedError( - "Conversion with zipmap operator is not implemented with hummingbird-ml.") + "Conversion with zipmap operator is not implemented with hummingbird-ml." + ) from hummingbird.ml import convert, constants + extra_config = {} # extra_config[constants.ONNX_INITIAL_TYPES] = initial_types extra_config[constants.ONNX_OUTPUT_MODEL_NAME] = name diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py index 7ca59985..fc374e04 100644 --- a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py +++ b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py @@ -8,7 +8,12 @@ import numpy as np from onnx import TensorProto from ...common._apply_operation import ( - apply_div, apply_reshape, apply_sub, apply_cast, apply_identity) + apply_div, + apply_reshape, + apply_sub, + apply_cast, + apply_identity, +) from ...common._registration import register_converter from ...common.tree_ensemble import get_default_tree_classifier_attribute_pairs from ....proto import onnx_proto @@ -17,6 +22,7 @@ def has_tqdm(): try: from tqdm import tqdm # noqa + return True except ImportError: return False @@ -25,22 +31,22 @@ def has_tqdm(): def _translate_split_criterion(criterion): # If the criterion is true, LightGBM use the left child. # Otherwise, right child is selected. - if criterion == '<=': - return 'BRANCH_LEQ' - elif criterion == '<': - return 'BRANCH_LT' - elif criterion == '>=': - return 'BRANCH_GTE' - elif criterion == '>': - return 'BRANCH_GT' - elif criterion == '==': - return 'BRANCH_EQ' - elif criterion == '!=': - return 'BRANCH_NEQ' + if criterion == "<=": + return "BRANCH_LEQ" + elif criterion == "<": + return "BRANCH_LT" + elif criterion == ">=": + return "BRANCH_GTE" + elif criterion == ">": + return "BRANCH_GT" + elif criterion == "==": + return "BRANCH_EQ" + elif criterion == "!=": + return "BRANCH_NEQ" else: raise ValueError( - 'Unsupported splitting criterion: %s. Only <=, ' - '<, >=, and > are allowed.') + "Unsupported splitting criterion: %s. Only <=, " "<, >=, and > are allowed." + ) def _create_node_id(node_id_pool): @@ -51,8 +57,7 @@ def _create_node_id(node_id_pool): return i -def _parse_tree_structure(tree_id, class_id, learning_rate, - tree_structure, attrs): +def _parse_tree_structure(tree_id, class_id, learning_rate, tree_structure, attrs): """ The pool of all nodes' indexes created when parsing a single tree. Different tree use different pools. @@ -64,14 +69,21 @@ def _parse_tree_structure(tree_id, class_id, learning_rate, node_pyid_pool[id(tree_structure)] = node_id # The root node is a leaf node. - if ('left_child' not in tree_structure or - 'right_child' not in tree_structure): - _parse_node(tree_id, class_id, node_id, node_id_pool, node_pyid_pool, - learning_rate, tree_structure, attrs) + if "left_child" not in tree_structure or "right_child" not in tree_structure: + _parse_node( + tree_id, + class_id, + node_id, + node_id_pool, + node_pyid_pool, + learning_rate, + tree_structure, + attrs, + ) return - left_pyid = id(tree_structure['left_child']) - right_pyid = id(tree_structure['right_child']) + left_pyid = id(tree_structure["left_child"]) + right_pyid = id(tree_structure["right_child"]) if left_pyid in node_pyid_pool: left_id = node_pyid_pool[left_pyid] @@ -89,56 +101,78 @@ def _parse_tree_structure(tree_id, class_id, learning_rate, node_pyid_pool[right_pyid] = right_id right_parse = True - attrs['nodes_treeids'].append(tree_id) - attrs['nodes_nodeids'].append(node_id) + attrs["nodes_treeids"].append(tree_id) + attrs["nodes_nodeids"].append(node_id) - attrs['nodes_featureids'].append(tree_structure['split_feature']) - attrs['nodes_modes'].append( - _translate_split_criterion(tree_structure['decision_type'])) - if isinstance(tree_structure['threshold'], str): + attrs["nodes_featureids"].append(tree_structure["split_feature"]) + attrs["nodes_modes"].append( + _translate_split_criterion(tree_structure["decision_type"]) + ) + if isinstance(tree_structure["threshold"], str): try: - attrs['nodes_values'].append(float(tree_structure['threshold'])) + attrs["nodes_values"].append(float(tree_structure["threshold"])) except ValueError: import pprint + text = pprint.pformat(tree_structure) if len(text) > 100000: text = text[:100000] + "\n..." - raise TypeError("threshold must be a number not '{}'" - "\n{}".format(tree_structure['threshold'], text)) + raise TypeError( + "threshold must be a number not '{}'" + "\n{}".format(tree_structure["threshold"], text) + ) else: - attrs['nodes_values'].append(tree_structure['threshold']) + attrs["nodes_values"].append(tree_structure["threshold"]) # Assume left is the true branch and right is the false branch - attrs['nodes_truenodeids'].append(left_id) - attrs['nodes_falsenodeids'].append(right_id) - if tree_structure['default_left']: - if tree_structure["missing_type"] == 'None' and float(tree_structure['threshold']) < 0.0: - attrs['nodes_missing_value_tracks_true'].append(0) + attrs["nodes_truenodeids"].append(left_id) + attrs["nodes_falsenodeids"].append(right_id) + if tree_structure["default_left"]: + if ( + tree_structure["missing_type"] == "None" + and float(tree_structure["threshold"]) < 0.0 + ): + attrs["nodes_missing_value_tracks_true"].append(0) else: - attrs['nodes_missing_value_tracks_true'].append(1) + attrs["nodes_missing_value_tracks_true"].append(1) else: - attrs['nodes_missing_value_tracks_true'].append(0) - attrs['nodes_hitrates'].append(1.) + attrs["nodes_missing_value_tracks_true"].append(0) + attrs["nodes_hitrates"].append(1.0) if left_parse: _parse_node( - tree_id, class_id, left_id, node_id_pool, node_pyid_pool, - learning_rate, tree_structure['left_child'], attrs) + tree_id, + class_id, + left_id, + node_id_pool, + node_pyid_pool, + learning_rate, + tree_structure["left_child"], + attrs, + ) if right_parse: _parse_node( - tree_id, class_id, right_id, node_id_pool, node_pyid_pool, - learning_rate, tree_structure['right_child'], attrs) + tree_id, + class_id, + right_id, + node_id_pool, + node_pyid_pool, + learning_rate, + tree_structure["right_child"], + attrs, + ) -def _parse_node(tree_id, class_id, node_id, node_id_pool, node_pyid_pool, - learning_rate, node, attrs): +def _parse_node( + tree_id, class_id, node_id, node_id_pool, node_pyid_pool, learning_rate, node, attrs +): """ Parses nodes. """ - if ((hasattr(node, 'left_child') and hasattr(node, 'right_child')) or - ('left_child' in node and 'right_child' in node)): - - left_pyid = id(node['left_child']) - right_pyid = id(node['right_child']) + if (hasattr(node, "left_child") and hasattr(node, "right_child")) or ( + "left_child" in node and "right_child" in node + ): + left_pyid = id(node["left_child"]) + right_pyid = id(node["right_child"]) if left_pyid in node_pyid_pool: left_id = node_pyid_pool[left_pyid] @@ -156,79 +190,95 @@ def _parse_node(tree_id, class_id, node_id, node_id_pool, node_pyid_pool, node_pyid_pool[right_pyid] = right_id right_parse = True - attrs['nodes_treeids'].append(tree_id) - attrs['nodes_nodeids'].append(node_id) + attrs["nodes_treeids"].append(tree_id) + attrs["nodes_nodeids"].append(node_id) - attrs['nodes_featureids'].append(node['split_feature']) - attrs['nodes_modes'].append( - _translate_split_criterion(node['decision_type'])) - if isinstance(node['threshold'], str): + attrs["nodes_featureids"].append(node["split_feature"]) + attrs["nodes_modes"].append(_translate_split_criterion(node["decision_type"])) + if isinstance(node["threshold"], str): try: - attrs['nodes_values'].append(float(node['threshold'])) + attrs["nodes_values"].append(float(node["threshold"])) except ValueError: import pprint + text = pprint.pformat(node) if len(text) > 100000: text = text[:100000] + "\n..." - raise TypeError("threshold must be a number not '{}'" - "\n{}".format(node['threshold'], text)) + raise TypeError( + "threshold must be a number not '{}'" + "\n{}".format(node["threshold"], text) + ) else: - attrs['nodes_values'].append(node['threshold']) + attrs["nodes_values"].append(node["threshold"]) # Assume left is the true branch # and right is the false branch - attrs['nodes_truenodeids'].append(left_id) - attrs['nodes_falsenodeids'].append(right_id) - if node['default_left']: - if node['missing_type'] == 'None' and float(node['threshold']) < 0.0: - attrs['nodes_missing_value_tracks_true'].append(0) + attrs["nodes_truenodeids"].append(left_id) + attrs["nodes_falsenodeids"].append(right_id) + if node["default_left"]: + if node["missing_type"] == "None" and float(node["threshold"]) < 0.0: + attrs["nodes_missing_value_tracks_true"].append(0) else: - attrs['nodes_missing_value_tracks_true'].append(1) + attrs["nodes_missing_value_tracks_true"].append(1) else: - attrs['nodes_missing_value_tracks_true'].append(0) - attrs['nodes_hitrates'].append(1.) + attrs["nodes_missing_value_tracks_true"].append(0) + attrs["nodes_hitrates"].append(1.0) # Recursively dive into the child nodes if left_parse: _parse_node( - tree_id, class_id, left_id, node_id_pool, node_pyid_pool, - learning_rate, node['left_child'], attrs) + tree_id, + class_id, + left_id, + node_id_pool, + node_pyid_pool, + learning_rate, + node["left_child"], + attrs, + ) if right_parse: _parse_node( - tree_id, class_id, right_id, node_id_pool, node_pyid_pool, - learning_rate, node['right_child'], attrs) - elif hasattr(node, 'left_child') or hasattr(node, 'right_child'): - raise ValueError('Need two branches') + tree_id, + class_id, + right_id, + node_id_pool, + node_pyid_pool, + learning_rate, + node["right_child"], + attrs, + ) + elif hasattr(node, "left_child") or hasattr(node, "right_child"): + raise ValueError("Need two branches") else: # Node attributes - attrs['nodes_treeids'].append(tree_id) - attrs['nodes_nodeids'].append(node_id) - attrs['nodes_featureids'].append(0) - attrs['nodes_modes'].append('LEAF') + attrs["nodes_treeids"].append(tree_id) + attrs["nodes_nodeids"].append(node_id) + attrs["nodes_featureids"].append(0) + attrs["nodes_modes"].append("LEAF") # Leaf node has no threshold. # A zero is appended but it will never be used. - attrs['nodes_values'].append(0.) + attrs["nodes_values"].append(0.0) # Leaf node has no child. # A zero is appended but it will never be used. - attrs['nodes_truenodeids'].append(0) + attrs["nodes_truenodeids"].append(0) # Leaf node has no child. # A zero is appended but it will never be used. - attrs['nodes_falsenodeids'].append(0) + attrs["nodes_falsenodeids"].append(0) # Leaf node has no split function. # A zero is appended but it will never be used. - attrs['nodes_missing_value_tracks_true'].append(0) - attrs['nodes_hitrates'].append(1.) + attrs["nodes_missing_value_tracks_true"].append(0) + attrs["nodes_hitrates"].append(1.0) # Leaf attributes - attrs['class_treeids'].append(tree_id) - attrs['class_nodeids'].append(node_id) - attrs['class_ids'].append(class_id) - attrs['class_weights'].append( - float(node['leaf_value']) * learning_rate) + attrs["class_treeids"].append(tree_id) + attrs["class_nodeids"].append(node_id) + attrs["class_ids"].append(class_id) + attrs["class_weights"].append(float(node["leaf_value"]) * learning_rate) -def dump_booster_model(self, num_iteration=None, start_iteration=0, - importance_type='split', verbose=0): +def dump_booster_model( + self, num_iteration=None, start_iteration=0, importance_type="split", verbose=0 +): """ Dumps Booster to JSON format. @@ -262,9 +312,10 @@ def dump_booster_model(self, num_iteration=None, start_iteration=0, into ONNX of such model. The function overwrites the `json.load` to fastly extract nodes. """ - if getattr(self, 'is_mock', False): + if getattr(self, "is_mock", False): return self.dump_model(), None from lightgbm.basic import _LIB, _safe_call + try: # lightgbm >= 4.0 from lightgbm.basic import ( @@ -290,34 +341,59 @@ def dump_booster_model(self, num_iteration=None, start_iteration=0, handle = self._handle except AttributeError: handle = self.handle - _safe_call(_LIB.LGBM_BoosterDumpModel( - handle, - ctypes.c_int(start_iteration), - ctypes.c_int(num_iteration), - ctypes.c_int(importance_type_int), - ctypes.c_int64(buffer_len), - ctypes.byref(tmp_out_len), - ptr_string_buffer)) + _safe_call( + _LIB.LGBM_BoosterDumpModel( + handle, + ctypes.c_int(start_iteration), + ctypes.c_int(num_iteration), + ctypes.c_int(importance_type_int), + ctypes.c_int64(buffer_len), + ctypes.byref(tmp_out_len), + ptr_string_buffer, + ) + ) actual_len = tmp_out_len.value # if buffer length is not long enough, reallocate a buffer if actual_len > buffer_len: string_buffer = ctypes.create_string_buffer(actual_len) - ptr_string_buffer = ctypes.c_char_p( - *[ctypes.addressof(string_buffer)]) + ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)]) try: # lightgbm >= 4.0 handle = self._handle except AttributeError: # lightgbm < 4.0 handle = self.handle - _safe_call(_LIB.LGBM_BoosterDumpModel( - handle, - ctypes.c_int(start_iteration), - ctypes.c_int(num_iteration), - ctypes.c_int(importance_type_int), - ctypes.c_int64(actual_len), - ctypes.byref(tmp_out_len), - ptr_string_buffer)) + _safe_call( + _LIB.LGBM_BoosterDumpModel( + handle, + ctypes.c_int(start_iteration), + ctypes.c_int(num_iteration), + ctypes.c_int(importance_type_int), + ctypes.c_int64(buffer_len), + ctypes.byref(tmp_out_len), + ptr_string_buffer, + ) + ) + actual_len = tmp_out_len.value + # if buffer length is not long enough, reallocate a buffer + if actual_len > buffer_len: + string_buffer = ctypes.create_string_buffer(actual_len) + ptr_string_buffer = ctypes.c_char_p(*[ctypes.addressof(string_buffer)]) + try: + handle = self._handle + except AttributeError: + handle = self.handle + _safe_call( + _LIB.LGBM_BoosterDumpModel( + handle, + ctypes.c_int(start_iteration), + ctypes.c_int(num_iteration), + ctypes.c_int(importance_type_int), + ctypes.c_int64(actual_len), + ctypes.byref(tmp_out_len), + ptr_string_buffer, + ) + ) class Hook(json.JSONDecoder): """ @@ -325,10 +401,9 @@ class Hook(json.JSONDecoder): a decision into a different container in order to walk through all nodes in a much faster way than going through the architecture. """ - def __init__(self, *args, info=None, n_trees=None, verbose=0, - **kwargs): - json.JSONDecoder.__init__( - self, object_hook=self.hook, *args, **kwargs) + + def __init__(self, *args, info=None, n_trees=None, verbose=0, **kwargs): + json.JSONDecoder.__init__(self, object_hook=self.hook, *args, **kwargs) self.nodes = [] self.buffer = [] self.info = info @@ -337,6 +412,7 @@ def __init__(self, *args, info=None, n_trees=None, verbose=0, self.stored = 0 if verbose >= 2 and n_trees is not None and has_tqdm(): from tqdm import tqdm + self.loop = tqdm(total=n_trees) self.loop.set_description("dump_booster") else: @@ -349,23 +425,25 @@ def hook(self, obj): a decision into a different container. """ # Every obj goes through this function from the leaves to the root. - if 'tree_info' in obj: - self.info['decision_nodes'] = self.nodes + if "tree_info" in obj: + self.info["decision_nodes"] = self.nodes if self.n_trees is not None and len(self.nodes) != self.n_trees: raise RuntimeError( - "Unexpected number of trees %d (expecting %d)." % ( - len(self.nodes), self.n_trees)) + "Unexpected number of trees %d (expecting %d)." + % (len(self.nodes), self.n_trees) + ) self.nodes = [] if self.loop is not None: self.loop.close() - if 'tree_structure' in obj: + if "tree_structure" in obj: self.nodes.append(self.buffer) if self.loop is not None: self.loop.update(len(self.nodes)) if len(self.nodes) % 10 == 0: self.loop.set_description( - "dump_booster: %d/%d trees, %d nodes" % ( - len(self.nodes), self.n_trees, self.stored)) + "dump_booster: %d/%d trees, %d nodes" + % (len(self.nodes), self.n_trees, self.stored) + ) self.buffer = [] if "decision_type" in obj: self.buffer.append(obj) @@ -375,11 +453,16 @@ def hook(self, obj): if verbose >= 2: print("[dump_booster_model] to_json") info = {} - ret = json.loads(string_buffer.value.decode('utf-8'), cls=Hook, - info=info, n_trees=self.num_trees(), verbose=verbose) - ret['pandas_categorical'] = json.loads( - json.dumps(self.pandas_categorical, - default=jdwn)) + ret = json.loads( + string_buffer.value.decode("utf-8"), + cls=Hook, + info=info, + n_trees=self.num_trees(), + verbose=verbose, + ) + ret["pandas_categorical"] = json.loads( + json.dumps(self.pandas_categorical, default=jdwn) + ) if verbose >= 2: print("[dump_booster_model] end.") return ret, info @@ -390,39 +473,47 @@ def _split_tree_ensemble_atts(attrs, split): Splits the attributes of a TreeEnsembleRegressor into multiple trees in order to do the summation in double instead of floats. """ - trees_id = list(sorted(set(attrs['nodes_treeids']))) + trees_id = list(sorted(set(attrs["nodes_treeids"]))) results = [] index = 0 while index < len(trees_id): index2 = min(index + split, len(trees_id)) - subset = set(trees_id[index: index2]) + subset = set(trees_id[index:index2]) indices_node = [] indices_target = [] - for j, v in enumerate(attrs['nodes_treeids']): + for j, v in enumerate(attrs["nodes_treeids"]): if v in subset: indices_node.append(j) - for j, v in enumerate(attrs['target_treeids']): + for j, v in enumerate(attrs["target_treeids"]): if v in subset: indices_target.append(j) - if (len(indices_node) >= len(attrs['nodes_treeids']) or - len(indices_target) >= len(attrs['target_treeids'])): + if len(indices_node) >= len(attrs["nodes_treeids"]) or len( + indices_target + ) >= len(attrs["target_treeids"]): raise RuntimeError( # pragma: no cover "Initial attributes are not consistant." "\nindex=%r index2=%r subset=%r" "\nnodes_treeids=%r\ntarget_treeids=%r" - "\nindices_node=%r\nindices_target=%r" % ( - index, index2, subset, - attrs['nodes_treeids'], attrs['target_treeids'], - indices_node, indices_target)) + "\nindices_node=%r\nindices_target=%r" + % ( + index, + index2, + subset, + attrs["nodes_treeids"], + attrs["target_treeids"], + indices_node, + indices_target, + ) + ) ats = {} for name, att in attrs.items(): - if name == 'nodes_treeids': + if name == "nodes_treeids": new_att = [att[i] for i in indices_node] new_att = [i - att[0] for i in new_att] - elif name == 'target_treeids': + elif name == "target_treeids": new_att = [att[i] for i in indices_target] new_att = [i - att[0] for i in new_att] elif name.startswith("nodes_"): @@ -431,7 +522,7 @@ def _split_tree_ensemble_atts(attrs, split): elif name.startswith("target_"): new_att = [att[i] for i in indices_target] assert len(new_att) == len(indices_target) - elif name == 'name': + elif name == "name": new_att = "%s%d" % (att, len(results)) else: new_att = att @@ -447,94 +538,96 @@ def convert_lightgbm(scope, operator, container): """ Converters for *lightgbm*. """ - verbose = getattr(container, 'verbose', 0) + verbose = getattr(container, "verbose", 0) gbm_model = operator.raw_operator gbm_text, info = dump_booster_model(gbm_model.booster_, verbose=verbose) modify_tree_for_rule_in_set(gbm_text, use_float=True, verbose=verbose, info=info) attrs = get_default_tree_classifier_attribute_pairs() - attrs['name'] = operator.full_name + attrs["name"] = operator.full_name # Create different attributes for classifier and # regressor, respectively post_transform = None - if gbm_text['objective'].startswith('binary'): + if gbm_text["objective"].startswith("binary"): n_classes = 1 - attrs['post_transform'] = 'LOGISTIC' - elif gbm_text['objective'].startswith('multiclass'): - n_classes = gbm_text['num_class'] - attrs['post_transform'] = 'SOFTMAX' - elif gbm_text['objective'].startswith(('regression', 'quantile')): + attrs["post_transform"] = "LOGISTIC" + elif gbm_text["objective"].startswith("multiclass"): + n_classes = gbm_text["num_class"] + attrs["post_transform"] = "SOFTMAX" + elif gbm_text["objective"].startswith(("regression", "quantile")): n_classes = 1 # Regressor has only one output variable - attrs['post_transform'] = 'NONE' - attrs['n_targets'] = n_classes - elif gbm_text['objective'].startswith(('poisson', 'gamma')): + attrs["post_transform"] = "NONE" + attrs["n_targets"] = n_classes + elif gbm_text["objective"].startswith(("poisson", "gamma")): n_classes = 1 # Regressor has only one output variable - attrs['n_targets'] = n_classes + attrs["n_targets"] = n_classes # 'Exp' is not a supported post_transform value in the ONNX spec yet, # so we need to add an 'Exp' post transform node to the model - attrs['post_transform'] = 'NONE' + attrs["post_transform"] = "NONE" post_transform = "Exp" else: raise RuntimeError( "LightGBM objective should be cleaned already not '{}'.".format( - gbm_text['objective'])) + gbm_text["objective"] + ) + ) # Use the same algorithm to parse the tree - for i, tree in enumerate(gbm_text['tree_info']): + for i, tree in enumerate(gbm_text["tree_info"]): tree_id = i class_id = tree_id % n_classes # tree['shrinkage'] --> LightGbm provides figures with it already. - learning_rate = 1. + learning_rate = 1.0 _parse_tree_structure( - tree_id, class_id, learning_rate, tree['tree_structure'], attrs) + tree_id, class_id, learning_rate, tree["tree_structure"], attrs + ) # Sort nodes_* attributes. For one tree, its node indexes # should appear in an ascent order in nodes_nodeids. Nodes # from a tree with a smaller tree index should appear # before trees with larger indexes in nodes_nodeids. - node_numbers_per_tree = Counter(attrs['nodes_treeids']) + node_numbers_per_tree = Counter(attrs["nodes_treeids"]) tree_number = len(node_numbers_per_tree.keys()) accumulated_node_numbers = [0] * tree_number for i in range(1, tree_number): accumulated_node_numbers[i] = ( - accumulated_node_numbers[i - 1] + node_numbers_per_tree[i - 1]) + accumulated_node_numbers[i - 1] + node_numbers_per_tree[i - 1] + ) global_node_indexes = [] - for i in range(len(attrs['nodes_nodeids'])): - tree_id = attrs['nodes_treeids'][i] - node_id = attrs['nodes_nodeids'][i] - global_node_indexes.append( - accumulated_node_numbers[tree_id] + node_id) + for i in range(len(attrs["nodes_nodeids"])): + tree_id = attrs["nodes_treeids"][i] + node_id = attrs["nodes_nodeids"][i] + global_node_indexes.append(accumulated_node_numbers[tree_id] + node_id) for k, v in attrs.items(): - if k.startswith('nodes_'): - merged_indexes = zip( - copy.deepcopy(global_node_indexes), v) - sorted_list = [pair[1] - for pair in sorted(merged_indexes, - key=lambda x: x[0])] + if k.startswith("nodes_"): + merged_indexes = zip(copy.deepcopy(global_node_indexes), v) + sorted_list = [ + pair[1] for pair in sorted(merged_indexes, key=lambda x: x[0]) + ] attrs[k] = sorted_list # Create ONNX object - if (gbm_text['objective'].startswith('binary') or - gbm_text['objective'].startswith('multiclass')): + if gbm_text["objective"].startswith("binary") or gbm_text["objective"].startswith( + "multiclass" + ): # Prepare label information for both of TreeEnsembleClassifier class_type = onnx_proto.TensorProto.STRING - if all(isinstance(i, (numbers.Real, bool, np.bool_)) - for i in gbm_model.classes_): + if all( + isinstance(i, (numbers.Real, bool, np.bool_)) for i in gbm_model.classes_ + ): class_type = onnx_proto.TensorProto.INT64 class_labels = [int(i) for i in gbm_model.classes_] - attrs['classlabels_int64s'] = class_labels + attrs["classlabels_int64s"] = class_labels elif all(isinstance(i, str) for i in gbm_model.classes_): class_labels = [str(i) for i in gbm_model.classes_] - attrs['classlabels_strings'] = class_labels + attrs["classlabels_strings"] = class_labels else: - raise ValueError( - 'Only string and integer class labels are allowed') + raise ValueError("Only string and integer class labels are allowed") # Create tree classifier - probability_tensor_name = scope.get_unique_variable_name( - 'probability_tensor') - label_tensor_name = scope.get_unique_variable_name('label_tensor') + probability_tensor_name = scope.get_unique_variable_name("probability_tensor") + label_tensor_name = scope.get_unique_variable_name("label_tensor") # onnx does not support int and float values for a float tensor update = {} @@ -548,88 +641,116 @@ def convert_lightgbm(scope, operator, container): attrs.update(update) container.add_node( - 'TreeEnsembleClassifier', operator.input_full_names, + "TreeEnsembleClassifier", + operator.input_full_names, [label_tensor_name, probability_tensor_name], - op_domain='ai.onnx.ml', **attrs) + op_domain="ai.onnx.ml", + **attrs + ) prob_tensor = probability_tensor_name - if gbm_model.boosting_type == 'rf': - col_index_name = scope.get_unique_variable_name('col_index') - first_col_name = scope.get_unique_variable_name('first_col') - zeroth_col_name = scope.get_unique_variable_name('zeroth_col') - denominator_name = scope.get_unique_variable_name('denominator') + if gbm_model.boosting_type == "rf": + col_index_name = scope.get_unique_variable_name("col_index") + first_col_name = scope.get_unique_variable_name("first_col") + zeroth_col_name = scope.get_unique_variable_name("zeroth_col") + denominator_name = scope.get_unique_variable_name("denominator") modified_first_col_name = scope.get_unique_variable_name( - 'modified_first_col') - unit_float_tensor_name = scope.get_unique_variable_name( - 'unit_float_tensor') - merged_prob_name = scope.get_unique_variable_name('merged_prob') - predicted_label_name = scope.get_unique_variable_name( - 'predicted_label') - classes_name = scope.get_unique_variable_name('classes') - final_label_name = scope.get_unique_variable_name('final_label') + "modified_first_col" + ) + unit_float_tensor_name = scope.get_unique_variable_name("unit_float_tensor") + merged_prob_name = scope.get_unique_variable_name("merged_prob") + predicted_label_name = scope.get_unique_variable_name("predicted_label") + classes_name = scope.get_unique_variable_name("classes") + final_label_name = scope.get_unique_variable_name("final_label") container.add_initializer( - col_index_name, onnx_proto.TensorProto.INT64, [], [1]) + col_index_name, onnx_proto.TensorProto.INT64, [], [1] + ) container.add_initializer( - unit_float_tensor_name, onnx_proto.TensorProto.FLOAT, - [], [1.0]) + unit_float_tensor_name, onnx_proto.TensorProto.FLOAT, [], [1.0] + ) container.add_initializer( - denominator_name, onnx_proto.TensorProto.FLOAT, [], - [100.0]) - container.add_initializer(classes_name, class_type, - [len(class_labels)], class_labels) + denominator_name, onnx_proto.TensorProto.FLOAT, [], [100.0] + ) + container.add_initializer( + classes_name, class_type, [len(class_labels)], class_labels + ) container.add_node( - 'ArrayFeatureExtractor', + "ArrayFeatureExtractor", [probability_tensor_name, col_index_name], first_col_name, - name=scope.get_unique_operator_name( - 'ArrayFeatureExtractor'), - op_domain='ai.onnx.ml') - apply_div(scope, [first_col_name, denominator_name], - modified_first_col_name, container, broadcast=1) + name=scope.get_unique_operator_name("ArrayFeatureExtractor"), + op_domain="ai.onnx.ml", + ) + apply_div( + scope, + [first_col_name, denominator_name], + modified_first_col_name, + container, + broadcast=1, + ) apply_sub( - scope, [unit_float_tensor_name, modified_first_col_name], - zeroth_col_name, container, broadcast=1) + scope, + [unit_float_tensor_name, modified_first_col_name], + zeroth_col_name, + container, + broadcast=1, + ) container.add_node( - 'Concat', [zeroth_col_name, modified_first_col_name], + "Concat", + [zeroth_col_name, modified_first_col_name], merged_prob_name, - name=scope.get_unique_operator_name('Concat'), axis=1) + name=scope.get_unique_operator_name("Concat"), + axis=1, + ) container.add_node( - 'ArgMax', merged_prob_name, + "ArgMax", + merged_prob_name, predicted_label_name, - name=scope.get_unique_operator_name('ArgMax'), axis=1) + name=scope.get_unique_operator_name("ArgMax"), + axis=1, + ) container.add_node( - 'ArrayFeatureExtractor', [classes_name, predicted_label_name], + "ArrayFeatureExtractor", + [classes_name, predicted_label_name], + final_label_name, + name=scope.get_unique_operator_name("ArrayFeatureExtractor"), + op_domain="ai.onnx.ml", + ) + apply_reshape( + scope, final_label_name, - name=scope.get_unique_operator_name('ArrayFeatureExtractor'), - op_domain='ai.onnx.ml') - apply_reshape(scope, final_label_name, - operator.outputs[0].full_name, - container, desired_shape=[-1, ]) + operator.outputs[0].full_name, + container, + desired_shape=[ + -1, + ], + ) prob_tensor = merged_prob_name else: - container.add_node('Identity', label_tensor_name, - operator.outputs[0].full_name, - name=scope.get_unique_operator_name('Identity')) + container.add_node( + "Identity", + label_tensor_name, + operator.outputs[0].full_name, + name=scope.get_unique_operator_name("Identity"), + ) # Convert probability tensor to probability map # (keys are labels while values are the associated probabilities) - container.add_node('Identity', prob_tensor, - operator.outputs[1].full_name) + container.add_node("Identity", prob_tensor, operator.outputs[1].full_name) else: # Create tree regressor - output_name = scope.get_unique_variable_name('output') + output_name = scope.get_unique_variable_name("output") - keys_to_be_renamed = list( - k for k in attrs if k.startswith('class_')) + keys_to_be_renamed = list(k for k in attrs if k.startswith("class_")) for k in keys_to_be_renamed: # Rename class_* attribute to target_* # because TreeEnsebmleClassifier # and TreeEnsembleClassifier have different ONNX attributes - attrs['target' + k[5:]] = copy.deepcopy(attrs[k]) + attrs["target" + k[5:]] = copy.deepcopy(attrs[k]) del attrs[k] # onnx does not support int and float values for a float tensor @@ -643,55 +764,83 @@ def convert_lightgbm(scope, operator, container): update[k] = [float(x) for x in v] attrs.update(update) - split = getattr(operator, 'split', None) + split = getattr(operator, "split", None) if split in (None, -1): container.add_node( - 'TreeEnsembleRegressor', operator.input_full_names, - output_name, op_domain='ai.onnx.ml', **attrs) + "TreeEnsembleRegressor", + operator.input_full_names, + output_name, + op_domain="ai.onnx.ml", + **attrs + ) else: tree_attrs = _split_tree_ensemble_atts(attrs, split) tree_nodes = [] for i, ats in enumerate(tree_attrs): - tree_name = scope.get_unique_variable_name('tree%d' % i) + tree_name = scope.get_unique_variable_name("tree%d" % i) container.add_node( - 'TreeEnsembleRegressor', operator.input_full_names, - tree_name, op_domain='ai.onnx.ml', **ats) - cast_name = scope.get_unique_variable_name('dtree%d' % i) + "TreeEnsembleRegressor", + operator.input_full_names, + tree_name, + op_domain="ai.onnx.ml", + **ats + ) + cast_name = scope.get_unique_variable_name("dtree%d" % i) container.add_node( - 'Cast', tree_name, cast_name, to=TensorProto.DOUBLE, # pylint: disable=E1101 - name=scope.get_unique_operator_name("dtree%d" % i)) + "Cast", + tree_name, + cast_name, + to=TensorProto.DOUBLE, # pylint: disable=E1101 + name=scope.get_unique_operator_name("dtree%d" % i), + ) tree_nodes.append(cast_name) - cast_name = scope.get_unique_variable_name('ftrees') + cast_name = scope.get_unique_variable_name("ftrees") container.add_node( - 'Sum', tree_nodes, cast_name, - name=scope.get_unique_operator_name("sumtree%d" % len(tree_nodes))) + "Sum", + tree_nodes, + cast_name, + name=scope.get_unique_operator_name("sumtree%d" % len(tree_nodes)), + ) container.add_node( - 'Cast', cast_name, output_name, to=TensorProto.FLOAT, # pylint: disable=E1101 - name=scope.get_unique_operator_name("dtree%d" % i)) - if gbm_model.boosting_type == 'rf': - denominator_name = scope.get_unique_variable_name('denominator') + "Cast", + cast_name, + output_name, + to=TensorProto.FLOAT, # pylint: disable=E1101 + name=scope.get_unique_operator_name("dtree%d" % i), + ) + if gbm_model.boosting_type == "rf": + denominator_name = scope.get_unique_variable_name("denominator") container.add_initializer( - denominator_name, onnx_proto.TensorProto.FLOAT, [], [100.0]) + denominator_name, onnx_proto.TensorProto.FLOAT, [], [100.0] + ) - apply_div(scope, [output_name, denominator_name], - operator.output_full_names, container, broadcast=1) + apply_div( + scope, + [output_name, denominator_name], + operator.output_full_names, + container, + broadcast=1, + ) elif post_transform: container.add_node( post_transform, output_name, operator.output_full_names, - name=scope.get_unique_operator_name( - post_transform), + name=scope.get_unique_operator_name(post_transform), ) else: - container.add_node('Identity', output_name, - operator.output_full_names, - name=scope.get_unique_operator_name('Identity')) + container.add_node( + "Identity", + output_name, + operator.output_full_names, + name=scope.get_unique_operator_name("Identity"), + ) -def modify_tree_for_rule_in_set(gbm, use_float=False, verbose=0, count=0, # pylint: disable=R1710 - info=None): +def modify_tree_for_rule_in_set( + gbm, use_float=False, verbose=0, count=0, info=None # pylint: disable=R1710 +): """ LightGBM produces sometimes a tree with a node set to use rule ``==`` to a set of values (= in set), @@ -706,32 +855,39 @@ def modify_tree_for_rule_in_set(gbm, use_float=False, verbose=0, count=0, # pyl :param info: addition information to speed up this search :return: number of changed nodes (include *count*) """ - if 'tree_info' in gbm: + if "tree_info" in gbm: if info is not None: - dec_nodes = info['decision_nodes'] + dec_nodes = info["decision_nodes"] else: dec_nodes = None if verbose >= 2 and has_tqdm(): from tqdm import tqdm - loop = tqdm(gbm['tree_info']) + + loop = tqdm(gbm["tree_info"]) for i, tree in enumerate(loop): loop.set_description("rules tree %d c=%d" % (i, count)) count = modify_tree_for_rule_in_set( - tree, use_float=use_float, count=count, - info=None if dec_nodes is None else dec_nodes[i]) + tree, + use_float=use_float, + count=count, + info=None if dec_nodes is None else dec_nodes[i], + ) else: - for i, tree in enumerate(gbm['tree_info']): + for i, tree in enumerate(gbm["tree_info"]): count = modify_tree_for_rule_in_set( - tree, use_float=use_float, count=count, - info=None if dec_nodes is None else dec_nodes[i]) + tree, + use_float=use_float, + count=count, + info=None if dec_nodes is None else dec_nodes[i], + ) return count - if 'tree_structure' in gbm: + if "tree_structure" in gbm: return modify_tree_for_rule_in_set( - gbm['tree_structure'], use_float=use_float, count=count, - info=info) + gbm["tree_structure"], use_float=use_float, count=count, info=info + ) - if 'decision_type' not in gbm: + if "decision_type" not in gbm: return count def str2number(val): @@ -746,34 +902,34 @@ def str2number(val): if info is None: def recursive_call(this, c): - if 'left_child' in this: - c = process_node(this['left_child'], count=c) - if 'right_child' in this: - c = process_node(this['right_child'], count=c) + if "left_child" in this: + c = process_node(this["left_child"], count=c) + if "right_child" in this: + c = process_node(this["right_child"], count=c) return c def process_node(node, count): - if 'decision_type' not in node: + if "decision_type" not in node: return count - if node['decision_type'] != '==': + if node["decision_type"] != "==": return recursive_call(node, count) - th = node['threshold'] + th = node["threshold"] if not isinstance(th, str): return recursive_call(node, count) - pos = th.find('||') + pos = th.find("||") if pos == -1: return recursive_call(node, count) th1 = str2number(th[:pos]) def doit(): - rest = th[pos + 2:] - if '||' not in rest: + rest = th[pos + 2 :] + if "||" not in rest: rest = str2number(rest) - node['threshold'] = th1 + node["threshold"] = th1 new_node = node.copy() - node['right_child'] = new_node - new_node['threshold'] = rest + node["right_child"] = new_node + new_node["threshold"] = rest doit() return recursive_call(node, count + 1) @@ -785,34 +941,34 @@ def doit(): def split_node(node, th, pos): th1 = str2number(th[:pos]) - rest = th[pos + 2:] - if '||' not in rest: + rest = th[pos + 2 :] + if "||" not in rest: rest = str2number(rest) app = False else: app = True - node['threshold'] = th1 + node["threshold"] = th1 new_node = node.copy() - node['right_child'] = new_node - new_node['threshold'] = rest + node["right_child"] = new_node + new_node["threshold"] = rest return new_node, app stack = deque(info) while len(stack) > 0: node = stack.pop() - if 'decision_type' not in node: + if "decision_type" not in node: continue # leave - if node['decision_type'] != '==': + if node["decision_type"] != "==": continue - th = node['threshold'] + th = node["threshold"] if not isinstance(th, str): continue - pos = th.find('||') + pos = th.find("||") if pos == -1: continue @@ -825,38 +981,50 @@ def split_node(node, th, pos): def convert_lgbm_zipmap(scope, operator, container): - zipmap_attrs = {'name': scope.get_unique_operator_name('ZipMap')} - if hasattr(operator, 'classlabels_int64s'): - zipmap_attrs['classlabels_int64s'] = operator.classlabels_int64s + zipmap_attrs = {"name": scope.get_unique_operator_name("ZipMap")} + if hasattr(operator, "classlabels_int64s"): + zipmap_attrs["classlabels_int64s"] = operator.classlabels_int64s to_type = onnx_proto.TensorProto.INT64 - elif hasattr(operator, 'classlabels_strings'): - zipmap_attrs['classlabels_strings'] = operator.classlabels_strings + elif hasattr(operator, "classlabels_strings"): + zipmap_attrs["classlabels_strings"] = operator.classlabels_strings to_type = onnx_proto.TensorProto.STRING else: raise RuntimeError("Unknown class type.") if to_type == onnx_proto.TensorProto.STRING: - apply_identity(scope, operator.inputs[0].full_name, - operator.outputs[0].full_name, container) + apply_identity( + scope, + operator.inputs[0].full_name, + operator.outputs[0].full_name, + container, + ) else: - apply_cast(scope, operator.inputs[0].full_name, - operator.outputs[0].full_name, container, to=to_type) + apply_cast( + scope, + operator.inputs[0].full_name, + operator.outputs[0].full_name, + container, + to=to_type, + ) if operator.zipmap: - container.add_node('ZipMap', operator.inputs[1].full_name, - operator.outputs[1].full_name, - op_domain='ai.onnx.ml', **zipmap_attrs) + container.add_node( + "ZipMap", + operator.inputs[1].full_name, + operator.outputs[1].full_name, + op_domain="ai.onnx.ml", + **zipmap_attrs + ) else: # onnxconverter-common when trying to remove identity nodes # if node identity is used. - one = scope.get_unique_variable_name('one') + one = scope.get_unique_variable_name("one") - container.add_initializer( - one, onnx_proto.TensorProto.FLOAT, [], [1]) + container.add_initializer(one, onnx_proto.TensorProto.FLOAT, [], [1]) container.add_node( - 'Mul', [operator.inputs[1].full_name, one], - operator.outputs[1].full_name) + "Mul", [operator.inputs[1].full_name, one], operator.outputs[1].full_name + ) -register_converter('LgbmClassifier', convert_lightgbm) -register_converter('LgbmRegressor', convert_lightgbm) -register_converter('LgbmZipMap', convert_lgbm_zipmap) +register_converter("LgbmClassifier", convert_lightgbm) +register_converter("LgbmRegressor", convert_lightgbm) +register_converter("LgbmZipMap", convert_lgbm_zipmap) diff --git a/onnxmltools/convert/lightgbm/shape_calculators/Classifier.py b/onnxmltools/convert/lightgbm/shape_calculators/Classifier.py index ad9e9676..6b39b414 100644 --- a/onnxmltools/convert/lightgbm/shape_calculators/Classifier.py +++ b/onnxmltools/convert/lightgbm/shape_calculators/Classifier.py @@ -4,25 +4,32 @@ from ...common._registration import register_shape_calculator from ...common.utils import check_input_and_output_numbers, check_input_and_output_types from ...common.data_types import ( - FloatTensorType, Int64TensorType, + FloatTensorType, + Int64TensorType, StringTensorType, ) def calculate_lightgbm_classifier_output_shapes(operator): - ''' - This operator maps an input feature vector into a scalar label if the number of outputs is one. If two outputs - appear in this operator's output list, we should further generate a map storing all classes' probabilities. + """ + This operator maps an input feature vector into a + scalar label if the number of outputs is one. If two outputs + appear in this operator's output list, we should + further generate a map storing all classes' probabilities. Allowed input/output patterns are 1. [N, C] ---> [N, 1], A sequence of map Note that the second case is not allowed as long as ZipMap only produces dictionary. - ''' - check_input_and_output_numbers(operator, input_count_range=1, output_count_range=[1, 2]) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + """ + check_input_and_output_numbers( + operator, input_count_range=1, output_count_range=[1, 2] + ) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) if len(operator.inputs[0].type.shape) != 2: - raise RuntimeError('Input must be a [N, C]-tensor') + raise RuntimeError("Input must be a [N, C]-tensor") N = operator.inputs[0].type.shape[0] @@ -43,5 +50,5 @@ def calculate_lgbm_zipmap(operator): check_input_and_output_numbers(operator, output_count_range=2) -register_shape_calculator('LgbmClassifier', calculate_lightgbm_classifier_output_shapes) -register_shape_calculator('LgbmZipMap', calculate_lgbm_zipmap) +register_shape_calculator("LgbmClassifier", calculate_lightgbm_classifier_output_shapes) +register_shape_calculator("LgbmZipMap", calculate_lgbm_zipmap) diff --git a/onnxmltools/convert/lightgbm/shape_calculators/Regressor.py b/onnxmltools/convert/lightgbm/shape_calculators/Regressor.py index 9e7e900f..23274276 100644 --- a/onnxmltools/convert/lightgbm/shape_calculators/Regressor.py +++ b/onnxmltools/convert/lightgbm/shape_calculators/Regressor.py @@ -3,4 +3,4 @@ from ...common._registration import register_shape_calculator from ...common.shape_calculator import calculate_linear_regressor_output_shapes -register_shape_calculator('LgbmRegressor', calculate_linear_regressor_output_shapes) +register_shape_calculator("LgbmRegressor", calculate_linear_regressor_output_shapes) diff --git a/onnxmltools/convert/main.py b/onnxmltools/convert/main.py index 2e5cde12..4e677080 100644 --- a/onnxmltools/convert/main.py +++ b/onnxmltools/convert/main.py @@ -6,59 +6,103 @@ from .common import utils -def convert_coreml(model, name=None, initial_types=None, doc_string='', target_opset=None, - targeted_onnx=None, custom_conversion_functions=None, custom_shape_calculators=None): +def convert_coreml( + model, + name=None, + initial_types=None, + doc_string="", + target_opset=None, + targeted_onnx=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): if targeted_onnx is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.coreml_installed(): - raise RuntimeError('coremltools is not installed. Please install coremltools to use this feature.') + raise RuntimeError( + "coremltools is not installed. Please install coremltools to use this feature." + ) from .coreml.convert import convert - return convert(model, name, initial_types, doc_string, target_opset, targeted_onnx, - custom_conversion_functions, custom_shape_calculators) - - -def convert_keras(model, name=None, - initial_types=None, - doc_string='', - target_opset=None, - targeted_onnx=None, - channel_first_inputs=None, - custom_conversion_functions=None, - custom_shape_calculators=None, - default_batch_size=1): + + return convert( + model, + name, + initial_types, + doc_string, + target_opset, + targeted_onnx, + custom_conversion_functions, + custom_shape_calculators, + ) + + +def convert_keras( + model, + name=None, + initial_types=None, + doc_string="", + target_opset=None, + targeted_onnx=None, + channel_first_inputs=None, + custom_conversion_functions=None, + custom_shape_calculators=None, + default_batch_size=1, +): """ .. versionchanged:: 1.9.0 The conversion is now using *tf2onnx*. """ if targeted_onnx is not None: - warnings.warn("targeted_onnx is deprecated and unused. Use target_opset.", DeprecationWarning) + warnings.warn( + "targeted_onnx is deprecated and unused. Use target_opset.", + DeprecationWarning, + ) import tensorflow as tf - if pv.Version(tf.__version__) < pv.Version('2.0'): + + if pv.Version(tf.__version__) < pv.Version("2.0"): # Former converter for tensorflow<2.0. from keras2onnx import convert_keras as convert + return convert(model, name, doc_string, target_opset, channel_first_inputs) else: # For tensorflow>=2.0, new converter based on tf2onnx. import tf2onnx if not utils.tf2onnx_installed(): - raise RuntimeError('tf2onnx is not installed. Please install it to use this feature.') + raise RuntimeError( + "tf2onnx is not installed. Please install it to use this feature." + ) if custom_conversion_functions is not None: - warnings.warn('custom_conversion_functions is not supported any more. Please set it to None.') + warnings.warn( + "custom_conversion_functions is not supported any more. Please set it to None." + ) if custom_shape_calculators is not None: - warnings.warn('custom_shape_calculators is not supported any more. Please set it to None.') + warnings.warn( + "custom_shape_calculators is not supported any more. Please set it to None." + ) if default_batch_size != 1: - warnings.warn('default_batch_size is not supported any more. Please set it to 1.') + warnings.warn( + "default_batch_size is not supported any more. Please set it to 1." + ) if default_batch_size != 1: - warnings.warn('default_batch_size is not supported any more. Please set it to 1.') + warnings.warn( + "default_batch_size is not supported any more. Please set it to 1." + ) if initial_types is not None: from onnxconverter_common import ( - FloatTensorType, DoubleTensorType, - Int64TensorType, Int32TensorType, - StringTensorType, BooleanTensorType) + FloatTensorType, + DoubleTensorType, + Int64TensorType, + Int32TensorType, + StringTensorType, + BooleanTensorType, + ) + spec = [] for name, kind in initial_types: if isinstance(kind, FloatTensorType): @@ -75,7 +119,8 @@ def convert_keras(model, name=None, dtype = tf.bool else: raise TypeError( - "Unexpected type %r, cannot infer tensorflow type." % type(kind)) + "Unexpected type %r, cannot infer tensorflow type." % type(kind) + ) spec.append(tf.TensorSpec(tuple(kind.shape), dtype, name=name)) input_signature = tuple(spec) else: @@ -93,97 +138,207 @@ def convert_keras(model, name=None, shape_override=None, target=None, large_model=False, - output_path=None) + output_path=None, + ) if external_tensor_storage is not None: - warnings.warn("The current API does not expose the second result 'external_tensor_storage'. " - "Use tf2onnx directly to get it.") + warnings.warn( + "The current API does not expose the second result 'external_tensor_storage'. " + "Use tf2onnx directly to get it." + ) model_proto.doc_string = doc_string return model_proto -def convert_libsvm(model, name=None, initial_types=None, doc_string='', target_opset=None, - targeted_onnx=None, custom_conversion_functions=None, custom_shape_calculators=None): +def convert_libsvm( + model, + name=None, + initial_types=None, + doc_string="", + target_opset=None, + targeted_onnx=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): if targeted_onnx is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.libsvm_installed(): - raise RuntimeError('libsvm is not installed. Please install libsvm to use this feature.') + raise RuntimeError( + "libsvm is not installed. Please install libsvm to use this feature." + ) from .libsvm.convert import convert - return convert(model, name, initial_types, doc_string, target_opset, targeted_onnx, - custom_conversion_functions, custom_shape_calculators) - -def convert_catboost(model, name=None, initial_types=None, doc_string='', target_opset=None): + return convert( + model, + name, + initial_types, + doc_string, + target_opset, + targeted_onnx, + custom_conversion_functions, + custom_shape_calculators, + ) + + +def convert_catboost( + model, name=None, initial_types=None, doc_string="", target_opset=None +): try: from catboost.utils import convert_to_onnx_object except ImportError: - raise RuntimeError('CatBoost is not installed or needs to be updated. ' - 'Please install/upgrade CatBoost to use this feature.') - - return convert_to_onnx_object(model, export_parameters={'onnx_doc_string': doc_string, 'onnx_graph_name': name}, - initial_types=initial_types, target_opset=target_opset) - - -def convert_lightgbm(model, name=None, initial_types=None, doc_string='', target_opset=None, - targeted_onnx=None, custom_conversion_functions=None, - custom_shape_calculators=None, without_onnx_ml=False, zipmap=True, - split=None): + raise RuntimeError( + "CatBoost is not installed or needs to be updated. " + "Please install/upgrade CatBoost to use this feature." + ) + + return convert_to_onnx_object( + model, + export_parameters={"onnx_doc_string": doc_string, "onnx_graph_name": name}, + initial_types=initial_types, + target_opset=target_opset, + ) + + +def convert_lightgbm( + model, + name=None, + initial_types=None, + doc_string="", + target_opset=None, + targeted_onnx=None, + custom_conversion_functions=None, + custom_shape_calculators=None, + without_onnx_ml=False, + zipmap=True, + split=None, +): if targeted_onnx is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.lightgbm_installed(): - raise RuntimeError('lightgbm is not installed. Please install lightgbm to use this feature.') + raise RuntimeError( + "lightgbm is not installed. Please install lightgbm to use this feature." + ) from .lightgbm.convert import convert - return convert(model, name, initial_types, doc_string, target_opset, targeted_onnx, - custom_conversion_functions, custom_shape_calculators, without_onnx_ml, - zipmap=zipmap, split=split) - -def convert_sklearn(model, name=None, initial_types=None, doc_string='', target_opset=None, - targeted_onnx=None, custom_conversion_functions=None, custom_shape_calculators=None): + return convert( + model, + name, + initial_types, + doc_string, + target_opset, + targeted_onnx, + custom_conversion_functions, + custom_shape_calculators, + without_onnx_ml, + zipmap=zipmap, + split=split, + ) + + +def convert_sklearn( + model, + name=None, + initial_types=None, + doc_string="", + target_opset=None, + targeted_onnx=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): if targeted_onnx is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.sklearn_installed(): - raise RuntimeError('scikit-learn is not installed. Please install scikit-learn to use this feature.') + raise RuntimeError( + "scikit-learn is not installed. Please install scikit-learn to use this feature." + ) if not utils.skl2onnx_installed(): - raise RuntimeError('skl2onnx is not installed. Please install skl2onnx to use this feature.') + raise RuntimeError( + "skl2onnx is not installed. Please install skl2onnx to use this feature." + ) from skl2onnx.convert import convert_sklearn as convert_skl2onnx - return convert_skl2onnx(model, name, initial_types, doc_string, target_opset, - custom_conversion_functions, custom_shape_calculators) - -def convert_sparkml(model, name=None, initial_types=None, doc_string='', target_opset=None, - targeted_onnx=None, custom_conversion_functions=None, - custom_shape_calculators=None, spark_session=None): + return convert_skl2onnx( + model, + name, + initial_types, + doc_string, + target_opset, + custom_conversion_functions, + custom_shape_calculators, + ) + + +def convert_sparkml( + model, + name=None, + initial_types=None, + doc_string="", + target_opset=None, + targeted_onnx=None, + custom_conversion_functions=None, + custom_shape_calculators=None, + spark_session=None, +): if targeted_onnx is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.sparkml_installed(): - raise RuntimeError('Spark is not installed. Please install Spark to use this feature.') + raise RuntimeError( + "Spark is not installed. Please install Spark to use this feature." + ) from .sparkml.convert import convert - return convert(model, name, initial_types, doc_string, target_opset, targeted_onnx, - custom_conversion_functions, custom_shape_calculators, spark_session) + + return convert( + model, + name, + initial_types, + doc_string, + target_opset, + targeted_onnx, + custom_conversion_functions, + custom_shape_calculators, + spark_session, + ) def convert_xgboost(*args, **kwargs): - if kwargs.get('targeted_onnx', None) is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + if kwargs.get("targeted_onnx", None) is not None: + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.xgboost_installed(): - raise RuntimeError('xgboost is not installed. Please install xgboost to use this feature.') + raise RuntimeError( + "xgboost is not installed. Please install xgboost to use this feature." + ) from .xgboost.convert import convert + return convert(*args, **kwargs) def convert_h2o(*args, **kwargs): - if kwargs.get('targeted_onnx', None) is not None: - warnings.warn("targeted_onnx is deprecated. Use target_opset.", DeprecationWarning) + if kwargs.get("targeted_onnx", None) is not None: + warnings.warn( + "targeted_onnx is deprecated. Use target_opset.", DeprecationWarning + ) if not utils.h2o_installed(): - raise RuntimeError('h2o is not installed. Please install h2o to use this feature.') + raise RuntimeError( + "h2o is not installed. Please install h2o to use this feature." + ) from .h2o.convert import convert + return convert(*args, **kwargs) @@ -194,7 +349,7 @@ def _collect_input_nodes(graph, outputs): while node_inputs: nd_ = node_inputs[0] del node_inputs[0] - if nd_.type in ['Placeholder', "PlaceholderV2", 'PlaceholderWithDefault']: + if nd_.type in ["Placeholder", "PlaceholderV2", "PlaceholderWithDefault"]: input_nodes.add(nd_) if nd_ in nodes_to_keep: continue @@ -205,13 +360,18 @@ def _collect_input_nodes(graph, outputs): return input_nodes, nodes_to_keep -def _convert_tf_wrapper(frozen_graph_def, - name=None, input_names=None, output_names=None, - doc_string='', - target_opset=None, - channel_first_inputs=None, - debug_mode=False, custom_op_conversions=None, - **kwargs): +def _convert_tf_wrapper( + frozen_graph_def, + name=None, + input_names=None, + output_names=None, + doc_string="", + target_opset=None, + channel_first_inputs=None, + debug_mode=False, + custom_op_conversions=None, + **kwargs +): """ convert a tensorflow graph def into a ONNX model proto, just like how keras does. :param graph_def: the frozen tensorflow graph @@ -235,21 +395,25 @@ def _convert_tf_wrapper(frozen_graph_def, if not doc_string: doc_string = "converted from {}".format(name) - tf_graph_def = tf2onnx.tfonnx.tf_optimize(input_names, output_names, frozen_graph_def, True) + tf_graph_def = tf2onnx.tfonnx.tf_optimize( + input_names, output_names, frozen_graph_def, True + ) with tf.Graph().as_default() as tf_graph: - tf.import_graph_def(tf_graph_def, name='') + tf.import_graph_def(tf_graph_def, name="") if not input_names: input_nodes = list(_collect_input_nodes(tf_graph, output_names)[0]) input_names = [nd_.outputs[0].name for nd_ in input_nodes] - g = tf2onnx.tfonnx.process_tf_graph(tf_graph, - continue_on_error=debug_mode, - opset=target_opset, - custom_op_handlers=custom_op_conversions, - inputs_as_nchw=channel_first_inputs, - output_names=output_names, - input_names=input_names, - **kwargs) + g = tf2onnx.tfonnx.process_tf_graph( + tf_graph, + continue_on_error=debug_mode, + opset=target_opset, + custom_op_handlers=custom_op_conversions, + inputs_as_nchw=channel_first_inputs, + output_names=output_names, + input_names=input_names, + **kwargs + ) onnx_graph = tf2onnx.optimizer.optimize_graph(g) model_proto = onnx_graph.make_model(doc_string) @@ -257,17 +421,34 @@ def _convert_tf_wrapper(frozen_graph_def, return model_proto -def convert_tensorflow(frozen_graph_def, - name=None, input_names=None, output_names=None, - doc_string='', - target_opset=None, - channel_first_inputs=None, - debug_mode=False, custom_op_conversions=None, - **kwargs): +def convert_tensorflow( + frozen_graph_def, + name=None, + input_names=None, + output_names=None, + doc_string="", + target_opset=None, + channel_first_inputs=None, + debug_mode=False, + custom_op_conversions=None, + **kwargs +): import pkgutil - if not pkgutil.find_loader('tf2onnx'): - raise RuntimeError('tf2onnx is not installed, please install it before calling this function.') - return _convert_tf_wrapper(frozen_graph_def, name, input_names, output_names, doc_string, - target_opset, channel_first_inputs, debug_mode, custom_op_conversions, - **kwargs) + if not pkgutil.find_loader("tf2onnx"): + raise RuntimeError( + "tf2onnx is not installed, please install it before calling this function." + ) + + return _convert_tf_wrapper( + frozen_graph_def, + name, + input_names, + output_names, + doc_string, + target_opset, + channel_first_inputs, + debug_mode, + custom_op_conversions, + **kwargs + ) diff --git a/onnxmltools/convert/sparkml/_parse.py b/onnxmltools/convert/sparkml/_parse.py index 22685a5e..326fd6c3 100644 --- a/onnxmltools/convert/sparkml/_parse.py +++ b/onnxmltools/convert/sparkml/_parse.py @@ -11,19 +11,21 @@ def _get_variable_for_input(scope, input_name, global_inputs, output_dict): - ''' + """ Find the corresponding Variable for a given raw operator (model) name - The variable is either supplied as graph/global inputs or has been generated as output by previous ops + The variable is either supplied as graph/global + nputs or has been generated as output by previous ops + :param input_name: :param global_inputs: :param output_dict: :return: - ''' + """ if input_name in output_dict: value = output_dict[input_name] ref_count = value[0] variable = value[1] - output_dict[input_name] = [ref_count+1, variable] + output_dict[input_name] = [ref_count + 1, variable] return variable matches = [x for x in global_inputs if x.raw_name == input_name] @@ -36,7 +38,7 @@ def _get_variable_for_input(scope, input_name, global_inputs, output_dict): def _parse_sparkml_simple_model(spark, scope, model, global_inputs, output_dict): - ''' + """ This function handles all non-pipeline models. :param scope: Scope object @@ -44,11 +46,16 @@ def _parse_sparkml_simple_model(spark, scope, model, global_inputs, output_dict) :param global_inputs: A list of variables :param output_dict: An accumulated list of output_original_name->(ref_count, variable) :return: A list of output variables which will be passed to next stage - ''' - this_operator = scope.declare_local_operator(get_sparkml_operator_name(type(model)), model) - this_operator.raw_params = {'SparkSession': spark} + """ + this_operator = scope.declare_local_operator( + get_sparkml_operator_name(type(model)), model + ) + this_operator.raw_params = {"SparkSession": spark} raw_input_names = get_input_names(model) - this_operator.inputs = [_get_variable_for_input(scope, x, global_inputs, output_dict) for x in raw_input_names] + this_operator.inputs = [ + _get_variable_for_input(scope, x, global_inputs, output_dict) + for x in raw_input_names + ] raw_output_names = get_output_names(model) for output_name in raw_output_names: variable = scope.declare_local_variable(output_name, FloatTensorType()) @@ -57,7 +64,7 @@ def _parse_sparkml_simple_model(spark, scope, model, global_inputs, output_dict) def _parse_sparkml_pipeline(spark, scope, model, global_inputs, output_dict): - ''' + """ The basic ideas of spark-ml parsing: 1. Sequentially go though all stages defined in the considered spark-ml pipeline 2. The output variables of one stage will be fed into its next stage as the inputs. @@ -67,49 +74,67 @@ def _parse_sparkml_pipeline(spark, scope, model, global_inputs, output_dict): :param global_inputs: A list of Variable objects :param output_dict: An accumulated list of output_original_name->(ref_count, variable) :return: A list of output variables produced by the input pipeline - ''' + """ for stage in model.stages: _parse_sparkml(spark, scope, stage, global_inputs, output_dict) def _parse_sparkml(spark, scope, model, global_inputs, output_dict): - ''' - This is a delegate function. It doesn't nothing but invoke the correct parsing function according to the input + """ + This is a delegate function. It doesn't nothing but + invoke the correct parsing function according to the input model's type. + :param scope: Scope object :param model: A spark-ml object (e.g., OneHotEncoder and LogisticRegression) :param inputs: A list of variables :return: The output variables produced by the input model - ''' + """ if isinstance(model, PipelineModel): return _parse_sparkml_pipeline(spark, scope, model, global_inputs, output_dict) else: - return _parse_sparkml_simple_model(spark, scope, model, global_inputs, output_dict) - - -def parse_sparkml(spark, model, initial_types=None, target_opset=None, - custom_conversion_functions=None, custom_shape_calculators=None): - # Put spark-ml object into an abstract container so that our framework can work seamlessly on models created + return _parse_sparkml_simple_model( + spark, scope, model, global_inputs, output_dict + ) + + +def parse_sparkml( + spark, + model, + initial_types=None, + target_opset=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): + # Put spark-ml object into an abstract container + # so that our framework can work seamlessly on models created # with different machine learning tools. raw_model_container = SparkmlModelContainer(model) - # Declare a computational graph. It will become a representation of the input spark-ml model after parsing. - topology = Topology(raw_model_container, default_batch_size='None', - initial_types=initial_types, - target_opset=target_opset, - custom_conversion_functions=custom_conversion_functions, - custom_shape_calculators=custom_shape_calculators) - - # Declare an object to provide variables' and operators' naming mechanism. In contrast to CoreML, one global scope + # Declare a computational graph. It will become + # a representation of the input spark-ml model after parsing. + topology = Topology( + raw_model_container, + default_batch_size="None", + initial_types=initial_types, + target_opset=target_opset, + custom_conversion_functions=custom_conversion_functions, + custom_shape_calculators=custom_shape_calculators, + ) + + # Declare an object to provide variables' and operators' + # naming mechanism. In contrast to CoreML, one global scope # is enough for parsing spark-ml models. - scope = topology.declare_scope('__root__') + scope = topology.declare_scope("__root__") - # Declare input variables. They should be the inputs of the spark-ml model you want to convert into ONNX + # Declare input variables. They should be the inputs of the + # spark-ml model you want to convert into ONNX inputs = [] for var_name, initial_type in initial_types: inputs.append(scope.declare_local_variable(var_name, initial_type)) - # The object raw_model_container is a part of the topology we're going to return. We use it to store the inputs of + # The object raw_model_container is a part of the topology + # we're going to return. We use it to store the inputs of # the spark-ml's computational graph. for variable in inputs: raw_model_container.add_input(variable) @@ -119,10 +144,11 @@ def parse_sparkml(spark, model, initial_types=None, target_opset=None, _parse_sparkml(spark, scope, model, inputs, output_dict) outputs = [] for k, v in output_dict.items(): - if v[0] == 0: # ref count is zero + if v[0] == 0: # ref count is zero outputs.append(v[1]) - # THe object raw_model_container is a part of the topology we're going to return. We use it to store the outputs of + # THe object raw_model_container is a part of the topology + # we're going to return. We use it to store the outputs of # the spark-ml's computational graph. for variable in outputs: raw_model_container.add_output(variable) diff --git a/onnxmltools/convert/sparkml/convert.py b/onnxmltools/convert/sparkml/convert.py index bf8a3521..abd350c0 100644 --- a/onnxmltools/convert/sparkml/convert.py +++ b/onnxmltools/convert/sparkml/convert.py @@ -5,14 +5,22 @@ from ..common.onnx_ex import get_maximum_opset_supported from ..common._topology import convert_topology from ._parse import parse_sparkml -from . import operator_converters -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, - spark_session=None): - ''' - This function produces an equivalent ONNX model of the given spark-ml model. The supported spark-ml +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, + spark_session=None, +): + """ + This function produces an equivalent ONNX model + of the given spark-ml model. The supported spark-ml modules are listed below. * Preprocessings and transformations: @@ -30,44 +38,67 @@ def convert(model, name=None, initial_types=None, doc_string='', target_opset=No * pipeline 29. pipeline.Pipeline - For pipeline conversion, user needs to make sure each component is one of our supported items (1)-(24). + For pipeline conversion, user needs to make sure + each component is one of our supported items (1)-(24). - This function converts the specified spark-ml model into its ONNX counterpart. Notice that for all conversions, + This function converts the specified spark-ml model into its + ONNX counterpart. Notice that for all conversions, initial types are required. ONNX model name can also be specified. :param model: A spark-ml model - :param initial_types: a python list. Each element is a tuple of a variable name and a type defined in data_types.py - :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 name: The name of the graph (type: GraphProto) + in the produced ONNX model (type: ModelProto) :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 + :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 spark-ml model Example of initial_types: - Assume that the specified spark-ml model takes a heterogeneous list as its input. If the first 5 elements are - floats and the last 10 elements are integers, we need to specify initial types as below. The [1] in [1, 5] indicates + + Assume that the specified spark-ml model takes a heterogeneous + list as its input. If the first 5 elements are + floats and the last 10 elements are integers, we need to + specify initial types as below. The [1] in [1, 5] indicates the batch size here is 1. >>> from onnxmltools.convert.common.data_types import FloatTensorType, Int64TensorType - >>> initial_type = [('float_input', FloatTensorType([1, 5])), ('int64_input', Int64TensorType([1, 10]))] - ''' + >>> initial_type = [('float_input', FloatTensorType([1, 5])), + ('int64_input', Int64TensorType([1, 10]))] + """ if initial_types is None: - raise ValueError('Initial types are required. See usage of convert(...) in \ - onnxmltools.convert.sparkml.convert for details') + raise ValueError( + "Initial types are required. See usage of convert(...) in \ + onnxmltools.convert.sparkml.convert for details" + ) if name is None: name = str(uuid4().hex) target_opset = target_opset if target_opset else get_maximum_opset_supported() # Parse spark-ml model as our internal data structure (i.e., Topology) - topology = parse_sparkml(spark_session, model, initial_types, target_opset, custom_conversion_functions, custom_shape_calculators) + topology = parse_sparkml( + spark_session, + model, + initial_types, + target_opset, + custom_conversion_functions, + custom_shape_calculators, + ) # Infer variable shapes topology.compile() # Convert our Topology object into ONNX. The outcome is an ONNX model. - onnx_model = convert_topology(topology, name, doc_string, target_opset, targeted_onnx) + onnx_model = convert_topology( + topology, name, doc_string, target_opset, targeted_onnx + ) return onnx_model diff --git a/onnxmltools/convert/sparkml/operator_converters/aft_survival_regression.py b/onnxmltools/convert/sparkml/operator_converters/aft_survival_regression.py index c1e4512d..284d5472 100644 --- a/onnxmltools/convert/sparkml/operator_converters/aft_survival_regression.py +++ b/onnxmltools/convert/sparkml/operator_converters/aft_survival_regression.py @@ -13,23 +13,38 @@ def convert_aft_survival_regression(scope, operator, container): op = operator.raw_operator coefficients = op.coefficients.toArray().astype(float) - coefficients_tensor = scope.get_unique_variable_name('coefficients_tensor') - container.add_initializer(coefficients_tensor, onnx_proto.TensorProto.FLOAT, [1, len(coefficients)], coefficients) + coefficients_tensor = scope.get_unique_variable_name("coefficients_tensor") + container.add_initializer( + coefficients_tensor, + onnx_proto.TensorProto.FLOAT, + [1, len(coefficients)], + coefficients, + ) intercepts = ( op.intercept.astype(float) if isinstance(op.intercept, collections.abc.Iterable) - else [float(op.intercept)]) - intercepts_tensor = scope.get_unique_variable_name('intercepts_tensor') - container.add_initializer(intercepts_tensor, onnx_proto.TensorProto.FLOAT, [len(intercepts)], intercepts) - - matmul_result = scope.get_unique_variable_name('matmul_result_tensor') - apply_matmul(scope, [operator.input_full_names[0], coefficients_tensor], matmul_result, container) - add_result = scope.get_unique_variable_name('intercept_added_tensor') + else [float(op.intercept)] + ) + intercepts_tensor = scope.get_unique_variable_name("intercepts_tensor") + container.add_initializer( + intercepts_tensor, onnx_proto.TensorProto.FLOAT, [len(intercepts)], intercepts + ) + + matmul_result = scope.get_unique_variable_name("matmul_result_tensor") + apply_matmul( + scope, + [operator.input_full_names[0], coefficients_tensor], + matmul_result, + container, + ) + add_result = scope.get_unique_variable_name("intercept_added_tensor") apply_add(scope, [matmul_result, intercepts_tensor], add_result, container) apply_exp(scope, add_result, operator.output_full_names, container) -register_converter('pyspark.ml.regression.AFTSurvivalRegressionModel', convert_aft_survival_regression) +register_converter( + "pyspark.ml.regression.AFTSurvivalRegressionModel", convert_aft_survival_regression +) def calculate_aft_survival_regression_output_shapes(operator): @@ -39,5 +54,7 @@ def calculate_aft_survival_regression_output_shapes(operator): operator.outputs[0].type = FloatTensorType([N, 1]) -register_shape_calculator('pyspark.ml.regression.AFTSurvivalRegressionModel', - calculate_aft_survival_regression_output_shapes) +register_shape_calculator( + "pyspark.ml.regression.AFTSurvivalRegressionModel", + calculate_aft_survival_regression_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/binarizer.py b/onnxmltools/convert/sparkml/operator_converters/binarizer.py index 83133c75..66533d3b 100644 --- a/onnxmltools/convert/sparkml/operator_converters/binarizer.py +++ b/onnxmltools/convert/sparkml/operator_converters/binarizer.py @@ -10,21 +10,27 @@ def convert_sparkml_binarizer(scope, operator, container): op = operator.raw_operator input_name = op.getInputCol() - op_type = 'Binarizer' + op_type = "Binarizer" name = scope.get_unique_operator_name(op_type) - attrs = {'name': name, 'threshold': float(op.getThreshold())} - container.add_node(op_type, input_name, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + attrs = {"name": name, "threshold": float(op.getThreshold())} + container.add_node( + op_type, input_name, operator.output_full_names, op_domain="ai.onnx.ml", **attrs + ) -register_converter('pyspark.ml.feature.Binarizer', convert_sparkml_binarizer) +register_converter("pyspark.ml.feature.Binarizer", convert_sparkml_binarizer) def calculate_sparkml_binarizer_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) input_type = copy.deepcopy(operator.inputs[0].type) operator.outputs[0].type = input_type -register_shape_calculator('pyspark.ml.feature.Binarizer', calculate_sparkml_binarizer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.Binarizer", calculate_sparkml_binarizer_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/bucketed_random_projection_lsh.py b/onnxmltools/convert/sparkml/operator_converters/bucketed_random_projection_lsh.py index 4ff3bea7..fa0c8057 100644 --- a/onnxmltools/convert/sparkml/operator_converters/bucketed_random_projection_lsh.py +++ b/onnxmltools/convert/sparkml/operator_converters/bucketed_random_projection_lsh.py @@ -14,30 +14,48 @@ def get_rand_vectors(operator): global g_rand_vectors if not g_rand_vectors: - g_rand_vectors = save_read_sparkml_model_data( - operator.raw_params['SparkSession'], operator.raw_operator - ).first()[0].toArray().transpose().astype(numpy.float32) + g_rand_vectors = ( + save_read_sparkml_model_data( + operator.raw_params["SparkSession"], operator.raw_operator + ) + .first()[0] + .toArray() + .transpose() + .astype(numpy.float32) + ) return g_rand_vectors def convert_min_hash_lsh(scope, operator, container): rand_vectors = get_rand_vectors(operator) - bucket_length = float(operator.raw_operator.getOrDefault('bucketLength')) - - rand_vectors_tensor = scope.get_unique_variable_name('rand_vectors_tensor') - container.add_initializer(rand_vectors_tensor, onnx_proto.TensorProto.FLOAT, - rand_vectors.shape, rand_vectors.flatten()) - matmul_result = scope.get_unique_variable_name('matmul_result_tensor') - apply_matmul(scope, [operator.input_full_names[0], rand_vectors_tensor], matmul_result, container) - bucket_length_tensor = scope.get_unique_variable_name('bucket_length_tensor') - container.add_initializer(bucket_length_tensor, onnx_proto.TensorProto.FLOAT, - [1], [bucket_length]) - div_result = scope.get_unique_variable_name('div_result_tensor') + bucket_length = float(operator.raw_operator.getOrDefault("bucketLength")) + + rand_vectors_tensor = scope.get_unique_variable_name("rand_vectors_tensor") + container.add_initializer( + rand_vectors_tensor, + onnx_proto.TensorProto.FLOAT, + rand_vectors.shape, + rand_vectors.flatten(), + ) + matmul_result = scope.get_unique_variable_name("matmul_result_tensor") + apply_matmul( + scope, + [operator.input_full_names[0], rand_vectors_tensor], + matmul_result, + container, + ) + bucket_length_tensor = scope.get_unique_variable_name("bucket_length_tensor") + container.add_initializer( + bucket_length_tensor, onnx_proto.TensorProto.FLOAT, [1], [bucket_length] + ) + div_result = scope.get_unique_variable_name("div_result_tensor") apply_div(scope, [matmul_result, bucket_length_tensor], div_result, container) apply_floor(scope, div_result, operator.output_full_names[0], container) -register_converter('pyspark.ml.feature.BucketedRandomProjectionLSHModel', convert_min_hash_lsh) +register_converter( + "pyspark.ml.feature.BucketedRandomProjectionLSHModel", convert_min_hash_lsh +) def calculate_min_hash_lsh_output_shapes(operator): @@ -49,4 +67,7 @@ def calculate_min_hash_lsh_output_shapes(operator): operator.outputs[0].type = FloatTensorType([N, C]) -register_shape_calculator('pyspark.ml.feature.BucketedRandomProjectionLSHModel', calculate_min_hash_lsh_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.BucketedRandomProjectionLSHModel", + calculate_min_hash_lsh_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/bucketizer.py b/onnxmltools/convert/sparkml/operator_converters/bucketizer.py index 37f1e385..09a925bc 100644 --- a/onnxmltools/convert/sparkml/operator_converters/bucketizer.py +++ b/onnxmltools/convert/sparkml/operator_converters/bucketizer.py @@ -10,58 +10,90 @@ def convert_bucketizer(scope, operator, container): op = operator.raw_operator splits = op.getSplits() - if splits[0] != -float("inf") or splits[-1] != float('inf'): + if splits[0] != -float("inf") or splits[-1] != float("inf"): raise RuntimeError("the Splits must include positive/negative infinity") input_shape = operator.inputs[0].type.shape - reshape_data = scope.get_unique_variable_name('reshape_info_tensor') - container.add_initializer(reshape_data, onnx_proto.TensorProto.INT64, [3], - [-1, input_shape[1], 1]) + reshape_data = scope.get_unique_variable_name("reshape_info_tensor") + container.add_initializer( + reshape_data, onnx_proto.TensorProto.INT64, [3], [-1, input_shape[1], 1] + ) outputs = [] for split in splits[1:]: - less_output = scope.get_unique_variable_name('less_output_tensor') - initializer_name = 'initializer_' + str(split) - container.add_initializer(initializer_name, onnx_proto.TensorProto.FLOAT, [1], - [split]) - container.add_node('Less', [operator.inputs[0].full_name, initializer_name], less_output, - name=scope.get_unique_operator_name('Less'), - op_version=7) - casted_output = scope.get_unique_variable_name('cast_output_tensor') - container.add_node('Cast', less_output, casted_output, - name=scope.get_unique_operator_name('Cast'), - op_version=6, - to=1) - redim_output = scope.get_unique_variable_name('reshape_output_tensor') - container.add_node('Reshape', [casted_output, reshape_data], redim_output, - name=scope.get_unique_operator_name('Reshape'), - op_version=5) + less_output = scope.get_unique_variable_name("less_output_tensor") + initializer_name = "initializer_" + str(split) + container.add_initializer( + initializer_name, onnx_proto.TensorProto.FLOAT, [1], [split] + ) + container.add_node( + "Less", + [operator.inputs[0].full_name, initializer_name], + less_output, + name=scope.get_unique_operator_name("Less"), + op_version=7, + ) + casted_output = scope.get_unique_variable_name("cast_output_tensor") + container.add_node( + "Cast", + less_output, + casted_output, + name=scope.get_unique_operator_name("Cast"), + op_version=6, + to=1, + ) + redim_output = scope.get_unique_variable_name("reshape_output_tensor") + container.add_node( + "Reshape", + [casted_output, reshape_data], + redim_output, + name=scope.get_unique_operator_name("Reshape"), + op_version=5, + ) outputs.append(redim_output) - concat_output = scope.get_unique_variable_name('concat_output_tensor') - container.add_node('Concat', outputs, concat_output, - name=scope.get_unique_operator_name('Concat'), - op_version=1, - axis=2) - argmax_output = scope.get_unique_variable_name('argmax_output_tensor') - container.add_node('ArgMax', concat_output, argmax_output, - name=scope.get_unique_operator_name('ArgMax'), - op_version=1, - axis=2, - keepdims=0) - container.add_node('Cast', argmax_output, operator.output_full_names, - name=scope.get_unique_operator_name('Cast'), - op_version=6, - to=1) + concat_output = scope.get_unique_variable_name("concat_output_tensor") + container.add_node( + "Concat", + outputs, + concat_output, + name=scope.get_unique_operator_name("Concat"), + op_version=1, + axis=2, + ) + argmax_output = scope.get_unique_variable_name("argmax_output_tensor") + container.add_node( + "ArgMax", + concat_output, + argmax_output, + name=scope.get_unique_operator_name("ArgMax"), + op_version=1, + axis=2, + keepdims=0, + ) + container.add_node( + "Cast", + argmax_output, + operator.output_full_names, + name=scope.get_unique_operator_name("Cast"), + op_version=6, + to=1, + ) + + +register_converter("pyspark.ml.feature.Bucketizer", convert_bucketizer) -register_converter('pyspark.ml.feature.Bucketizer', convert_bucketizer) def calculate_bucketizer_output_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, - good_input_types=[FloatTensorType], - good_output_types=[FloatTensorType]) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType], + good_output_types=[FloatTensorType], + ) input_type = copy.deepcopy(operator.inputs[0].type) for output in operator.outputs: output.type = input_type -register_shape_calculator('pyspark.ml.feature.Bucketizer', calculate_bucketizer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.Bucketizer", calculate_bucketizer_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/chi_sq_selector.py b/onnxmltools/convert/sparkml/operator_converters/chi_sq_selector.py index b13fab71..240d4763 100644 --- a/onnxmltools/convert/sparkml/operator_converters/chi_sq_selector.py +++ b/onnxmltools/convert/sparkml/operator_converters/chi_sq_selector.py @@ -1,31 +1,40 @@ # SPDX-License-Identifier: Apache-2.0 import copy +import onnx from ...common._registration import register_converter, register_shape_calculator from ...common.utils import check_input_and_output_numbers, check_input_and_output_types -from ...common.data_types import * +from ...common.data_types import FloatTensorType, Int64TensorType, StringTensorType def convert_chi_sq_selector(scope, operator, container): op = operator.raw_operator indices = op.selectedFeatures - indices_tensor = 'indices_tensor' - container.add_initializer(indices_tensor, onnx_proto.TensorProto.INT64, [len(indices)], indices) - container.add_node('ArrayFeatureExtractor', - [operator.input_full_names[0], indices_tensor], operator.output_full_names, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('ArrayFeatureExtractor')) + indices_tensor = "indices_tensor" + container.add_initializer( + indices_tensor, onnx.TensorProto.INT64, [len(indices)], indices + ) + container.add_node( + "ArrayFeatureExtractor", + [operator.input_full_names[0], indices_tensor], + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("ArrayFeatureExtractor"), + ) -register_converter('pyspark.ml.feature.ChiSqSelectorModel', convert_chi_sq_selector) +register_converter("pyspark.ml.feature.ChiSqSelectorModel", convert_chi_sq_selector) def calculate_chi_sq_selector_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, - good_input_types=[FloatTensorType, Int64TensorType, StringTensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType, StringTensorType] + ) operator.outputs[0].type = copy.deepcopy(operator.inputs[0].type) operator.outputs[0].type.shape[1] = len(operator.raw_operator.selectedFeatures) -register_shape_calculator('pyspark.ml.feature.ChiSqSelectorModel', calculate_chi_sq_selector_shapes) +register_shape_calculator( + "pyspark.ml.feature.ChiSqSelectorModel", calculate_chi_sq_selector_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/common.py b/onnxmltools/convert/sparkml/operator_converters/common.py index dc8cb782..ff9a38f8 100644 --- a/onnxmltools/convert/sparkml/operator_converters/common.py +++ b/onnxmltools/convert/sparkml/operator_converters/common.py @@ -1,27 +1,48 @@ # SPDX-License-Identifier: Apache-2.0 -from ...common.data_types import Int64TensorType, Int64Type, FloatTensorType, FloatType, StringType +from ...common.data_types import ( + Int64TensorType, + Int64Type, + FloatTensorType, + FloatType, + StringType, +) def convert_integer_to_float(scope, variable, container): - op_type = 'Scaler' - scaled_name = scope.get_unique_variable_name(variable.full_name + '_scaled') - scaler_attrs = {'name': scope.get_unique_operator_name(op_type), 'scale': [1.], 'offset': [0.]} - container.add_node('Scaler', variable.full_name, scaled_name, op_domain='ai.onnx.ml', **scaler_attrs) + op_type = "Scaler" + scaled_name = scope.get_unique_variable_name(variable.full_name + "_scaled") + scaler_attrs = { + "name": scope.get_unique_operator_name(op_type), + "scale": [1.0], + "offset": [0.0], + } + container.add_node( + "Scaler", + variable.full_name, + scaled_name, + op_domain="ai.onnx.ml", + **scaler_attrs + ) return scaled_name def concatenate_variables(scope, variables, container): - ''' - This function allocate operators to from a float tensor by concatenating all input variables. Notice that if all + """ + This function allocate operators to from a float + tensor by concatenating all input variables. Notice that if all integer inputs would be converted to floats before concatenation. - ''' + """ # Check if it's possible to concatenate those inputs. type_set = set(type(variable.type) for variable in variables) number_type_set = {FloatType, FloatTensorType, Int64Type, Int64TensorType} - if StringType in type_set and any(number_type in type_set for number_type in number_type_set): - raise RuntimeError('We are not able to concatenate numerical tensor(s) and string tensor(s)') + if StringType in type_set and any( + number_type in type_set for number_type in number_type_set + ): + raise RuntimeError( + "We are not able to concatenate numerical tensor(s) and string tensor(s)" + ) input_names = [] # input variables' names we want to concatenate input_dims = [] # dimensions of the variables that is going to be concatenated @@ -40,11 +61,16 @@ def concatenate_variables(scope, variables, container): return input_names[0] else: # To combine all inputs, we need a FeatureVectorizer - op_type = 'FeatureVectorizer' - attrs = {'name': scope.get_unique_operator_name(op_type), 'inputdimensions': input_dims} + op_type = "FeatureVectorizer" + attrs = { + "name": scope.get_unique_operator_name(op_type), + "inputdimensions": input_dims, + } # Create a variable name to capture feature vectorizer's output - concatenated_name = scope.get_unique_variable_name('concatenated') + concatenated_name = scope.get_unique_variable_name("concatenated") # Set up our FeatureVectorizer - container.add_node(op_type, input_names, concatenated_name, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, input_names, concatenated_name, op_domain="ai.onnx.ml", **attrs + ) return concatenated_name diff --git a/onnxmltools/convert/sparkml/operator_converters/count_vectorizer.py b/onnxmltools/convert/sparkml/operator_converters/count_vectorizer.py index 502ee556..516eadd2 100644 --- a/onnxmltools/convert/sparkml/operator_converters/count_vectorizer.py +++ b/onnxmltools/convert/sparkml/operator_converters/count_vectorizer.py @@ -7,25 +7,37 @@ from pyspark.ml.feature import CountVectorizerModel -def convert_count_vectorizer(scope: Scope, operator: Operator, container: ModelComponentContainer): +def convert_count_vectorizer( + scope: Scope, operator: Operator, container: ModelComponentContainer +): op: CountVectorizerModel = operator.raw_operator - vocab, minTF, binary = op.vocabulary, op.getOrDefault("minTF"), op.getOrDefault("binary") + vocab, minTF, binary = ( + op.vocabulary, + op.getOrDefault("minTF"), + op.getOrDefault("binary"), + ) if minTF < 1.0: - raise NotImplementedError("Converting to ONNX for CountVectorizerModel is not supported when minTF < 1.0") + raise NotImplementedError( + "Converting to ONNX for CountVectorizerModel is not supported when minTF < 1.0" + ) min_opset = 9 if not binary: - # If binary is False, then we need the ThresholdedRelu operator which is only available since opset 10. + # If binary is False, then we need the ThresholdedRelu + # operator which is only available since opset 10. min_opset = 10 if container.target_opset < min_opset: raise NotImplementedError( - f"Converting to ONNX for CountVectorizerModel is not supported in opset < {min_opset}" + f"Converting to ONNX for CountVectorizerModel " + f"is not supported in opset < {min_opset}" ) # Create a TfIdfVectorizer node with gram length set to 1 and mode set to "TF". - vectorizer_output_variable_name = scope.get_unique_variable_name("vectorizer_output") + vectorizer_output_variable_name = scope.get_unique_variable_name( + "vectorizer_output" + ) tfIdfVectorizer_attrs = { "name": scope.get_unique_operator_name("tfIdfVectorizer"), "min_gram_length": 1, @@ -45,9 +57,12 @@ def convert_count_vectorizer(scope: Scope, operator: Operator, container: ModelC **tfIdfVectorizer_attrs, ) - # In Spark's CountVectorizerModel, the comparison with minTF is inclusive, - # but in ThresholdedRelu (or Binarizer) node, the comparison with `alpha` (or `threshold`) is exclusive. - # So, we need to subtract epsilon from minTF to make the comparison with `alpha` (or `threshold`) effectively inclusive. + # In Spark's CountVectorizerModel, the comparison + # with minTF is inclusive, + # but in ThresholdedRelu (or Binarizer) node, the comparison + # with `alpha` (or `threshold`) is exclusive. + # So, we need to subtract epsilon from minTF to make the comparison + # with `alpha` (or `threshold`) effectively inclusive. epsilon = 1e-6 if binary: # Create a Binarizer node with threshold set to minTF - epsilon. @@ -82,4 +97,6 @@ def calculate_count_vectorizer_output_shapes(operator): operator.outputs[0].type = FloatTensorType([N, C]) -register_shape_calculator("pyspark.ml.feature.CountVectorizerModel", calculate_count_vectorizer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.CountVectorizerModel", calculate_count_vectorizer_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/dct.py b/onnxmltools/convert/sparkml/operator_converters/dct.py index 4c23b50f..27ccfbc9 100644 --- a/onnxmltools/convert/sparkml/operator_converters/dct.py +++ b/onnxmltools/convert/sparkml/operator_converters/dct.py @@ -16,24 +16,41 @@ def convert_sparkml_dct(scope, operator, container): # op = operator.raw_operator # inverse = op.getInverse() K = operator.inputs[0].type.shape[1] - two_sqrt_n = math.sqrt(2. / K) - sqrt_n = math.sqrt(1. / K) + two_sqrt_n = math.sqrt(2.0 / K) + sqrt_n = math.sqrt(1.0 / K) cosine_matrix = numpy.array( - [[math.cos(math.pi * (n + .5) * k / K) for n in range(0, K)] for k in range(0, K)] + [ + [math.cos(math.pi * (n + 0.5) * k / K) for n in range(0, K)] + for k in range(0, K) + ] ).transpose() - cosine_tensor = scope.get_unique_variable_name('cosine_tensor') - container.add_initializer(cosine_tensor, onnx_proto.TensorProto.FLOAT, - cosine_matrix.shape, list(cosine_matrix.flatten())) - matmul_result = scope.get_unique_variable_name('matmul_tensor') - apply_matmul(scope, [operator.inputs[0].full_name, cosine_tensor], matmul_result, container) + cosine_tensor = scope.get_unique_variable_name("cosine_tensor") + container.add_initializer( + cosine_tensor, + onnx_proto.TensorProto.FLOAT, + cosine_matrix.shape, + list(cosine_matrix.flatten()), + ) + matmul_result = scope.get_unique_variable_name("matmul_tensor") + apply_matmul( + scope, [operator.inputs[0].full_name, cosine_tensor], matmul_result, container + ) scale_vector = [sqrt_n if (k == 0) else two_sqrt_n for k in range(0, K)] - scale_tensor = scope.get_unique_variable_name('scale_tensor') - container.add_initializer(scale_tensor, onnx_proto.TensorProto.FLOAT, - [1, len(scale_vector)], scale_vector) - apply_mul(scope, [matmul_result, scale_tensor], operator.output_full_names, container, axis=1, broadcast=True) + scale_tensor = scope.get_unique_variable_name("scale_tensor") + container.add_initializer( + scale_tensor, onnx_proto.TensorProto.FLOAT, [1, len(scale_vector)], scale_vector + ) + apply_mul( + scope, + [matmul_result, scale_tensor], + operator.output_full_names, + container, + axis=1, + broadcast=True, + ) -register_converter('pyspark.ml.feature.DCT', convert_sparkml_dct) +register_converter("pyspark.ml.feature.DCT", convert_sparkml_dct) def calculate_sparkml_dct_output_shapes(operator): @@ -42,4 +59,4 @@ def calculate_sparkml_dct_output_shapes(operator): operator.outputs[0].type = copy.deepcopy(operator.inputs[0].type) -register_shape_calculator('pyspark.ml.feature.DCT', calculate_sparkml_dct_output_shapes) +register_shape_calculator("pyspark.ml.feature.DCT", calculate_sparkml_dct_output_shapes) diff --git a/onnxmltools/convert/sparkml/operator_converters/decision_tree_classifier.py b/onnxmltools/convert/sparkml/operator_converters/decision_tree_classifier.py index 0694c357..b80aa63c 100644 --- a/onnxmltools/convert/sparkml/operator_converters/decision_tree_classifier.py +++ b/onnxmltools/convert/sparkml/operator_converters/decision_tree_classifier.py @@ -1,13 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 import logging -import numpy as np from ...common.data_types import Int64TensorType, FloatTensorType from ...common.utils import check_input_and_output_numbers, check_input_and_output_types from ...common._registration import register_converter, register_shape_calculator from .tree_ensemble_common import ( - save_read_sparkml_model_data, sparkml_tree_dataset_to_sklearn, - add_tree_to_attribute_pairs, get_default_tree_classifier_attribute_pairs) + save_read_sparkml_model_data, + sparkml_tree_dataset_to_sklearn, + add_tree_to_attribute_pairs, + get_default_tree_classifier_attribute_pairs, +) from .tree_helper import rewrite_ids_and_process logger = logging.getLogger("onnxmltools") @@ -15,36 +17,52 @@ def convert_decision_tree_classifier(scope, operator, container): op = operator.raw_operator - op_type = 'TreeEnsembleClassifier' + op_type = "TreeEnsembleClassifier" attrs = get_default_tree_classifier_attribute_pairs() - attrs['name'] = scope.get_unique_operator_name(op_type) + attrs["name"] = scope.get_unique_operator_name(op_type) attrs["classlabels_int64s"] = list(range(0, op.numClasses)) logger.info("[convert_decision_tree_classifier] save_read_sparkml_model_data") - tree_df = save_read_sparkml_model_data(operator.raw_params['SparkSession'], op) + tree_df = save_read_sparkml_model_data(operator.raw_params["SparkSession"], op) logger.info("[convert_decision_tree_classifier] sparkml_tree_dataset_to_sklearn") tree = sparkml_tree_dataset_to_sklearn(tree_df, is_classifier=True) logger.info("[convert_decision_tree_classifier] add_tree_to_attribute_pairs") - add_tree_to_attribute_pairs(attrs, True, tree, 0, 1., 0, leaf_weights_are_counts=True) - logger.info("[convert_decision_tree_classifier] n_nodes=%d", len(attrs['nodes_nodeids'])) + add_tree_to_attribute_pairs( + attrs, True, tree, 0, 1.0, 0, leaf_weights_are_counts=True + ) + logger.info( + "[convert_decision_tree_classifier] n_nodes=%d", len(attrs["nodes_nodeids"]) + ) # Some values appear in an array of one element instead of a float. new_attrs = rewrite_ids_and_process(attrs, logger) - container.add_node(op_type, operator.input_full_names, [operator.outputs[0].full_name, - operator.outputs[1].full_name], op_domain='ai.onnx.ml', **new_attrs) + container.add_node( + op_type, + operator.input_full_names, + [operator.outputs[0].full_name, operator.outputs[1].full_name], + op_domain="ai.onnx.ml", + **new_attrs + ) -register_converter('pyspark.ml.classification.DecisionTreeClassificationModel', convert_decision_tree_classifier) +register_converter( + "pyspark.ml.classification.DecisionTreeClassificationModel", + convert_decision_tree_classifier, +) def calculate_decision_tree_classifier_output_shapes(operator): - check_input_and_output_numbers(operator, input_count_range=1, output_count_range=[1, 2]) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_numbers( + operator, input_count_range=1, output_count_range=[1, 2] + ) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) if len(operator.inputs[0].type.shape) != 2: - raise RuntimeError('Input must be a [N, C]-tensor') + raise RuntimeError("Input must be a [N, C]-tensor") N = operator.inputs[0].type.shape[0] @@ -53,5 +71,7 @@ def calculate_decision_tree_classifier_output_shapes(operator): operator.outputs[1].type = FloatTensorType([N, class_count]) -register_shape_calculator('pyspark.ml.classification.DecisionTreeClassificationModel', - calculate_decision_tree_classifier_output_shapes) +register_shape_calculator( + "pyspark.ml.classification.DecisionTreeClassificationModel", + calculate_decision_tree_classifier_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/decision_tree_regressor.py b/onnxmltools/convert/sparkml/operator_converters/decision_tree_regressor.py index a96188df..4bd76863 100644 --- a/onnxmltools/convert/sparkml/operator_converters/decision_tree_regressor.py +++ b/onnxmltools/convert/sparkml/operator_converters/decision_tree_regressor.py @@ -1,11 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 import logging from ...common.data_types import FloatTensorType -from ...common.tree_ensemble import add_tree_to_attribute_pairs, \ - get_default_tree_regressor_attribute_pairs +from ...common.tree_ensemble import ( + add_tree_to_attribute_pairs, + get_default_tree_regressor_attribute_pairs, +) from ...common.utils import check_input_and_output_numbers -from ...sparkml.operator_converters.decision_tree_classifier import save_read_sparkml_model_data -from ...sparkml.operator_converters.tree_ensemble_common import sparkml_tree_dataset_to_sklearn +from ...sparkml.operator_converters.decision_tree_classifier import ( + save_read_sparkml_model_data, +) +from ...sparkml.operator_converters.tree_ensemble_common import ( + sparkml_tree_dataset_to_sklearn, +) from ...common._registration import register_converter, register_shape_calculator from .tree_helper import rewrite_ids_and_process @@ -14,22 +20,29 @@ def convert_decision_tree_regressor(scope, operator, container): op = operator.raw_operator - op_type = 'TreeEnsembleRegressor' + op_type = "TreeEnsembleRegressor" attrs = get_default_tree_regressor_attribute_pairs() - attrs['name'] = scope.get_unique_operator_name(op_type) - attrs['n_targets'] = 1 + attrs["name"] = scope.get_unique_operator_name(op_type) + attrs["n_targets"] = 1 - tree_df = save_read_sparkml_model_data(operator.raw_params['SparkSession'], op) + tree_df = save_read_sparkml_model_data(operator.raw_params["SparkSession"], op) tree = sparkml_tree_dataset_to_sklearn(tree_df, is_classifier=False) - add_tree_to_attribute_pairs(attrs, False, tree, 0, 1., 0, False) + add_tree_to_attribute_pairs(attrs, False, tree, 0, 1.0, 0, False) new_attrs = rewrite_ids_and_process(attrs, logger) - container.add_node(op_type, operator.input_full_names, operator.output_full_names, - op_domain='ai.onnx.ml', **new_attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **new_attrs + ) -register_converter('pyspark.ml.regression.DecisionTreeRegressionModel', convert_decision_tree_regressor) +register_converter( + "pyspark.ml.regression.DecisionTreeRegressionModel", convert_decision_tree_regressor +) def calculate_decision_tree_regressor_output_shapes(operator): @@ -38,5 +51,7 @@ def calculate_decision_tree_regressor_output_shapes(operator): operator.outputs[0].type = FloatTensorType(shape=[N, 1]) -register_shape_calculator('pyspark.ml.regression.DecisionTreeRegressionModel', - calculate_decision_tree_regressor_output_shapes) +register_shape_calculator( + "pyspark.ml.regression.DecisionTreeRegressionModel", + calculate_decision_tree_regressor_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/element_wise_product.py b/onnxmltools/convert/sparkml/operator_converters/element_wise_product.py index 6404cc53..068111fb 100644 --- a/onnxmltools/convert/sparkml/operator_converters/element_wise_product.py +++ b/onnxmltools/convert/sparkml/operator_converters/element_wise_product.py @@ -9,13 +9,24 @@ def convert_element_wise_product(scope, operator, container): op = operator.raw_operator - scaling_vector = scope.get_unique_variable_name('scaling_vector') - container.add_initializer(scaling_vector, onnx_proto.TensorProto.FLOAT, - [1, len(op.getScalingVec())], op.getScalingVec()) - apply_mul(scope, [operator.inputs[0].full_name, scaling_vector], operator.output_full_names, container) - - -register_converter('pyspark.ml.feature.ElementwiseProduct', convert_element_wise_product) + scaling_vector = scope.get_unique_variable_name("scaling_vector") + container.add_initializer( + scaling_vector, + onnx_proto.TensorProto.FLOAT, + [1, len(op.getScalingVec())], + op.getScalingVec(), + ) + apply_mul( + scope, + [operator.inputs[0].full_name, scaling_vector], + operator.output_full_names, + container, + ) + + +register_converter( + "pyspark.ml.feature.ElementwiseProduct", convert_element_wise_product +) def calculate_element_wise_product_output_shapes(operator): @@ -25,4 +36,7 @@ def calculate_element_wise_product_output_shapes(operator): operator.outputs[0].type = FloatTensorType([N, operator.inputs[0].type.shape[1]]) -register_shape_calculator('pyspark.ml.feature.ElementwiseProduct', calculate_element_wise_product_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.ElementwiseProduct", + calculate_element_wise_product_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/gbt_classifier.py b/onnxmltools/convert/sparkml/operator_converters/gbt_classifier.py index 20e90121..84d798f4 100644 --- a/onnxmltools/convert/sparkml/operator_converters/gbt_classifier.py +++ b/onnxmltools/convert/sparkml/operator_converters/gbt_classifier.py @@ -3,8 +3,15 @@ from onnx import onnx_pb as onnx_proto from pyspark.ml.classification import GBTClassificationModel -from ...common._apply_operation import apply_neg, apply_concat, apply_mul, apply_exp, apply_add, \ - apply_argmax, apply_matmul +from ...common._apply_operation import ( + apply_neg, + apply_concat, + apply_mul, + apply_exp, + apply_add, + apply_argmax, + apply_matmul, +) from ...common.data_types import Int64TensorType, FloatTensorType from ...common.utils import check_input_and_output_numbers, check_input_and_output_types from ...common._registration import register_converter, register_shape_calculator @@ -15,59 +22,100 @@ def convert_gbt_classifier(scope, operator, container): op = operator.raw_operator regressor_output_names = [] - # spark implementation uses DecisionTreeRegressor (and not Classifier) for each tree in this forest + # spark implementation uses DecisionTreeRegressor + # (and not Classifier) for each tree in this forest for tree_model in op.trees: - regressor_op = scope.declare_local_operator(get_sparkml_operator_name(type(tree_model)), tree_model) + regressor_op = scope.declare_local_operator( + get_sparkml_operator_name(type(tree_model)), tree_model + ) regressor_op.raw_params = operator.raw_params regressor_op.inputs = operator.inputs - regressor_output = scope.declare_local_variable('regressor_prediction', FloatTensorType()) + regressor_output = scope.declare_local_variable( + "regressor_prediction", FloatTensorType() + ) regressor_output_names.append(regressor_output.full_name) regressor_op.outputs.append(regressor_output) convert_decision_tree_regressor(scope, regressor_op, container) regressor_op.is_evaluated = True - targets_tensor = scope.get_unique_variable_name('target_tensor') - weights_tensor = scope.get_unique_variable_name('weights_tensor') - container.add_initializer(weights_tensor, onnx_proto.TensorProto.FLOAT, [len(op.treeWeights), 1], op.treeWeights) - concatenated_predictions = scope.get_unique_variable_name('concatenated_predictions_tensor') - apply_concat(scope, regressor_output_names, concatenated_predictions, container, axis=1) - apply_matmul(scope, [concatenated_predictions, weights_tensor], targets_tensor, container) + targets_tensor = scope.get_unique_variable_name("target_tensor") + weights_tensor = scope.get_unique_variable_name("weights_tensor") + container.add_initializer( + weights_tensor, + onnx_proto.TensorProto.FLOAT, + [len(op.treeWeights), 1], + op.treeWeights, + ) + concatenated_predictions = scope.get_unique_variable_name( + "concatenated_predictions_tensor" + ) + apply_concat( + scope, regressor_output_names, concatenated_predictions, container, axis=1 + ) + apply_matmul( + scope, [concatenated_predictions, weights_tensor], targets_tensor, container + ) - # this is to calculate prediction and probability given the raw_prediction (= [-target, target]) - targets_neg_tensor = scope.get_unique_variable_name('target_neg_tensor') + # this is to calculate prediction and probability + # given the raw_prediction (= [-target, target]) + targets_neg_tensor = scope.get_unique_variable_name("target_neg_tensor") apply_neg(scope, targets_tensor, targets_neg_tensor, container) - raw_prediction_tensor = scope.get_unique_variable_name('raw_prediction_tensor') - apply_concat(scope, [targets_neg_tensor, targets_tensor], raw_prediction_tensor, container, - axis=1) + raw_prediction_tensor = scope.get_unique_variable_name("raw_prediction_tensor") + apply_concat( + scope, + [targets_neg_tensor, targets_tensor], + raw_prediction_tensor, + container, + axis=1, + ) if isinstance(op, GBTClassificationModel): - # this section is only for the classifier; for the regressor we don't calculate the probability - minus_two = scope.get_unique_variable_name('minus_two_tensor') + # this section is only for the classifier; + # for the regressor we don't calculate the probability + minus_two = scope.get_unique_variable_name("minus_two_tensor") container.add_initializer(minus_two, onnx_proto.TensorProto.FLOAT, [1], [-2.0]) - mul_output_tensor = scope.get_unique_variable_name('mul_output_tensor') - apply_mul(scope, [raw_prediction_tensor, minus_two], mul_output_tensor, container) - exp_output_tensor = scope.get_unique_variable_name('exp_output_tensor') + mul_output_tensor = scope.get_unique_variable_name("mul_output_tensor") + apply_mul( + scope, [raw_prediction_tensor, minus_two], mul_output_tensor, container + ) + exp_output_tensor = scope.get_unique_variable_name("exp_output_tensor") apply_exp(scope, mul_output_tensor, exp_output_tensor, container) - one_tensor = scope.get_unique_variable_name('one_tensor') + one_tensor = scope.get_unique_variable_name("one_tensor") container.add_initializer(one_tensor, onnx_proto.TensorProto.FLOAT, [1], [1.0]) - add_output_tensor = scope.get_unique_variable_name('add_output_tensor') + add_output_tensor = scope.get_unique_variable_name("add_output_tensor") apply_add(scope, [exp_output_tensor, one_tensor], add_output_tensor, container) - container.add_node('Reciprocal', add_output_tensor, operator.outputs[1].full_name, - name=scope.get_unique_operator_name('Reciprocal'), - op_version=6) + container.add_node( + "Reciprocal", + add_output_tensor, + operator.outputs[1].full_name, + name=scope.get_unique_operator_name("Reciprocal"), + op_version=6, + ) # to get Prediction from rawPrediction (or probability) - apply_argmax(scope, raw_prediction_tensor, operator.outputs[0].full_name, container, - axis=1, keepdims=0) + apply_argmax( + scope, + raw_prediction_tensor, + operator.outputs[0].full_name, + container, + axis=1, + keepdims=0, + ) -register_converter('pyspark.ml.classification.GBTClassificationModel', convert_gbt_classifier) -register_converter('pyspark.ml.regression.GBTRegressionModel', convert_gbt_classifier) +register_converter( + "pyspark.ml.classification.GBTClassificationModel", convert_gbt_classifier +) +register_converter("pyspark.ml.regression.GBTRegressionModel", convert_gbt_classifier) def calculate_gbt_classifier_output_shapes(operator): - check_input_and_output_numbers(operator, input_count_range=1, output_count_range=[1, 2]) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_numbers( + operator, input_count_range=1, output_count_range=[1, 2] + ) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) if len(operator.inputs[0].type.shape) != 2: - raise RuntimeError('Input must be a [N, C]-tensor') + raise RuntimeError("Input must be a [N, C]-tensor") N = operator.inputs[0].type.shape[0] operator.outputs[0].type = Int64TensorType(shape=[N]) @@ -76,7 +124,10 @@ def calculate_gbt_classifier_output_shapes(operator): operator.outputs[1].type = FloatTensorType([N, class_count]) -register_shape_calculator('pyspark.ml.classification.GBTClassificationModel', - calculate_gbt_classifier_output_shapes) -register_shape_calculator('pyspark.ml.regression.GBTRegressionModel', - calculate_gbt_classifier_output_shapes) +register_shape_calculator( + "pyspark.ml.classification.GBTClassificationModel", + calculate_gbt_classifier_output_shapes, +) +register_shape_calculator( + "pyspark.ml.regression.GBTRegressionModel", calculate_gbt_classifier_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/imputer.py b/onnxmltools/convert/sparkml/operator_converters/imputer.py index 4e74c7a3..7a50fa17 100644 --- a/onnxmltools/convert/sparkml/operator_converters/imputer.py +++ b/onnxmltools/convert/sparkml/operator_converters/imputer.py @@ -7,54 +7,80 @@ from ...common._registration import register_converter, register_shape_calculator from ...common._topology import Operator, Scope from pyspark.ml.feature import ImputerModel -from typing import List + def convert_imputer(scope: Scope, operator: Operator, container): op: ImputerModel = operator.raw_operator - op_type = 'Imputer' + op_type = "Imputer" name = scope.get_unique_operator_name(op_type) - attrs = {'name': name} + attrs = {"name": name} input_type = operator.inputs[0].type surrogates = op.surrogateDF.toPandas().values[0].tolist() - value = op.getOrDefault('missingValue') - + value = op.getOrDefault("missingValue") + if isinstance(input_type, FloatTensorType): - attrs['imputed_value_floats'] = surrogates - attrs['replaced_value_float'] = value + attrs["imputed_value_floats"] = surrogates + attrs["replaced_value_float"] = value elif isinstance(input_type, Int64TensorType): - attrs['imputed_value_int64s'] = [int(x) for x in surrogates] - attrs['replaced_value_int64'] = int(value) + attrs["imputed_value_int64s"] = [int(x) for x in surrogates] + attrs["replaced_value_int64"] = int(value) else: raise RuntimeError("Invalid input type: " + input_type) if len(operator.inputs) > 1: - concatenated_output = scope.get_unique_variable_name('concat_tensor') - container.add_node('Concat', operator.input_full_names, concatenated_output, - name=scope.get_unique_operator_name('Concat'), - op_version=4, - axis=1) - imputed_output = scope.get_unique_variable_name('imputed_tensor') - container.add_node(op_type, concatenated_output, imputed_output, op_domain='ai.onnx.ml', **attrs) - container.add_node('Split', imputed_output, operator.output_full_names, - name=scope.get_unique_operator_name('Split'), - op_version=2, - axis=1, - split=[1] * len(operator.output_full_names)) + concatenated_output = scope.get_unique_variable_name("concat_tensor") + container.add_node( + "Concat", + operator.input_full_names, + concatenated_output, + name=scope.get_unique_operator_name("Concat"), + op_version=4, + axis=1, + ) + imputed_output = scope.get_unique_variable_name("imputed_tensor") + container.add_node( + op_type, + concatenated_output, + imputed_output, + op_domain="ai.onnx.ml", + **attrs + ) + container.add_node( + "Split", + imputed_output, + operator.output_full_names, + name=scope.get_unique_operator_name("Split"), + op_version=2, + axis=1, + split=[1] * len(operator.output_full_names), + ) else: - container.add_node(op_type, operator.inputs[0].full_name, operator.output_full_names[0], - op_domain='ai.onnx.ml', - **attrs) + container.add_node( + op_type, + operator.inputs[0].full_name, + operator.output_full_names[0], + op_domain="ai.onnx.ml", + **attrs + ) + + +register_converter("pyspark.ml.feature.ImputerModel", convert_imputer) -register_converter('pyspark.ml.feature.ImputerModel', convert_imputer) def calculate_imputer_output_shapes(operator): - check_input_and_output_numbers(operator, output_count_range=[1, len(operator.outputs)]) - check_input_and_output_types(operator, - good_input_types=[FloatTensorType, Int64TensorType], - good_output_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_numbers( + operator, output_count_range=[1, len(operator.outputs)] + ) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType, Int64TensorType], + good_output_types=[FloatTensorType, Int64TensorType], + ) input_type = copy.deepcopy(operator.inputs[0].type) for output in operator.outputs: output.type = input_type -register_shape_calculator('pyspark.ml.feature.ImputerModel', calculate_imputer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.ImputerModel", calculate_imputer_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/index_to_string.py b/onnxmltools/convert/sparkml/operator_converters/index_to_string.py index 8e87d5b3..0e726f0a 100644 --- a/onnxmltools/convert/sparkml/operator_converters/index_to_string.py +++ b/onnxmltools/convert/sparkml/operator_converters/index_to_string.py @@ -10,26 +10,34 @@ def convert_index_to_string(scope, operator, container): op = operator.raw_operator - op_type = 'LabelEncoder' - if not op.isDefined('labels') or len(op.getLabels()) == 0: - raise SparkMlConversionError('Labels must be specified for IndexToString Transformer') + op_type = "LabelEncoder" + if not op.isDefined("labels") or len(op.getLabels()) == 0: + raise SparkMlConversionError( + "Labels must be specified for IndexToString Transformer" + ) attrs = { - 'name': scope.get_unique_operator_name(op_type), - 'classes_strings': [str(c) for c in op.getLabels()], - 'default_string': '__unknown__' + "name": scope.get_unique_operator_name(op_type), + "classes_strings": [str(c) for c in op.getLabels()], + "default_string": "__unknown__", } - container.add_node(op_type, operator.input_full_names, operator.output_full_names, - op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('pyspark.ml.feature.IndexToString', convert_index_to_string) +register_converter("pyspark.ml.feature.IndexToString", convert_index_to_string) def calculate_index_to_string_output_shapes(operator): - ''' - This function just copy the input shape to the output because label encoder only alters input features' values, not + """ + This function just copy the input shape to the output + because label encoder only alters input features' values, not their shape. - ''' + """ check_input_and_output_numbers(operator, output_count_range=1) check_input_and_output_types(operator, good_input_types=[Int64TensorType]) @@ -37,4 +45,6 @@ def calculate_index_to_string_output_shapes(operator): operator.outputs[0].type = StringTensorType(input_shape) -register_shape_calculator('pyspark.ml.feature.IndexToString', calculate_index_to_string_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.IndexToString", calculate_index_to_string_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/k_means.py b/onnxmltools/convert/sparkml/operator_converters/k_means.py index 45909394..e29f474b 100644 --- a/onnxmltools/convert/sparkml/operator_converters/k_means.py +++ b/onnxmltools/convert/sparkml/operator_converters/k_means.py @@ -4,46 +4,59 @@ from ...common._topology import Operator, Scope, ModelComponentContainer from ....proto import onnx_proto from pyspark.ml.clustering import KMeansModel -from typing import List import numpy as np -def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: ModelComponentContainer): + +def convert_sparkml_k_means_model( + scope: Scope, operator: Operator, container: ModelComponentContainer +): if container.target_opset < 7: - raise NotImplementedError("Converting to ONNX for KMeansModel is not supported in opset < 7") - + raise NotImplementedError( + "Converting to ONNX for KMeansModel is not supported in opset < 7" + ) + op: KMeansModel = operator.raw_operator centers: np.ndarray = np.vstack(op.clusterCenters()) - K = centers.shape[0] # number of clusters - C = operator.inputs[0].type.shape[1] # Number of features from input - + K = centers.shape[0] # number of clusters + C = operator.inputs[0].type.shape[1] # Number of features from input + if centers.shape[1] != C: - raise ValueError(f"Number of features {centers.shape[1]} in input does not match number of features in centers {C}") + raise ValueError( + f"Number of features {centers.shape[1]} " + f"in input does not match number of features in centers {C}" + ) # [K x C] centers_variable_name = scope.get_unique_variable_name("centers") container.add_initializer( - centers_variable_name, - onnx_proto.TensorProto.FLOAT, - centers.shape, - centers.flatten().astype(np.float32) - ) + centers_variable_name, + onnx_proto.TensorProto.FLOAT, + centers.shape, + centers.flatten().astype(np.float32), + ) distance_output_variable_name = scope.get_unique_variable_name("distance_output") if op.getDistanceMeasure() == "euclidean": # [1 x K] - centers_row_squared_sum_variable_name = scope.get_unique_variable_name("centers_row_squared_sum") - centers_row_squared_sum = np.sum(centers**2,axis=-1).flatten().astype(np.float32) + centers_row_squared_sum_variable_name = scope.get_unique_variable_name( + "centers_row_squared_sum" + ) + centers_row_squared_sum = ( + np.sum(centers**2, axis=-1).flatten().astype(np.float32) + ) container.add_initializer( - centers_row_squared_sum_variable_name, - onnx_proto.TensorProto.FLOAT, - [1, K], - centers_row_squared_sum - ) + centers_row_squared_sum_variable_name, + onnx_proto.TensorProto.FLOAT, + [1, K], + centers_row_squared_sum, + ) # input_row_squared_sum: [N x 1] - input_row_squared_sum_variable_name = scope.get_unique_variable_name("input_row_squared_sum") + input_row_squared_sum_variable_name = scope.get_unique_variable_name( + "input_row_squared_sum" + ) reduce_sum_square_attrs = { "name": scope.get_unique_operator_name("input_row_squared_sum"), "axes": [1], @@ -53,7 +66,7 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M op_type="ReduceSumSquare", inputs=[operator.inputs[0].full_name], outputs=[input_row_squared_sum_variable_name], - **reduce_sum_square_attrs + **reduce_sum_square_attrs, ) # -2 * input * Transpose(Center) + input_row_squared_sum: [N x K] @@ -65,14 +78,19 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M "transB": 1, } container.add_node( - op_type="Gemm", - inputs=[operator.inputs[0].full_name, centers_variable_name, input_row_squared_sum_variable_name], + op_type="Gemm", + inputs=[ + operator.inputs[0].full_name, + centers_variable_name, + input_row_squared_sum_variable_name, + ], outputs=[gemm_output_variable_name], op_version=7, - **gemm_attrs + **gemm_attrs, ) - - # Euclidean Distance Squared = input_row_squared_sum - 2 * input * Transpose(Center) + Transpose(centers_row_squared_sum) + + # Euclidean Distance Squared = input_row_squared_sum - 2 * + # input * Transpose(Center) + Transpose(centers_row_squared_sum) # [N x K] container.add_node( op_type="Add", @@ -82,17 +100,23 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M ) elif op.getDistanceMeasure() == "cosine": # centers_row_norm2: [1 x K] - centers_row_norm2_variable_name = scope.get_unique_variable_name("centers_row_norm2") - centers_row_norm2 = np.linalg.norm(centers, ord = 2, axis=1).flatten().astype(np.float32) + centers_row_norm2_variable_name = scope.get_unique_variable_name( + "centers_row_norm2" + ) + centers_row_norm2 = ( + np.linalg.norm(centers, ord=2, axis=1).flatten().astype(np.float32) + ) container.add_initializer( - centers_row_norm2_variable_name, - onnx_proto.TensorProto.FLOAT, - [1, K], - centers_row_norm2 - ) - + centers_row_norm2_variable_name, + onnx_proto.TensorProto.FLOAT, + [1, K], + centers_row_norm2, + ) + # input_row_norm2: [N x 1] - input_row_norm2_variable_name = scope.get_unique_variable_name("input_row_norm2") + input_row_norm2_variable_name = scope.get_unique_variable_name( + "input_row_norm2" + ) reduce_l2_attrs = { "name": scope.get_unique_operator_name("input_row_norm2"), "axes": [1], @@ -102,16 +126,16 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M op_type="ReduceL2", inputs=[operator.inputs[0].full_name], outputs=[input_row_norm2_variable_name], - **reduce_l2_attrs + **reduce_l2_attrs, ) # input * Transpose(Center): [N x K] zeros_variable_name = scope.get_unique_variable_name("zeros") container.add_initializer( - zeros_variable_name, - onnx_proto.TensorProto.FLOAT, - [1, K], - np.zeros([1, K]).flatten().astype(np.float32) + zeros_variable_name, + onnx_proto.TensorProto.FLOAT, + [1, K], + np.zeros([1, K]).flatten().astype(np.float32), ) gemm_output_variable_name = scope.get_unique_variable_name("gemm_output") gemm_attrs = { @@ -120,11 +144,15 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M "transB": 1, } container.add_node( - op_type="Gemm", - inputs=[operator.inputs[0].full_name, centers_variable_name, zeros_variable_name], + op_type="Gemm", + inputs=[ + operator.inputs[0].full_name, + centers_variable_name, + zeros_variable_name, + ], outputs=[gemm_output_variable_name], op_version=7, - **gemm_attrs + **gemm_attrs, ) # Cosine similarity = gemm_output / input_row_norm2 / centers_row_norm2: [N x K] @@ -135,7 +163,9 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M outputs=[div_output_variable_name], op_version=7, ) - cosine_similarity_output_variable_name = scope.get_unique_variable_name("cosine_similarity_output") + cosine_similarity_output_variable_name = scope.get_unique_variable_name( + "cosine_similarity_output" + ) container.add_node( op_type="Div", inputs=[div_output_variable_name, centers_row_norm2_variable_name], @@ -161,20 +191,23 @@ def convert_sparkml_k_means_model(scope: Scope, operator: Operator, container: M op_type="ArgMin", inputs=[distance_output_variable_name], outputs=[operator.outputs[0].full_name], - **argmin_attrs + **argmin_attrs, ) -register_converter('pyspark.ml.clustering.KMeansModel', convert_sparkml_k_means_model) + +register_converter("pyspark.ml.clustering.KMeansModel", convert_sparkml_k_means_model) def calculate_k_means_model_output_shapes(operator: Operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) if len(operator.inputs[0].type.shape) != 2: - raise RuntimeError('Input must be a [N, C]-tensor') + raise RuntimeError("Input must be a [N, C]-tensor") N = operator.inputs[0].type.shape[0] operator.outputs[0].type = Int64TensorType(shape=[N]) -register_shape_calculator('pyspark.ml.clustering.KMeansModel', calculate_k_means_model_output_shapes) +register_shape_calculator( + "pyspark.ml.clustering.KMeansModel", calculate_k_means_model_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/linear_classifier.py b/onnxmltools/convert/sparkml/operator_converters/linear_classifier.py index 5e0aeb27..126ba1d9 100644 --- a/onnxmltools/convert/sparkml/operator_converters/linear_classifier.py +++ b/onnxmltools/convert/sparkml/operator_converters/linear_classifier.py @@ -9,8 +9,8 @@ def convert_sparkml_linear_classifier(scope, operator, container): op = operator.raw_operator - op_type = 'LinearClassifier' - attrs = {'name': scope.get_unique_operator_name(op_type)} + op_type = "LinearClassifier" + attrs = {"name": scope.get_unique_operator_name(op_type)} if op.numClasses == 2: coefficients = op.coefficients.toArray().tolist() @@ -25,51 +25,77 @@ def convert_sparkml_linear_classifier(scope, operator, container): if op.numClasses > 2: coefficients = op.coefficientMatrix.toArray().ravel().tolist() intercepts = op.interceptVector.toArray().ravel().tolist() - attrs['post_transform'] = 'LOGISTIC' + attrs["post_transform"] = "LOGISTIC" elif isinstance(op, LinearSVCModel): - attrs['post_transform'] = 'NONE' + attrs["post_transform"] = "NONE" else: if op.numClasses >= 2: - attrs['post_transform'] = 'SOFTMAX' + attrs["post_transform"] = "SOFTMAX" else: - attrs['post_transform'] = 'LOGISTIC' + attrs["post_transform"] = "LOGISTIC" - attrs['coefficients'] = coefficients - attrs['intercepts'] = intercepts - attrs['multi_class'] = 1 if op.numClasses >= 2 else 0 + attrs["coefficients"] = coefficients + attrs["intercepts"] = intercepts + attrs["multi_class"] = 1 if op.numClasses >= 2 else 0 attrs["classlabels_ints"] = list(range(0, op.numClasses)) import pprint + pprint.pprint(attrs) label_name = operator.outputs[0].full_name if not isinstance(operator.raw_operator, LinearSVCModel): - probability_tensor_name = scope.get_unique_variable_name('probability_tensor') - container.add_node(op_type, operator.inputs[0].full_name, - [label_name, probability_tensor_name], - op_domain='ai.onnx.ml', **attrs) + probability_tensor_name = scope.get_unique_variable_name("probability_tensor") + container.add_node( + op_type, + operator.inputs[0].full_name, + [label_name, probability_tensor_name], + op_domain="ai.onnx.ml", + **attrs + ) # Make sure the probability sum is 1 over all classes - normalizer_type = 'Normalizer' - normalizer_attrs = {'name': scope.get_unique_operator_name(normalizer_type), 'norm': 'L1'} - container.add_node(normalizer_type, probability_tensor_name, operator.outputs[1].full_name, - op_domain='ai.onnx.ml', **normalizer_attrs) + normalizer_type = "Normalizer" + normalizer_attrs = { + "name": scope.get_unique_operator_name(normalizer_type), + "norm": "L1", + } + container.add_node( + normalizer_type, + probability_tensor_name, + operator.outputs[1].full_name, + op_domain="ai.onnx.ml", + **normalizer_attrs + ) else: # add a dummy output variable since onnx LinearClassifier has 2 - unused_probabilities_output = scope.get_unique_variable_name('probabilities') - container.add_node(op_type, operator.inputs[0].full_name, - [label_name,unused_probabilities_output], - op_domain='ai.onnx.ml', **attrs) + unused_probabilities_output = scope.get_unique_variable_name("probabilities") + container.add_node( + op_type, + operator.inputs[0].full_name, + [label_name, unused_probabilities_output], + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('pyspark.ml.classification.LogisticRegressionModel', convert_sparkml_linear_classifier) -register_converter('pyspark.ml.classification.LinearSVCModel', convert_sparkml_linear_classifier) +register_converter( + "pyspark.ml.classification.LogisticRegressionModel", + convert_sparkml_linear_classifier, +) +register_converter( + "pyspark.ml.classification.LinearSVCModel", convert_sparkml_linear_classifier +) def calculate_linear_classifier_output_shapes(operator): - check_input_and_output_numbers(operator, input_count_range=1, output_count_range=[1, 2]) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_numbers( + operator, input_count_range=1, output_count_range=[1, 2] + ) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) if len(operator.inputs[0].type.shape) != 2: - raise RuntimeError('Input must be a [N, C]-tensor') + raise RuntimeError("Input must be a [N, C]-tensor") N = operator.inputs[0].type.shape[0] operator.outputs[0].type = Int64TensorType(shape=[N]) @@ -78,5 +104,11 @@ def calculate_linear_classifier_output_shapes(operator): operator.outputs[1].type = FloatTensorType([N, class_count]) -register_shape_calculator('pyspark.ml.classification.LogisticRegressionModel', calculate_linear_classifier_output_shapes) -register_shape_calculator('pyspark.ml.classification.LinearSVCModel', calculate_linear_classifier_output_shapes) +register_shape_calculator( + "pyspark.ml.classification.LogisticRegressionModel", + calculate_linear_classifier_output_shapes, +) +register_shape_calculator( + "pyspark.ml.classification.LinearSVCModel", + calculate_linear_classifier_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/linear_regressor.py b/onnxmltools/convert/sparkml/operator_converters/linear_regressor.py index bf09b9aa..17c658c2 100644 --- a/onnxmltools/convert/sparkml/operator_converters/linear_regressor.py +++ b/onnxmltools/convert/sparkml/operator_converters/linear_regressor.py @@ -9,20 +9,32 @@ def convert_sparkml_linear_regressor(scope, operator, container): op = operator.raw_operator - op_type = 'LinearRegressor' + op_type = "LinearRegressor" attrs = { - 'name': scope.get_unique_operator_name(op_type), - 'coefficients': op.coefficients.astype(float), - 'intercepts': ( + "name": scope.get_unique_operator_name(op_type), + "coefficients": op.coefficients.astype(float), + "intercepts": ( op.intercept.astype(float) if isinstance(op.intercept, collections.abc.Iterable) - else [float(op.intercept)]) + else [float(op.intercept)] + ), } - container.add_node(op_type, operator.input_full_names, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + **attrs + ) -register_converter('pyspark.ml.regression.LinearRegressionModel', convert_sparkml_linear_regressor) -register_converter('pyspark.ml.regression.GeneralizedLinearRegressionModel', convert_sparkml_linear_regressor) +register_converter( + "pyspark.ml.regression.LinearRegressionModel", convert_sparkml_linear_regressor +) +register_converter( + "pyspark.ml.regression.GeneralizedLinearRegressionModel", + convert_sparkml_linear_regressor, +) def calculate_linear_regressor_output_shapes(operator): @@ -32,6 +44,11 @@ def calculate_linear_regressor_output_shapes(operator): operator.outputs[0].type = FloatTensorType([N, 1]) -register_shape_calculator('pyspark.ml.regression.LinearRegressionModel', calculate_linear_regressor_output_shapes) -register_shape_calculator('pyspark.ml.regression.GeneralizedLinearRegressionModel', - calculate_linear_regressor_output_shapes) +register_shape_calculator( + "pyspark.ml.regression.LinearRegressionModel", + calculate_linear_regressor_output_shapes, +) +register_shape_calculator( + "pyspark.ml.regression.GeneralizedLinearRegressionModel", + calculate_linear_regressor_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/min_hash_lsh.py b/onnxmltools/convert/sparkml/operator_converters/min_hash_lsh.py index ab12f71e..977557d7 100644 --- a/onnxmltools/convert/sparkml/operator_converters/min_hash_lsh.py +++ b/onnxmltools/convert/sparkml/operator_converters/min_hash_lsh.py @@ -2,8 +2,13 @@ from onnx import onnx_pb as onnx_proto from ...common._apply_operation import ( - apply_add, apply_mul, apply_sum, apply_div, apply_sub, - apply_concat, apply_cast) + apply_add, + apply_mul, + apply_div, + apply_sub, + apply_concat, + apply_cast, +) from ...common._registration import register_converter, register_shape_calculator from ...common.data_types import FloatTensorType, DoubleTensorType from ...common.utils import check_input_and_output_numbers, check_input_and_output_types @@ -17,7 +22,7 @@ def get_rand_coefficients(operator): global g_rand_coefficients if not g_rand_coefficients: g_rand_coefficients = save_read_sparkml_model_data( - operator.raw_params['SparkSession'], operator.raw_operator + operator.raw_params["SparkSession"], operator.raw_operator ).first()[0] return g_rand_coefficients @@ -28,56 +33,71 @@ def convert_min_hash_lsh(scope, operator, container): coeffs = [] for i in range(0, len(rand_coefficients), 2): coeffs.append((rand_coefficients[i], rand_coefficients[i + 1])) - one = scope.get_unique_variable_name('one_tensor') + one = scope.get_unique_variable_name("one_tensor") container.add_initializer(one, int_type, [1], [1]) - prime = scope.get_unique_variable_name('prime_tensor') + prime = scope.get_unique_variable_name("prime_tensor") container.add_initializer(prime, int_type, [1], [MinHashLSH_HASH_PRIME]) - non_zeros_int = scope.get_unique_variable_name('non_zero_int_tensor') - container.add_node('NonZero', operator.input_full_names, non_zeros_int, op_version=9) - non_zeros = scope.get_unique_variable_name('non_zeros_tensor') + non_zeros_int = scope.get_unique_variable_name("non_zero_int_tensor") + container.add_node( + "NonZero", operator.input_full_names, non_zeros_int, op_version=9 + ) + non_zeros = scope.get_unique_variable_name("non_zeros_tensor") apply_cast(scope, non_zeros_int, non_zeros, container, to=int_type) remainders = [] for coeff in coeffs: - one_added = scope.get_unique_variable_name('one_added_tensor') + one_added = scope.get_unique_variable_name("one_added_tensor") apply_add(scope, [one, non_zeros], one_added, container) - a = scope.get_unique_variable_name('a_coeff_tensor') + a = scope.get_unique_variable_name("a_coeff_tensor") container.add_initializer(a, int_type, [1], [coeff[0]]) - b = scope.get_unique_variable_name('b_coeff_tensor') + b = scope.get_unique_variable_name("b_coeff_tensor") container.add_initializer(b, int_type, [1], [coeff[1]]) - coeff_0_times = scope.get_unique_variable_name('coeff_0_times_tensor') + coeff_0_times = scope.get_unique_variable_name("coeff_0_times_tensor") apply_mul(scope, [a, one_added], coeff_0_times, container) - coeff_1_added = scope.get_unique_variable_name('coeff_1_added_tensor') + coeff_1_added = scope.get_unique_variable_name("coeff_1_added_tensor") apply_add(scope, [b, coeff_0_times], coeff_1_added, container) # this is for Mod - div_by_prime = scope.get_unique_variable_name('div_by_prime_tensor') + div_by_prime = scope.get_unique_variable_name("div_by_prime_tensor") apply_div(scope, [coeff_1_added, prime], div_by_prime, container) - prime_x_floor = scope.get_unique_variable_name('prime_x_floor_tensor') + prime_x_floor = scope.get_unique_variable_name("prime_x_floor_tensor") apply_mul(scope, [div_by_prime, prime], prime_x_floor, container) - remainder = scope.get_unique_variable_name('remainder_tensor') + remainder = scope.get_unique_variable_name("remainder_tensor") apply_sub(scope, [coeff_1_added, prime_x_floor], remainder, container) - float_remainder = scope.get_unique_variable_name('float_remainder_tensor') - apply_cast(scope, remainder, float_remainder, container, to=onnx_proto.TensorProto.FLOAT) - min_remainder = scope.get_unique_variable_name('min_remainder_tensor') - container.add_node('ReduceMin', float_remainder, min_remainder, - op_version=1, - axes=[1], - keepdims=1) + float_remainder = scope.get_unique_variable_name("float_remainder_tensor") + apply_cast( + scope, + remainder, + float_remainder, + container, + to=onnx_proto.TensorProto.FLOAT, + ) + min_remainder = scope.get_unique_variable_name("min_remainder_tensor") + container.add_node( + "ReduceMin", + float_remainder, + min_remainder, + op_version=1, + axes=[1], + keepdims=1, + ) remainders.append(min_remainder) apply_concat(scope, remainders, operator.output_full_names, container, axis=1) -register_converter('pyspark.ml.feature.MinHashLSHModel', convert_min_hash_lsh) +register_converter("pyspark.ml.feature.MinHashLSHModel", convert_min_hash_lsh) def calculate_min_hash_lsh_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) check_input_and_output_types( - operator, good_input_types=[FloatTensorType, DoubleTensorType]) + operator, good_input_types=[FloatTensorType, DoubleTensorType] + ) N = operator.inputs[0].type.shape[0] C = len(get_rand_coefficients(operator)) // 2 operator.outputs[0].type = FloatTensorType([N, C]) -register_shape_calculator('pyspark.ml.feature.MinHashLSHModel', calculate_min_hash_lsh_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.MinHashLSHModel", calculate_min_hash_lsh_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/mlp_classifier.py b/onnxmltools/convert/sparkml/operator_converters/mlp_classifier.py index 786df1cd..0fab5ad7 100644 --- a/onnxmltools/convert/sparkml/operator_converters/mlp_classifier.py +++ b/onnxmltools/convert/sparkml/operator_converters/mlp_classifier.py @@ -11,7 +11,9 @@ import numpy as np -def convert_sparkml_mlp_classifier(scope: Scope, operator: Operator, container: ModelComponentContainer): +def convert_sparkml_mlp_classifier( + scope: Scope, operator: Operator, container: ModelComponentContainer +): op: MultilayerPerceptronClassificationModel = operator.raw_operator layers: List[int] = op.getLayers() weights: np.ndarray = op.weights.toArray() @@ -20,7 +22,9 @@ def convert_sparkml_mlp_classifier(scope: Scope, operator: Operator, container: input: str for i in range(len(layers) - 1): - weight_matrix = weights[offset : offset + layers[i] * layers[i + 1]].reshape([layers[i], layers[i + 1]]) + weight_matrix = weights[offset : offset + layers[i] * layers[i + 1]].reshape( + [layers[i], layers[i + 1]] + ) offset += layers[i] * layers[i + 1] bias_vector = weights[offset : offset + layers[i + 1]] offset += layers[i + 1] @@ -38,7 +42,10 @@ def convert_sparkml_mlp_classifier(scope: Scope, operator: Operator, container: bias_variable = scope.get_unique_variable_name("b") container.add_initializer( - bias_variable, onnx_proto.TensorProto.FLOAT, bias_vector.shape, bias_vector.astype(np.float32), + bias_variable, + onnx_proto.TensorProto.FLOAT, + bias_vector.shape, + bias_vector.astype(np.float32), ) gemm_output_variable = scope.get_unique_variable_name("gemm_output") @@ -74,18 +81,25 @@ def convert_sparkml_mlp_classifier(scope: Scope, operator: Operator, container: [operator.outputs[0].full_name], name=scope.get_unique_operator_name("ArgMax"), axis=1, - keepdims = 0, + keepdims=0, ) -register_converter("pyspark.ml.classification.MultilayerPerceptronClassificationModel", convert_sparkml_mlp_classifier) +register_converter( + "pyspark.ml.classification.MultilayerPerceptronClassificationModel", + convert_sparkml_mlp_classifier, +) def calculate_mlp_classifier_output_shapes(operator: Operator): op: MultilayerPerceptronClassificationModel = operator.raw_operator - check_input_and_output_numbers(operator, input_count_range=1, output_count_range=[1, 2]) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_numbers( + operator, input_count_range=1, output_count_range=[1, 2] + ) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) if len(operator.inputs[0].type.shape) != 2: raise RuntimeError("Input must be a [N, C]-tensor") @@ -97,5 +111,6 @@ def calculate_mlp_classifier_output_shapes(operator: Operator): register_shape_calculator( - "pyspark.ml.classification.MultilayerPerceptronClassificationModel", calculate_mlp_classifier_output_shapes + "pyspark.ml.classification.MultilayerPerceptronClassificationModel", + calculate_mlp_classifier_output_shapes, ) diff --git a/onnxmltools/convert/sparkml/operator_converters/naive_bayes.py b/onnxmltools/convert/sparkml/operator_converters/naive_bayes.py index 7cb3d185..1bf2f620 100644 --- a/onnxmltools/convert/sparkml/operator_converters/naive_bayes.py +++ b/onnxmltools/convert/sparkml/operator_converters/naive_bayes.py @@ -10,112 +10,188 @@ def convert_sparkml_naive_bayes(scope, operator, container): op = operator.raw_operator - model_type = op.getOrDefault('modelType') + model_type = op.getOrDefault("modelType") theta_np = numpy.array(op.theta.toArray()) transposed_theta = numpy.transpose(theta_np) - raw_predictions_tensor = scope.get_unique_variable_name('raw_predictions_tensor') - if model_type == 'bernoulli': + raw_predictions_tensor = scope.get_unique_variable_name("raw_predictions_tensor") + if model_type == "bernoulli": negTheta = numpy.log(1.0 - numpy.exp(theta_np)) thetaMinusLogTheta = numpy.transpose(theta_np - negTheta) ones = numpy.ones((theta_np.shape[1], 1)) negThetaSum = numpy.matmul(negTheta, ones).flatten() - thetaMinusLogTheta_tensor = scope.get_unique_variable_name('thetaMinusLogTheta_tensor') - container.add_initializer(thetaMinusLogTheta_tensor, onnx_proto.TensorProto.FLOAT, - list(thetaMinusLogTheta.shape), thetaMinusLogTheta.flatten().tolist()) - negThetaSum_tensor = scope.get_unique_variable_name('negThetaSum_tensor') - container.add_initializer(negThetaSum_tensor, onnx_proto.TensorProto.FLOAT, - list(negThetaSum.shape), negThetaSum.flatten().tolist()) - prior_tensor = scope.get_unique_variable_name('prior_tensor') - container.add_initializer(prior_tensor, onnx_proto.TensorProto.FLOAT, [len(op.pi)], op.pi.flatten().tolist()) - probability1_output = scope.get_unique_variable_name('temp_probability') - container.add_node('MatMul', [operator.input_full_names[0], thetaMinusLogTheta_tensor], probability1_output, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('MatMul'), - op_version=9) - probability2_output = scope.get_unique_variable_name('temp_probability') - container.add_node('Add', [probability1_output, prior_tensor], probability2_output, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Add'), - op_version=7) - container.add_node('Add', [probability2_output,negThetaSum_tensor], raw_predictions_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Add'), - op_version=7) + thetaMinusLogTheta_tensor = scope.get_unique_variable_name( + "thetaMinusLogTheta_tensor" + ) + container.add_initializer( + thetaMinusLogTheta_tensor, + onnx_proto.TensorProto.FLOAT, + list(thetaMinusLogTheta.shape), + thetaMinusLogTheta.flatten().tolist(), + ) + negThetaSum_tensor = scope.get_unique_variable_name("negThetaSum_tensor") + container.add_initializer( + negThetaSum_tensor, + onnx_proto.TensorProto.FLOAT, + list(negThetaSum.shape), + negThetaSum.flatten().tolist(), + ) + prior_tensor = scope.get_unique_variable_name("prior_tensor") + container.add_initializer( + prior_tensor, + onnx_proto.TensorProto.FLOAT, + [len(op.pi)], + op.pi.flatten().tolist(), + ) + probability1_output = scope.get_unique_variable_name("temp_probability") + container.add_node( + "MatMul", + [operator.input_full_names[0], thetaMinusLogTheta_tensor], + probability1_output, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("MatMul"), + op_version=9, + ) + probability2_output = scope.get_unique_variable_name("temp_probability") + container.add_node( + "Add", + [probability1_output, prior_tensor], + probability2_output, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Add"), + op_version=7, + ) + container.add_node( + "Add", + [probability2_output, negThetaSum_tensor], + raw_predictions_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Add"), + op_version=7, + ) else: - probability1_output = scope.get_unique_variable_name('temp_probability') - theta_tensor = scope.get_unique_variable_name('theta_tensor') - container.add_initializer(theta_tensor, onnx_proto.TensorProto.FLOAT, - list(transposed_theta.shape), transposed_theta.flatten().tolist()) - container.add_node('MatMul', [operator.input_full_names[0], theta_tensor], probability1_output, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('MatMul'), - op_version=1) - prior_tensor = scope.get_unique_variable_name('raw_predictions_tensor') - container.add_initializer(prior_tensor, onnx_proto.TensorProto.FLOAT, [len(op.pi)], op.pi) - container.add_node('Add', [probability1_output, prior_tensor], raw_predictions_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Add'), - op_version=7) - argmax_tensor = scope.get_unique_variable_name('argmax_tensor') - container.add_node('ArgMax', raw_predictions_tensor, argmax_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('ArgMax'), - op_version=1, - axis=1) - container.add_node('Cast', argmax_tensor, operator.output_full_names[0], - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Cast'), - op_version=9, - to=1) + probability1_output = scope.get_unique_variable_name("temp_probability") + theta_tensor = scope.get_unique_variable_name("theta_tensor") + container.add_initializer( + theta_tensor, + onnx_proto.TensorProto.FLOAT, + list(transposed_theta.shape), + transposed_theta.flatten().tolist(), + ) + container.add_node( + "MatMul", + [operator.input_full_names[0], theta_tensor], + probability1_output, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("MatMul"), + op_version=1, + ) + prior_tensor = scope.get_unique_variable_name("raw_predictions_tensor") + container.add_initializer( + prior_tensor, onnx_proto.TensorProto.FLOAT, [len(op.pi)], op.pi + ) + container.add_node( + "Add", + [probability1_output, prior_tensor], + raw_predictions_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Add"), + op_version=7, + ) + argmax_tensor = scope.get_unique_variable_name("argmax_tensor") + container.add_node( + "ArgMax", + raw_predictions_tensor, + argmax_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("ArgMax"), + op_version=1, + axis=1, + ) + container.add_node( + "Cast", + argmax_tensor, + operator.output_full_names[0], + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Cast"), + op_version=9, + to=1, + ) # Now we need to calculate Probabilities from rawPredictions # print('prediction:', numpy.argmax(result, 1)) # max_log = numpy.max(result, 1).reshape(result.shape[0], 1) - max_prediction_tensor = scope.get_unique_variable_name('max_prediction_tensor') - container.add_node('ReduceMax', raw_predictions_tensor, max_prediction_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('ReduceMax'), - op_version=1, - axes=[1], - keepdims=1) + max_prediction_tensor = scope.get_unique_variable_name("max_prediction_tensor") + container.add_node( + "ReduceMax", + raw_predictions_tensor, + max_prediction_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("ReduceMax"), + op_version=1, + axes=[1], + keepdims=1, + ) # sub_result = result - max_log - raw_minus_max_tensor = scope.get_unique_variable_name('raw_minus_max_tensor') - container.add_node('Sub', [raw_predictions_tensor, max_prediction_tensor], raw_minus_max_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Sub'), - op_version=7) + raw_minus_max_tensor = scope.get_unique_variable_name("raw_minus_max_tensor") + container.add_node( + "Sub", + [raw_predictions_tensor, max_prediction_tensor], + raw_minus_max_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Sub"), + op_version=7, + ) # exp_result = numpy.exp(sub_result) - exp_tensor = scope.get_unique_variable_name('exp_tensor') - container.add_node('Exp', raw_minus_max_tensor, exp_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Exp'), - op_version=6) + exp_tensor = scope.get_unique_variable_name("exp_tensor") + container.add_node( + "Exp", + raw_minus_max_tensor, + exp_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Exp"), + op_version=6, + ) # sum_log = numpy.sum(exp_result, 1).reshape(result.shape[0], 1) - sum_prediction_tensor = scope.get_unique_variable_name('sum_prediction_tensor') - container.add_node('ReduceSum', exp_tensor, sum_prediction_tensor, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('ReduceSum'), - op_version=1, - axes=[1], - keepdims=1) + sum_prediction_tensor = scope.get_unique_variable_name("sum_prediction_tensor") + container.add_node( + "ReduceSum", + exp_tensor, + sum_prediction_tensor, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("ReduceSum"), + op_version=1, + axes=[1], + keepdims=1, + ) # probabilities = exp_result / sum_log - container.add_node('Div', [exp_tensor, sum_prediction_tensor], operator.output_full_names[1], - op_domain='ai.onnx', - name=scope.get_unique_operator_name('Div'), - op_version=7) + container.add_node( + "Div", + [exp_tensor, sum_prediction_tensor], + operator.output_full_names[1], + op_domain="ai.onnx", + name=scope.get_unique_operator_name("Div"), + op_version=7, + ) -register_converter('pyspark.ml.classification.NaiveBayesModel', convert_sparkml_naive_bayes) +register_converter( + "pyspark.ml.classification.NaiveBayesModel", convert_sparkml_naive_bayes +) def calculate_sparkml_naive_bayes_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=2) - check_input_and_output_types(operator, - good_input_types=[FloatTensorType], - good_output_types=[FloatTensorType,FloatTensorType]) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType], + good_output_types=[FloatTensorType, FloatTensorType], + ) N = operator.inputs[0].type.shape[0] C = operator.raw_operator.numClasses operator.outputs[0].type = FloatTensorType([N, 1]) operator.outputs[1].type = FloatTensorType([N, C]) -register_shape_calculator('pyspark.ml.classification.NaiveBayesModel', calculate_sparkml_naive_bayes_output_shapes) +register_shape_calculator( + "pyspark.ml.classification.NaiveBayesModel", + calculate_sparkml_naive_bayes_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/normalizer.py b/onnxmltools/convert/sparkml/operator_converters/normalizer.py index 88a350fc..8c6eca01 100644 --- a/onnxmltools/convert/sparkml/operator_converters/normalizer.py +++ b/onnxmltools/convert/sparkml/operator_converters/normalizer.py @@ -4,7 +4,8 @@ from ...common.data_types import Int64TensorType, FloatTensorType from ...common.utils import check_input_and_output_numbers, check_input_and_output_types -#from ..convert.sparkml import SparkMLConversionError + +# from ..convert.sparkml import SparkMLConversionError from ...common._registration import register_converter, register_shape_calculator @@ -12,28 +13,34 @@ def convert_sparkml_normalizer(scope, operator, container): op = operator.raw_operator input_name = op.getInputCol() - op_type = 'Normalizer' + op_type = "Normalizer" name = scope.get_unique_operator_name(op_type) - norm = 'L1' + norm = "L1" p = op.getP() if int(p) == 2: - norm = 'L2' + norm = "L2" elif int(p) != 1: raise ValueError("Unsupported Norm value: " + p) - attrs = {'name': name, 'norm': norm} - container.add_node(op_type, input_name, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + attrs = {"name": name, "norm": norm} + container.add_node( + op_type, input_name, operator.output_full_names, op_domain="ai.onnx.ml", **attrs + ) -register_converter('pyspark.ml.feature.Normalizer', convert_sparkml_normalizer) +register_converter("pyspark.ml.feature.Normalizer", convert_sparkml_normalizer) def calculate_sparkml_normalizer_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) - check_input_and_output_types(operator, - good_input_types=[FloatTensorType, Int64TensorType], - good_output_types=[FloatTensorType]) + check_input_and_output_types( + operator, + good_input_types=[FloatTensorType, Int64TensorType], + good_output_types=[FloatTensorType], + ) input_shape = copy.deepcopy(operator.inputs[0].type.shape) operator.outputs[0].type = FloatTensorType(input_shape) -register_shape_calculator('pyspark.ml.feature.Normalizer', calculate_sparkml_normalizer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.Normalizer", calculate_sparkml_normalizer_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/one_vs_rest.py b/onnxmltools/convert/sparkml/operator_converters/one_vs_rest.py index 5b7b0055..8d3e74a5 100644 --- a/onnxmltools/convert/sparkml/operator_converters/one_vs_rest.py +++ b/onnxmltools/convert/sparkml/operator_converters/one_vs_rest.py @@ -13,46 +13,68 @@ def convert_one_vs_rest(scope, operator, container): op = operator.raw_operator classifier_output_names = [] # initializer needed to extract the 2nd value of probability array - index_tensor = scope.get_unique_variable_name('index_tensor') + index_tensor = scope.get_unique_variable_name("index_tensor") container.add_initializer(index_tensor, onnx_proto.TensorProto.INT64, [1], [1]) # OneVsRest could have different classifiers # all must have at least Probability values for us to do the argmax for sub_model in op.models: - classifier_op = scope.declare_local_operator(get_sparkml_operator_name(type(sub_model)), sub_model) + classifier_op = scope.declare_local_operator( + get_sparkml_operator_name(type(sub_model)), sub_model + ) classifier_op.raw_params = operator.raw_params classifier_op.inputs = operator.inputs - classifier_prediction_output = scope.declare_local_variable('classifier_prediction', FloatTensorType()) - classifier_probability_output = scope.declare_local_variable('classifier_probability', FloatTensorType()) + classifier_prediction_output = scope.declare_local_variable( + "classifier_prediction", FloatTensorType() + ) + classifier_probability_output = scope.declare_local_variable( + "classifier_probability", FloatTensorType() + ) classifier_op.outputs.append(classifier_prediction_output) classifier_op.outputs.append(classifier_probability_output) convert_sparkml_linear_classifier(scope, classifier_op, container) classifier_op.is_evaluated = True - single_feature_tensor = scope.get_unique_variable_name('single_feature_tensor') - container.add_node('ArrayFeatureExtractor', [classifier_probability_output.full_name, index_tensor], - single_feature_tensor, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('ArrayFeatureExtractor')) + single_feature_tensor = scope.get_unique_variable_name("single_feature_tensor") + container.add_node( + "ArrayFeatureExtractor", + [classifier_probability_output.full_name, index_tensor], + single_feature_tensor, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("ArrayFeatureExtractor"), + ) classifier_output_names.append(single_feature_tensor) - concatenated_probabilities = scope.get_unique_variable_name('concatenated_predictions_tensor') - apply_concat(scope, classifier_output_names, concatenated_probabilities, container, axis=1) + concatenated_probabilities = scope.get_unique_variable_name( + "concatenated_predictions_tensor" + ) + apply_concat( + scope, classifier_output_names, concatenated_probabilities, container, axis=1 + ) # to get Prediction from probability - apply_argmax(scope, concatenated_probabilities, operator.outputs[0].full_name, container, - axis=1, keepdims=1) + apply_argmax( + scope, + concatenated_probabilities, + operator.outputs[0].full_name, + container, + axis=1, + keepdims=1, + ) -register_converter('pyspark.ml.classification.OneVsRestModel', convert_one_vs_rest) +register_converter("pyspark.ml.classification.OneVsRestModel", convert_one_vs_rest) def calculate_one_vs_rest_output_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) if len(operator.inputs[0].type.shape) != 2: - raise RuntimeError('Input must be a [N, C]-tensor') + raise RuntimeError("Input must be a [N, C]-tensor") N = operator.inputs[0].type.shape[0] operator.outputs[0].type = Int64TensorType(shape=[N]) -register_shape_calculator('pyspark.ml.classification.OneVsRestModel', - calculate_one_vs_rest_output_shapes) +register_shape_calculator( + "pyspark.ml.classification.OneVsRestModel", calculate_one_vs_rest_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/onehot_encoder.py b/onnxmltools/convert/sparkml/operator_converters/onehot_encoder.py index 880e9075..d36920b0 100644 --- a/onnxmltools/convert/sparkml/operator_converters/onehot_encoder.py +++ b/onnxmltools/convert/sparkml/operator_converters/onehot_encoder.py @@ -7,26 +7,35 @@ from pyspark.ml.feature import OneHotEncoderModel from typing import List + def _get_categories(operator: Operator) -> List[List[int]]: op: OneHotEncoderModel = operator.raw_operator categorySizes: List[int] = op.categorySizes # This is necessary to match SparkML's OneHotEncoder behavior. - # If handleInvalid is set to "keep", an extra "category" indicating invalid values is added as last category: - # - if dropLast is set to false, then an extra bit will be added for invalid input, which does not match ONNX's OneHotEncoder behavior. - # - if dropLast is set to true, then invalid values are encoded as all-zeros vector, which matches ONNX's OneHotEncoder behavior when "zeros" is set to 1. + # If handleInvalid is set to "keep", an extra "category" indicating + # invalid values is added as last category: + # - if dropLast is set to false, then an extra bit will be added for invalid input, + # which does not match ONNX's OneHotEncoder behavior. + # - if dropLast is set to true, then invalid values are encoded as all-zeros vector, + # which matches ONNX's OneHotEncoder behavior when "zeros" is set to 1. # If handleInvalid is set to "error", then: - # - if dropLast is set to false, then nothing will be added or dropped, and matches ONNX's OneHotEncoder behavior when "zeros" is set to 0. - # - if dropLast is set to true, then the last bit will be dropped, which does not match ONNX's OneHotEncoder behavior. - if (op.getHandleInvalid() == "keep" and op.getDropLast() == False) or ( - op.getHandleInvalid() == "error" and op.getDropLast() == True + # - if dropLast is set to false, then nothing will be added or dropped, + # and matches ONNX's OneHotEncoder behavior when "zeros" is set to 0. + # - if dropLast is set to true, then the last bit will be dropped, + # which does not match ONNX's OneHotEncoder behavior. + if (op.getHandleInvalid() == "keep" and not op.getDropLast()) or ( + op.getHandleInvalid() == "error" and op.getDropLast() ): raise RuntimeError( - f"The 'handleInvalid' and 'dropLast' parameters must be set to ('keep', True) or ('error', False), but got ('{op.getHandleInvalid()}', {op.getDropLast()}) instead." + f"The 'handleInvalid' and 'dropLast' parameters must be set to " + f"('keep', True) or ('error', False), but got " + f"('{op.getHandleInvalid()}', {op.getDropLast()}) instead." ) return [list(range(0, size)) for size in categorySizes] + def convert_sparkml_one_hot_encoder(scope: Scope, operator: Operator, container): categories = _get_categories(operator) N = operator.inputs[0].type.shape[0] or -1 @@ -36,14 +45,17 @@ def convert_sparkml_one_hot_encoder(scope: Scope, operator: Operator, container) for i in range(0, len(categories)): encoder_type = "OneHotEncoder" - # Set zeros to 0 to match the "error" handleInvalid behavior of SparkML's OneHotEncoder. + # Set zeros to 0 to match the "error" + # handleInvalid behavior of SparkML's OneHotEncoder. encoder_attrs = { "name": scope.get_unique_operator_name(encoder_type), "cats_int64s": categories[i], "zeros": zeros, } - encoded_feature_name = scope.get_unique_variable_name("encoded_feature_at_" + str(i)) + encoded_feature_name = scope.get_unique_variable_name( + "encoded_feature_at_" + str(i) + ) container.add_node( op_type=encoder_type, inputs=[operator.inputs[i].full_name], @@ -54,7 +66,12 @@ def convert_sparkml_one_hot_encoder(scope: Scope, operator: Operator, container) ) shape_variable_name = scope.get_unique_variable_name("shape_at_" + str(i)) - container.add_initializer(shape_variable_name, onnx_proto.TensorProto.INT64, [2], [N, len(categories[i])]) + container.add_initializer( + shape_variable_name, + onnx_proto.TensorProto.INT64, + [2], + [N, len(categories[i])], + ) container.add_node( op_type="Reshape", @@ -65,7 +82,9 @@ def convert_sparkml_one_hot_encoder(scope: Scope, operator: Operator, container) ) -register_converter('pyspark.ml.feature.OneHotEncoderModel', convert_sparkml_one_hot_encoder) +register_converter( + "pyspark.ml.feature.OneHotEncoderModel", convert_sparkml_one_hot_encoder +) def calculate_sparkml_one_hot_encoder_output_shapes(operator: Operator): @@ -75,5 +94,7 @@ def calculate_sparkml_one_hot_encoder_output_shapes(operator: Operator): output.type = FloatTensorType([N, len(categories[i])]) -register_shape_calculator('pyspark.ml.feature.OneHotEncoderModel', calculate_sparkml_one_hot_encoder_output_shapes) - +register_shape_calculator( + "pyspark.ml.feature.OneHotEncoderModel", + calculate_sparkml_one_hot_encoder_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/pca.py b/onnxmltools/convert/sparkml/operator_converters/pca.py index 5b606009..2f3539d8 100644 --- a/onnxmltools/convert/sparkml/operator_converters/pca.py +++ b/onnxmltools/convert/sparkml/operator_converters/pca.py @@ -9,20 +9,33 @@ def convert_sparkml_pca(scope, operator, container): op = operator.raw_operator - pc_tensor = scope.get_unique_variable_name('pc_tensor') - container.add_initializer(pc_tensor, onnx_proto.TensorProto.FLOAT, - [op.pc.numRows, op.pc.numCols], list(op.pc.toArray().flatten())) - apply_matmul(scope, [operator.inputs[0].full_name, pc_tensor], operator.output_full_names, container) + pc_tensor = scope.get_unique_variable_name("pc_tensor") + container.add_initializer( + pc_tensor, + onnx_proto.TensorProto.FLOAT, + [op.pc.numRows, op.pc.numCols], + list(op.pc.toArray().flatten()), + ) + apply_matmul( + scope, + [operator.inputs[0].full_name, pc_tensor], + operator.output_full_names, + container, + ) -register_converter('pyspark.ml.feature.PCAModel', convert_sparkml_pca) +register_converter("pyspark.ml.feature.PCAModel", convert_sparkml_pca) def calculate_sparkml_pca_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) check_input_and_output_types(operator, good_input_types=[FloatTensorType]) N = operator.inputs[0].type.shape[0] - operator.outputs[0].type = FloatTensorType([N, operator.raw_operator.getOrDefault('k')]) + operator.outputs[0].type = FloatTensorType( + [N, operator.raw_operator.getOrDefault("k")] + ) -register_shape_calculator('pyspark.ml.feature.PCAModel', calculate_sparkml_pca_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.PCAModel", calculate_sparkml_pca_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/polynomial_expansion.py b/onnxmltools/convert/sparkml/operator_converters/polynomial_expansion.py index 36ddca1d..4f422192 100644 --- a/onnxmltools/convert/sparkml/operator_converters/polynomial_expansion.py +++ b/onnxmltools/convert/sparkml/operator_converters/polynomial_expansion.py @@ -19,56 +19,78 @@ def convert_sparkml_polynomial_expansion(scope, operator, container): if comb is None: pass else: - comb_name = scope.get_unique_variable_name('comb') - col_name = scope.get_unique_variable_name('col') - prod_name = scope.get_unique_variable_name('prod') + comb_name = scope.get_unique_variable_name("comb") + col_name = scope.get_unique_variable_name("col") + prod_name = scope.get_unique_variable_name("prod") - container.add_initializer(comb_name, onnx_proto.TensorProto.INT64, - [len(comb)], list(comb)) + container.add_initializer( + comb_name, onnx_proto.TensorProto.INT64, [len(comb)], list(comb) + ) container.add_node( - 'ArrayFeatureExtractor', - [operator.inputs[0].full_name, comb_name], col_name, - name=scope.get_unique_operator_name('ArrayFeatureExtractor'), - op_domain='ai.onnx.ml') + "ArrayFeatureExtractor", + [operator.inputs[0].full_name, comb_name], + col_name, + name=scope.get_unique_operator_name("ArrayFeatureExtractor"), + op_domain="ai.onnx.ml", + ) reduce_prod_input = col_name - if (operator.inputs[0].type._get_element_onnx_type() - == onnx_proto.TensorProto.INT64): - float_col_name = scope.get_unique_variable_name('col') - container.add_node('Cast', col_name, float_col_name, - name=scope.get_unique_operator_name('Cast'), - to=onnx_proto.TensorProto.FLOAT) + if ( + operator.inputs[0].type._get_element_onnx_type() + == onnx_proto.TensorProto.INT64 + ): + float_col_name = scope.get_unique_variable_name("col") + container.add_node( + "Cast", + col_name, + float_col_name, + name=scope.get_unique_operator_name("Cast"), + to=onnx_proto.TensorProto.FLOAT, + ) reduce_prod_input = float_col_name container.add_node( - 'ReduceProd', reduce_prod_input, prod_name, - axes=[1], name=scope.get_unique_operator_name('ReduceProd')) + "ReduceProd", + reduce_prod_input, + prod_name, + axes=[1], + name=scope.get_unique_operator_name("ReduceProd"), + ) transformed_columns.append(prod_name) - if (operator.inputs[0].type._get_element_onnx_type() - == onnx_proto.TensorProto.INT64): - concat_result_name = scope.get_unique_variable_name('concat_result') - - apply_concat(scope, [t for t in transformed_columns], concat_result_name, - container, axis=1) - apply_cast(scope, concat_result_name, operator.outputs[0].full_name, - container, - to=onnx_proto.TensorProto.INT64) + if operator.inputs[0].type._get_element_onnx_type() == onnx_proto.TensorProto.INT64: + concat_result_name = scope.get_unique_variable_name("concat_result") + + apply_concat( + scope, + [t for t in transformed_columns], + concat_result_name, + container, + axis=1, + ) + apply_cast( + scope, + concat_result_name, + operator.outputs[0].full_name, + container, + to=onnx_proto.TensorProto.INT64, + ) else: - apply_concat(scope, [t for t in transformed_columns], operator.outputs[0].full_name, - container, axis=1) + apply_concat( + scope, + [t for t in transformed_columns], + operator.outputs[0].full_name, + container, + axis=1, + ) -register_converter('pyspark.ml.feature.PolynomialExpansion', convert_sparkml_polynomial_expansion) +register_converter( + "pyspark.ml.feature.PolynomialExpansion", convert_sparkml_polynomial_expansion +) -def expand_inner( - values, - last_index, - degree, - multiplier, - poly_values, - cur_poly_index): +def expand_inner(values, last_index, degree, multiplier, poly_values, cur_poly_index): if degree == 0 or last_index < 0: if cur_poly_index >= 0: # skip the very first 1 @@ -80,7 +102,9 @@ def expand_inner( i = 0 cur_start = cur_poly_index while i <= degree: - cur_start, multiplier = expand_inner(values, last_index1, degree - i, alpha, poly_values, cur_start) + cur_start, multiplier = expand_inner( + values, last_index1, degree - i, alpha, poly_values, cur_start + ) i += 1 alpha.append(v) return cur_poly_index + get_poly_size(last_index + 1, degree), multiplier @@ -95,11 +119,12 @@ def calc_combinations(feature_count, degree): def get_poly_size(feature_count, degree): - return get_combinations_count(feature_count+degree, degree) + return get_combinations_count(feature_count + degree, degree) def get_combinations_count(n, k): from math import factorial + if n == k or k == 0: return 1 if k == 1 or k == n - 1: @@ -109,13 +134,19 @@ def get_combinations_count(n, k): def calculate_sparkml_polynomial_expansion_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[ - FloatTensorType, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) N = operator.inputs[0].type.shape[0] - C = get_combinations_count(operator.inputs[0].type.shape[1], operator.raw_operator.getDegree()) + C = get_combinations_count( + operator.inputs[0].type.shape[1], operator.raw_operator.getDegree() + ) operator.outputs[0].type = copy.deepcopy(operator.inputs[0].type) operator.outputs[0].type.shape = [N, C] -register_shape_calculator('pyspark.ml.feature.PolynomialExpansion', calculate_sparkml_polynomial_expansion_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.PolynomialExpansion", + calculate_sparkml_polynomial_expansion_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/random_forest_classifier.py b/onnxmltools/convert/sparkml/operator_converters/random_forest_classifier.py index 02fb35b5..3a808c1c 100644 --- a/onnxmltools/convert/sparkml/operator_converters/random_forest_classifier.py +++ b/onnxmltools/convert/sparkml/operator_converters/random_forest_classifier.py @@ -1,10 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 import logging -from ...common.tree_ensemble import get_default_tree_classifier_attribute_pairs, \ - add_tree_to_attribute_pairs +from ...common.tree_ensemble import ( + get_default_tree_classifier_attribute_pairs, + add_tree_to_attribute_pairs, +) from ...common._registration import register_converter, register_shape_calculator -from .tree_ensemble_common import save_read_sparkml_model_data, sparkml_tree_dataset_to_sklearn +from .tree_ensemble_common import ( + save_read_sparkml_model_data, + sparkml_tree_dataset_to_sklearn, +) from .decision_tree_classifier import calculate_decision_tree_classifier_output_shapes from .tree_helper import rewrite_ids_and_process @@ -13,38 +18,49 @@ def convert_random_forest_classifier(scope, operator, container): op = operator.raw_operator - op_type = 'TreeEnsembleClassifier' + op_type = "TreeEnsembleClassifier" main_attr_pairs = get_default_tree_classifier_attribute_pairs() - main_attr_pairs['name'] = scope.get_unique_operator_name(op_type) - main_attr_pairs['classlabels_int64s'] = list(range(0, op.numClasses)) + main_attr_pairs["name"] = scope.get_unique_operator_name(op_type) + main_attr_pairs["classlabels_int64s"] = list(range(0, op.numClasses)) # random forest calculate the final score by averaging over all trees' # outcomes, so all trees' weights are identical. - tree_weight = 1. / op.getNumTrees + tree_weight = 1.0 / op.getNumTrees for tree_id in range(0, op.getNumTrees): tree_model = op.trees[tree_id] - tree_df = save_read_sparkml_model_data(operator.raw_params['SparkSession'], tree_model) + tree_df = save_read_sparkml_model_data( + operator.raw_params["SparkSession"], tree_model + ) tree = sparkml_tree_dataset_to_sklearn(tree_df, is_classifier=True) attr_pairs = get_default_tree_classifier_attribute_pairs() - attr_pairs['name'] = scope.get_unique_operator_name(op_type) - attr_pairs['classlabels_int64s'] = list(range(0, op.numClasses)) + attr_pairs["name"] = scope.get_unique_operator_name(op_type) + attr_pairs["classlabels_int64s"] = list(range(0, op.numClasses)) - add_tree_to_attribute_pairs(attr_pairs, True, tree, tree_id, - tree_weight, 0, True) + add_tree_to_attribute_pairs( + attr_pairs, True, tree, tree_id, tree_weight, 0, True + ) new_attrs = rewrite_ids_and_process(attr_pairs, logger) for k, v in new_attrs.items(): - if isinstance(v, list) and k not in {'classlabels_int64s'}: + if isinstance(v, list) and k not in {"classlabels_int64s"}: main_attr_pairs[k].extend(v) container.add_node( - op_type, operator.input_full_names, + op_type, + operator.input_full_names, [operator.outputs[0].full_name, operator.outputs[1].full_name], - op_domain='ai.onnx.ml', **main_attr_pairs) + op_domain="ai.onnx.ml", + **main_attr_pairs + ) -register_converter('pyspark.ml.classification.RandomForestClassificationModel', convert_random_forest_classifier) +register_converter( + "pyspark.ml.classification.RandomForestClassificationModel", + convert_random_forest_classifier, +) -register_shape_calculator('pyspark.ml.classification.RandomForestClassificationModel', - calculate_decision_tree_classifier_output_shapes) +register_shape_calculator( + "pyspark.ml.classification.RandomForestClassificationModel", + calculate_decision_tree_classifier_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/random_forest_regressor.py b/onnxmltools/convert/sparkml/operator_converters/random_forest_regressor.py index db0cb23c..cd367d95 100644 --- a/onnxmltools/convert/sparkml/operator_converters/random_forest_regressor.py +++ b/onnxmltools/convert/sparkml/operator_converters/random_forest_regressor.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 import logging -from ...common.tree_ensemble import add_tree_to_attribute_pairs, \ - get_default_tree_regressor_attribute_pairs +from ...common.tree_ensemble import ( + add_tree_to_attribute_pairs, + get_default_tree_regressor_attribute_pairs, +) from ...common._registration import register_converter, register_shape_calculator from .decision_tree_classifier import save_read_sparkml_model_data from .decision_tree_regressor import calculate_decision_tree_regressor_output_shapes @@ -14,36 +16,46 @@ def convert_random_forest_regressor(scope, operator, container): op = operator.raw_operator - op_type = 'TreeEnsembleRegressor' + op_type = "TreeEnsembleRegressor" main_attrs = get_default_tree_regressor_attribute_pairs() - main_attrs['name'] = scope.get_unique_operator_name(op_type) - main_attrs['n_targets'] = 1 + main_attrs["name"] = scope.get_unique_operator_name(op_type) + main_attrs["n_targets"] = 1 # random forest calculate the final score by averaging over all trees' # outcomes, so all trees' weights are identical. - tree_weight = 1. / op.getNumTrees + tree_weight = 1.0 / op.getNumTrees for tree_id in range(0, op.getNumTrees): tree_model = op.trees[tree_id] - tree_df = save_read_sparkml_model_data(operator.raw_params['SparkSession'], tree_model) + tree_df = save_read_sparkml_model_data( + operator.raw_params["SparkSession"], tree_model + ) tree = sparkml_tree_dataset_to_sklearn(tree_df, is_classifier=False) attrs = get_default_tree_regressor_attribute_pairs() - attrs['name'] = scope.get_unique_operator_name(op_type) - attrs['n_targets'] = 1 - add_tree_to_attribute_pairs(attrs, False, tree, tree_id, - tree_weight, 0, False) + attrs["name"] = scope.get_unique_operator_name(op_type) + attrs["n_targets"] = 1 + add_tree_to_attribute_pairs(attrs, False, tree, tree_id, tree_weight, 0, False) new_attrs = rewrite_ids_and_process(attrs, logger) for k, v in new_attrs.items(): if isinstance(v, list): main_attrs[k].extend(v) - container.add_node(op_type, operator.input_full_names, operator.output_full_names[0], - op_domain='ai.onnx.ml', **main_attrs) + container.add_node( + op_type, + operator.input_full_names, + operator.output_full_names[0], + op_domain="ai.onnx.ml", + **main_attrs + ) -register_converter('pyspark.ml.regression.RandomForestRegressionModel', convert_random_forest_regressor) +register_converter( + "pyspark.ml.regression.RandomForestRegressionModel", convert_random_forest_regressor +) -register_shape_calculator('pyspark.ml.regression.RandomForestRegressionModel', - calculate_decision_tree_regressor_output_shapes) +register_shape_calculator( + "pyspark.ml.regression.RandomForestRegressionModel", + calculate_decision_tree_regressor_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/scaler.py b/onnxmltools/convert/sparkml/operator_converters/scaler.py index 93a645b5..ac7f649c 100644 --- a/onnxmltools/convert/sparkml/operator_converters/scaler.py +++ b/onnxmltools/convert/sparkml/operator_converters/scaler.py @@ -14,45 +14,61 @@ def convert_sparkml_scaler(scope, operator, container): op = operator.raw_operator input_name = operator.inputs[0].full_name - op_type = 'Scaler' - attrs = {'name': scope.get_unique_operator_name(op_type)} + op_type = "Scaler" + attrs = {"name": scope.get_unique_operator_name(op_type)} if isinstance(op, StandardScalerModel): C = operator.inputs[0].type.shape[1] - attrs['offset'] = op.mean.toArray() if op.getOrDefault("withMean") else [0.0] * C - attrs['scale'] = [1.0 / x for x in op.std.toArray()] if op.getOrDefault("withStd") else [1.0] * C + attrs["offset"] = ( + op.mean.toArray() if op.getOrDefault("withMean") else [0.0] * C + ) + attrs["scale"] = ( + [1.0 / x for x in op.std.toArray()] + if op.getOrDefault("withStd") + else [1.0] * C + ) elif isinstance(op, MinMaxScalerModel): epsilon = 1.0e-8 # to avoid dividing by zero - attrs['offset'] = [x for x in op.originalMin] + attrs["offset"] = [x for x in op.originalMin] max_min = [a - b for a, b in zip(op.originalMax, op.originalMin)] - attrs['scale'] = [1.0 / (x + epsilon) for x in max_min] + attrs["scale"] = [1.0 / (x + epsilon) for x in max_min] elif isinstance(op, MaxAbsScalerModel): C = operator.inputs[0].type.shape[1] - attrs['offset'] = [0.] * C - attrs['scale'] = [1.0 / x for x in op.maxAbs] + attrs["offset"] = [0.0] * C + attrs["scale"] = [1.0 / x for x in op.maxAbs] else: - raise ValueError('Unsupported Scaler: %s' % type(op)) + raise ValueError("Unsupported Scaler: %s" % type(op)) # ONNX does not convert arrays of float32. for k in attrs: v = attrs[k] if isinstance(v, numpy.ndarray) and v.dtype == numpy.float32: attrs[k] = v.astype(numpy.float64) - container.add_node(op_type, input_name, operator.output_full_names, op_domain='ai.onnx.ml', **attrs) + container.add_node( + op_type, input_name, operator.output_full_names, op_domain="ai.onnx.ml", **attrs + ) -register_converter('pyspark.ml.feature.StandardScalerModel', convert_sparkml_scaler) -register_converter('pyspark.ml.feature.MaxAbsScalerModel', convert_sparkml_scaler) -register_converter('pyspark.ml.feature.MinMaxScalerModel', convert_sparkml_scaler) +register_converter("pyspark.ml.feature.StandardScalerModel", convert_sparkml_scaler) +register_converter("pyspark.ml.feature.MaxAbsScalerModel", convert_sparkml_scaler) +register_converter("pyspark.ml.feature.MinMaxScalerModel", convert_sparkml_scaler) def calculate_sparkml_scaler_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) input_shape = copy.deepcopy(operator.inputs[0].type.shape) operator.outputs[0].type = FloatTensorType(input_shape) -register_shape_calculator('pyspark.ml.feature.StandardScalerModel', calculate_sparkml_scaler_output_shapes) -register_shape_calculator('pyspark.ml.feature.MaxAbsScalerModel', calculate_sparkml_scaler_output_shapes) -register_shape_calculator('pyspark.ml.feature.MinMaxScalerModel', calculate_sparkml_scaler_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.StandardScalerModel", calculate_sparkml_scaler_output_shapes +) +register_shape_calculator( + "pyspark.ml.feature.MaxAbsScalerModel", calculate_sparkml_scaler_output_shapes +) +register_shape_calculator( + "pyspark.ml.feature.MinMaxScalerModel", calculate_sparkml_scaler_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/stop_words_remover.py b/onnxmltools/convert/sparkml/operator_converters/stop_words_remover.py index dd8e6aed..9cc8bdb4 100644 --- a/onnxmltools/convert/sparkml/operator_converters/stop_words_remover.py +++ b/onnxmltools/convert/sparkml/operator_converters/stop_words_remover.py @@ -10,24 +10,32 @@ def convert_sparkml_stop_words_remover(scope, operator, container): op = operator.raw_operator - container.add_node('StringNormalizer', operator.input_full_names[0], operator.output_full_names, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('StringNormalizer'), - op_version=10, - case_change_action='NONE', - is_case_sensitive=1 if op.getCaseSensitive() else 0, - stopwords=op.getStopWords()) + container.add_node( + "StringNormalizer", + operator.input_full_names[0], + operator.output_full_names, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("StringNormalizer"), + op_version=10, + case_change_action="NONE", + is_case_sensitive=1 if op.getCaseSensitive() else 0, + stopwords=op.getStopWords(), + ) -register_converter('pyspark.ml.feature.StopWordsRemover', convert_sparkml_stop_words_remover) +register_converter( + "pyspark.ml.feature.StopWordsRemover", convert_sparkml_stop_words_remover +) def calculate_sparkml_stop_words_remover_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) - check_input_and_output_types(operator, - good_input_types=[StringTensorType]) + check_input_and_output_types(operator, good_input_types=[StringTensorType]) input_shape = copy.deepcopy(operator.inputs[0].type.shape) operator.outputs[0].type = StringTensorType(input_shape) -register_shape_calculator('pyspark.ml.feature.StopWordsRemover', calculate_sparkml_stop_words_remover_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.StopWordsRemover", + calculate_sparkml_stop_words_remover_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/string_indexer.py b/onnxmltools/convert/sparkml/operator_converters/string_indexer.py index 4b629a25..4810ffbb 100644 --- a/onnxmltools/convert/sparkml/operator_converters/string_indexer.py +++ b/onnxmltools/convert/sparkml/operator_converters/string_indexer.py @@ -12,7 +12,9 @@ from ...common.utils import check_input_and_output_types -def convert_sparkml_string_indexer(scope: Scope, operator: Operator, container: ModelComponentContainer): +def convert_sparkml_string_indexer( + scope: Scope, operator: Operator, container: ModelComponentContainer +): op: StringIndexerModel = operator.raw_operator op_domain = "ai.onnx.ml" op_version = 2 @@ -42,20 +44,28 @@ def convert_sparkml_string_indexer(scope: Scope, operator: Operator, container: ) -register_converter("pyspark.ml.feature.StringIndexerModel", convert_sparkml_string_indexer) +register_converter( + "pyspark.ml.feature.StringIndexerModel", convert_sparkml_string_indexer +) def calculate_sparkml_string_indexer_output_shapes(operator: Operator): """ - This function just copy the input shape to the output because label encoder only alters input features' values, not + This function just copy the input shape to the output + because label encoder only alters input features' values, not their shape. """ - check_input_and_output_types(operator, good_input_types=[Int64TensorType, StringTensorType]) + check_input_and_output_types( + operator, good_input_types=[Int64TensorType, StringTensorType] + ) input: Variable output: Variable - for (input, output) in zip(operator.inputs, operator.outputs): + for input, output in zip(operator.inputs, operator.outputs): input_shape = copy.deepcopy(input.type.shape) output.type = Int64TensorType(input_shape) -register_shape_calculator("pyspark.ml.feature.StringIndexerModel", calculate_sparkml_string_indexer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.StringIndexerModel", + calculate_sparkml_string_indexer_output_shapes, +) diff --git a/onnxmltools/convert/sparkml/operator_converters/tokenizer.py b/onnxmltools/convert/sparkml/operator_converters/tokenizer.py index d27ec200..6b58c9f0 100644 --- a/onnxmltools/convert/sparkml/operator_converters/tokenizer.py +++ b/onnxmltools/convert/sparkml/operator_converters/tokenizer.py @@ -11,30 +11,39 @@ def convert_sparkml_tokenizer(scope, operator, container): raise RuntimeError("op cannot be None.") # the SPARK version converts text to lowercase and applies "\\s" regexp to it # Here we'll tokenize and then normalize (to convert to lowercase) - lowercase_output = scope.get_unique_variable_name('lowercase_tensor') - container.add_node('StringNormalizer', operator.input_full_names[0], lowercase_output, - op_domain='ai.onnx', - name=scope.get_unique_operator_name('StringNormalizer'), - op_version=10, - case_change_action='LOWER') - container.add_node('Tokenizer', lowercase_output, operator.output_full_names, - op_domain='com.microsoft', - name=scope.get_unique_operator_name('Tokenizer'), - op_version=1, - mark=0, - separators=[' ', '\t', '\r', '\n'], - pad_value='##ERROR##', - mincharnum=1) - - -register_converter('pyspark.ml.feature.Tokenizer', convert_sparkml_tokenizer) + lowercase_output = scope.get_unique_variable_name("lowercase_tensor") + container.add_node( + "StringNormalizer", + operator.input_full_names[0], + lowercase_output, + op_domain="ai.onnx", + name=scope.get_unique_operator_name("StringNormalizer"), + op_version=10, + case_change_action="LOWER", + ) + container.add_node( + "Tokenizer", + lowercase_output, + operator.output_full_names, + op_domain="com.microsoft", + name=scope.get_unique_operator_name("Tokenizer"), + op_version=1, + mark=0, + separators=[" ", "\t", "\r", "\n"], + pad_value="##ERROR##", + mincharnum=1, + ) + + +register_converter("pyspark.ml.feature.Tokenizer", convert_sparkml_tokenizer) def calculate_sparkml_tokenizer_output_shapes(operator): check_input_and_output_numbers(operator, output_count_range=1) - check_input_and_output_types(operator, - good_input_types=[StringTensorType]) + check_input_and_output_types(operator, good_input_types=[StringTensorType]) operator.outputs[0].type = StringTensorType() -register_shape_calculator('pyspark.ml.feature.Tokenizer', calculate_sparkml_tokenizer_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.Tokenizer", calculate_sparkml_tokenizer_output_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py b/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py index 9df6d983..9d95cb42 100644 --- a/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py +++ b/onnxmltools/convert/sparkml/operator_converters/tree_ensemble_common.py @@ -6,6 +6,7 @@ import numpy from pyspark.sql import SparkSession + class SparkMLTree(dict): pass @@ -27,10 +28,10 @@ def sparkml_tree_dataset_to_sklearn(tree_df, is_classifier): try: feature.append(item["featureIndex"]) threshold.append(item["leftCategoriesOrThreshold"]) - except KeyError as e: + except KeyError: raise RuntimeError(f"Unable to process {item}.") else: - tuple_item = tuple(item) + tuple(item) feature.append(item[0]) threshold.append(item[1][0] if len(item[1]) >= 1 else -1.0) @@ -48,115 +49,155 @@ def sparkml_tree_dataset_to_sklearn(tree_df, is_classifier): def save_read_sparkml_model_data(spark: SparkSession, model): tdir = tempfile.tempdir if tdir is None: - local_dir = spark._jvm.org.apache.spark.util.Utils.getLocalDir(spark._jsc.sc().conf()) - tdir = spark._jvm.org.apache.spark.util.Utils.createTempDir(local_dir, "onnx").getAbsolutePath() + local_dir = spark._jvm.org.apache.spark.util.Utils.getLocalDir( + spark._jsc.sc().conf() + ) + tdir = spark._jvm.org.apache.spark.util.Utils.createTempDir( + local_dir, "onnx" + ).getAbsolutePath() if tdir is None: raise FileNotFoundError( "Unable to create a temporary directory for model '{}'" - ".".format(type(model).__name__)) + ".".format(type(model).__name__) + ) path = os.path.join(tdir, type(model).__name__ + "_" + str(time.time())) model.write().overwrite().save(path) - df = spark.read.parquet(os.path.join(path, 'data')) + df = spark.read.parquet(os.path.join(path, "data")) return df def get_default_tree_classifier_attribute_pairs(): attrs = {} - attrs['post_transform'] = 'NONE' - attrs['nodes_treeids'] = [] - attrs['nodes_nodeids'] = [] - attrs['nodes_featureids'] = [] - attrs['nodes_modes'] = [] - attrs['nodes_values'] = [] - attrs['nodes_truenodeids'] = [] - attrs['nodes_falsenodeids'] = [] - attrs['nodes_missing_value_tracks_true'] = [] - attrs['nodes_hitrates'] = [] - attrs['class_treeids'] = [] - attrs['class_nodeids'] = [] - attrs['class_ids'] = [] - attrs['class_weights'] = [] + attrs["post_transform"] = "NONE" + attrs["nodes_treeids"] = [] + attrs["nodes_nodeids"] = [] + attrs["nodes_featureids"] = [] + attrs["nodes_modes"] = [] + attrs["nodes_values"] = [] + attrs["nodes_truenodeids"] = [] + attrs["nodes_falsenodeids"] = [] + attrs["nodes_missing_value_tracks_true"] = [] + attrs["nodes_hitrates"] = [] + attrs["class_treeids"] = [] + attrs["class_nodeids"] = [] + attrs["class_ids"] = [] + attrs["class_weights"] = [] return attrs def get_default_tree_regressor_attribute_pairs(): attrs = {} - attrs['post_transform'] = 'NONE' - attrs['n_targets'] = 0 - attrs['nodes_treeids'] = [] - attrs['nodes_nodeids'] = [] - attrs['nodes_featureids'] = [] - attrs['nodes_modes'] = [] - attrs['nodes_values'] = [] - attrs['nodes_truenodeids'] = [] - attrs['nodes_falsenodeids'] = [] - attrs['nodes_missing_value_tracks_true'] = [] - attrs['nodes_hitrates'] = [] - attrs['target_treeids'] = [] - attrs['target_nodeids'] = [] - attrs['target_ids'] = [] - attrs['target_weights'] = [] + attrs["post_transform"] = "NONE" + attrs["n_targets"] = 0 + attrs["nodes_treeids"] = [] + attrs["nodes_nodeids"] = [] + attrs["nodes_featureids"] = [] + attrs["nodes_modes"] = [] + attrs["nodes_values"] = [] + attrs["nodes_truenodeids"] = [] + attrs["nodes_falsenodeids"] = [] + attrs["nodes_missing_value_tracks_true"] = [] + attrs["nodes_hitrates"] = [] + attrs["target_treeids"] = [] + attrs["target_nodeids"] = [] + attrs["target_ids"] = [] + attrs["target_weights"] = [] return attrs -def add_node(attr_pairs, is_classifier, tree_id, tree_weight, node_id, feature_id, mode, value, true_child_id, - false_child_id, weights, weight_id_bias, leaf_weights_are_counts): - attr_pairs['nodes_treeids'].append(tree_id) - attr_pairs['nodes_nodeids'].append(node_id) - attr_pairs['nodes_featureids'].append(feature_id) - attr_pairs['nodes_modes'].append(mode) - attr_pairs['nodes_values'].append(value) - attr_pairs['nodes_truenodeids'].append(true_child_id) - attr_pairs['nodes_falsenodeids'].append(false_child_id) - attr_pairs['nodes_missing_value_tracks_true'].append(False) - attr_pairs['nodes_hitrates'].append(1.) +def add_node( + attr_pairs, + is_classifier, + tree_id, + tree_weight, + node_id, + feature_id, + mode, + value, + true_child_id, + false_child_id, + weights, + weight_id_bias, + leaf_weights_are_counts, +): + attr_pairs["nodes_treeids"].append(tree_id) + attr_pairs["nodes_nodeids"].append(node_id) + attr_pairs["nodes_featureids"].append(feature_id) + attr_pairs["nodes_modes"].append(mode) + attr_pairs["nodes_values"].append(value) + attr_pairs["nodes_truenodeids"].append(true_child_id) + attr_pairs["nodes_falsenodeids"].append(false_child_id) + attr_pairs["nodes_missing_value_tracks_true"].append(False) + attr_pairs["nodes_hitrates"].append(1.0) # Add leaf information for making prediction - if mode == 'LEAF': + if mode == "LEAF": flattened_weights = weights.flatten() factor = tree_weight - # If the values stored at leaves are counts of possible classes, we need convert them to probabilities by + # If the values stored at leaves are counts of possible classes, + # we need convert them to probabilities by # doing a normalization. if leaf_weights_are_counts: s = sum(flattened_weights) - factor /= float(s) if s != 0. else 1. + factor /= float(s) if s != 0.0 else 1.0 flattened_weights = [w * factor for w in flattened_weights] if len(flattened_weights) == 2 and is_classifier: flattened_weights = [flattened_weights[1]] - # Note that attribute names for making prediction are different for classifiers and regressors + # Note that attribute names for making prediction + # are different for classifiers and regressors if is_classifier: for i, w in enumerate(flattened_weights): - attr_pairs['class_treeids'].append(tree_id) - attr_pairs['class_nodeids'].append(node_id) - attr_pairs['class_ids'].append(i + weight_id_bias) - attr_pairs['class_weights'].append(w) + attr_pairs["class_treeids"].append(tree_id) + attr_pairs["class_nodeids"].append(node_id) + attr_pairs["class_ids"].append(i + weight_id_bias) + attr_pairs["class_weights"].append(w) else: for i, w in enumerate(flattened_weights): - attr_pairs['target_treeids'].append(tree_id) - attr_pairs['target_nodeids'].append(node_id) - attr_pairs['target_ids'].append(i + weight_id_bias) - attr_pairs['target_weights'].append(w) - - -def add_tree_to_attribute_pairs(attr_pairs, is_classifier, tree, tree_id, tree_weight, - weight_id_bias, leaf_weights_are_counts): + attr_pairs["target_treeids"].append(tree_id) + attr_pairs["target_nodeids"].append(node_id) + attr_pairs["target_ids"].append(i + weight_id_bias) + attr_pairs["target_weights"].append(w) + + +def add_tree_to_attribute_pairs( + attr_pairs, + is_classifier, + tree, + tree_id, + tree_weight, + weight_id_bias, + leaf_weights_are_counts, +): for i in range(tree.node_count): node_id = tree.nodes_ids[i] weight = tree.value[i] if tree.children_left[i] >= 0 or tree.children_right[i] >= 0: - mode = 'BRANCH_LEQ' + mode = "BRANCH_LEQ" feat_id = tree.feature[i] threshold = tree.threshold[i] left_child_id = int(tree.children_left[i]) right_child_id = int(tree.children_right[i]) else: - mode = 'LEAF' + mode = "LEAF" feat_id = 0 - threshold = 0. + threshold = 0.0 left_child_id = 0 right_child_id = 0 - add_node(attr_pairs, is_classifier, tree_id, tree_weight, node_id, feat_id, mode, threshold, - left_child_id, right_child_id, weight, weight_id_bias, leaf_weights_are_counts) + add_node( + attr_pairs, + is_classifier, + tree_id, + tree_weight, + node_id, + feat_id, + mode, + threshold, + left_child_id, + right_child_id, + weight, + weight_id_bias, + leaf_weights_are_counts, + ) diff --git a/onnxmltools/convert/sparkml/operator_converters/tree_helper.py b/onnxmltools/convert/sparkml/operator_converters/tree_helper.py index f8f1ffef..53a98e9e 100644 --- a/onnxmltools/convert/sparkml/operator_converters/tree_helper.py +++ b/onnxmltools/convert/sparkml/operator_converters/tree_helper.py @@ -4,7 +4,6 @@ class Node: - _names_classifier = [ "class_ids", "class_nodeids", @@ -90,7 +89,7 @@ def create(attrs): zip(attrs["target_treeids"], attrs["target_nodeids"]) ) if t == tid and c == nid - ] + ] for k, v in attrs.items(): if k in {"post_transform", "name", "domain", "n_targets"}: continue @@ -173,7 +172,6 @@ def _unfold_rule_or(self): return None def unfold_rule_or(self): - cond = True nodes = [] for node in self: if node.nodes_modes == "||": @@ -251,18 +249,26 @@ def to_attrs(self, **kwargs): # update numbers new_numbers = {} - for tid, nid, md in sorted(zip(attrs["nodes_treeids"], attrs["nodes_nodeids"], attrs["nodes_modes"])): + for tid, nid, md in sorted( + zip(attrs["nodes_treeids"], attrs["nodes_nodeids"], attrs["nodes_modes"]) + ): new_numbers[tid, nid] = len(new_numbers) - for k in ["nodes_truenodeids", "nodes_falsenodeids", "nodes_nodeids", "class_nodeids", "target_nodeids"]: + for k in [ + "nodes_truenodeids", + "nodes_falsenodeids", + "nodes_nodeids", + "class_nodeids", + "target_nodeids", + ]: if k not in attrs: continue if "class_" in k or "target_" in k: - field = k.split('_')[0] + "_treeids" + field = k.split("_")[0] + "_treeids" else: field = "nodes_treeids" for i in range(len(attrs[k])): nid = attrs[k][i] - if nid == 0 and k in {'nodes_truenodeids', 'nodes_falsenodeids'}: + if nid == 0 and k in {"nodes_truenodeids", "nodes_falsenodeids"}: continue tid = attrs[field][i] new_id = new_numbers[tid, nid] @@ -276,7 +282,10 @@ def rewrite_ids_and_process(attrs, logger): if isinstance(value, (np.ndarray, list)): in_sets_rules.append(i) - logger.info("[convert_decision_tree_classifier] in_set_rules has %d elements", len(in_sets_rules)) + logger.info( + "[convert_decision_tree_classifier] in_set_rules has %d elements", + len(in_sets_rules), + ) for i in in_sets_rules: attrs["nodes_modes"][i] = "||" logger.info("[convert_decision_tree_classifier] Node.create") @@ -286,21 +295,27 @@ def rewrite_ids_and_process(attrs, logger): logger.info("[convert_decision_tree_classifier] to_attrs") if "class_nodeids" in attrs: new_attrs = root.to_attrs( - post_transform=attrs['post_transform'], - classlabels_int64s=attrs["classlabels_int64s"], - name=attrs['name']) + post_transform=attrs["post_transform"], + classlabels_int64s=attrs["classlabels_int64s"], + name=attrs["name"], + ) else: new_attrs = root.to_attrs( - post_transform=attrs['post_transform'], - n_targets=attrs['n_targets'], - name=attrs['name']) - if len(attrs['nodes_nodeids']) > len(new_attrs['nodes_nodeids']): + post_transform=attrs["post_transform"], + n_targets=attrs["n_targets"], + name=attrs["name"], + ) + if len(attrs["nodes_nodeids"]) > len(new_attrs["nodes_nodeids"]): raise RuntimeError( f"The replacement fails as there are less nodes in the new tree, " f"{len(attrs['nodes_nodeids'])} > {len(new_attrs['nodes_nodeids'])}." ) if set(attrs) != set(new_attrs): - raise RuntimeError(f"Missing key: {list(sorted(attrs))} != {list(sorted(new_attrs))}.") - logger.info("[convert_decision_tree_classifier] n_nodes=%d", len(attrs['nodes_nodeids'])) + raise RuntimeError( + f"Missing key: {list(sorted(attrs))} != {list(sorted(new_attrs))}." + ) + logger.info( + "[convert_decision_tree_classifier] n_nodes=%d", len(attrs["nodes_nodeids"]) + ) logger.info("[convert_decision_tree_classifier] end") return new_attrs diff --git a/onnxmltools/convert/sparkml/operator_converters/vector_assembler.py b/onnxmltools/convert/sparkml/operator_converters/vector_assembler.py index 4e90b5a4..0c8dd197 100644 --- a/onnxmltools/convert/sparkml/operator_converters/vector_assembler.py +++ b/onnxmltools/convert/sparkml/operator_converters/vector_assembler.py @@ -3,17 +3,23 @@ from ...common._registration import register_converter from ...common._registration import register_shape_calculator from ...common.utils import check_input_and_output_numbers -from ...common.data_types import * +from ...common.data_types import FloatTensorType, Int64TensorType def convert_sparkml_vector_assembler(scope, operator, container): - container.add_node('Concat', [s for s in operator.input_full_names], operator.outputs[0].full_name, - name=scope.get_unique_operator_name('Concat'), - op_version=4, - axis=1) + container.add_node( + "Concat", + [s for s in operator.input_full_names], + operator.outputs[0].full_name, + name=scope.get_unique_operator_name("Concat"), + op_version=4, + axis=1, + ) -register_converter('pyspark.ml.feature.VectorAssembler', convert_sparkml_vector_assembler) +register_converter( + "pyspark.ml.feature.VectorAssembler", convert_sparkml_vector_assembler +) def calculate_vector_assembler_shapes(operator): @@ -32,4 +38,6 @@ def calculate_vector_assembler_shapes(operator): operator.outputs[0].type = col_type -register_shape_calculator('pyspark.ml.feature.VectorAssembler', calculate_vector_assembler_shapes) +register_shape_calculator( + "pyspark.ml.feature.VectorAssembler", calculate_vector_assembler_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/vector_indexer.py b/onnxmltools/convert/sparkml/operator_converters/vector_indexer.py index 319c1674..7da73258 100644 --- a/onnxmltools/convert/sparkml/operator_converters/vector_indexer.py +++ b/onnxmltools/convert/sparkml/operator_converters/vector_indexer.py @@ -3,54 +3,81 @@ from onnx import onnx_pb as onnx_proto from ...common._registration import register_converter, register_shape_calculator from ...common.utils import check_input_and_output_numbers -from ...common.data_types import * +from ...common.data_types import FloatTensorType def convert_sparkml_vector_indexer(scope, operator, container): feature_count = operator.raw_operator.numFeatures category_map = operator.raw_operator.categoryMaps - split_output_names = [ scope.get_unique_variable_name('split_tensor_%d' % i) for i in range(0, feature_count)] + split_output_names = [ + scope.get_unique_variable_name("split_tensor_%d" % i) + for i in range(0, feature_count) + ] if feature_count > 1: - container.add_node('Split', operator.inputs[0].full_name, split_output_names, - name=scope.get_unique_operator_name('Split'), - op_version=2, - axis=1, - split=[1]*feature_count) + container.add_node( + "Split", + operator.inputs[0].full_name, + split_output_names, + name=scope.get_unique_operator_name("Split"), + op_version=2, + axis=1, + split=[1] * feature_count, + ) else: split_output_names = operator.input_full_names concat_inputs = split_output_names.copy() for i in category_map.keys(): - converted_output = scope.get_unique_variable_name('converted_tensor_%d' % i) - container.add_node('Cast', split_output_names[i], converted_output, - name=scope.get_unique_operator_name('Cast'), - op_version=9, - to=onnx_proto.TensorProto.STRING) + converted_output = scope.get_unique_variable_name("converted_tensor_%d" % i) + container.add_node( + "Cast", + split_output_names[i], + converted_output, + name=scope.get_unique_operator_name("Cast"), + op_version=9, + to=onnx_proto.TensorProto.STRING, + ) attrs = { - 'name': scope.get_unique_operator_name('LabelEncoder'), - 'classes_strings': ['{0:g}'.format(c) for c in category_map[i].keys()], - 'default_string': '__unknown__' + "name": scope.get_unique_operator_name("LabelEncoder"), + "classes_strings": ["{0:g}".format(c) for c in category_map[i].keys()], + "default_string": "__unknown__", } - encoded_output_name = scope.get_unique_variable_name('indexed_tensor_%d' % i) - container.add_node('LabelEncoder', converted_output, encoded_output_name, - op_domain='ai.onnx.ml', - **attrs) - converted_float_output = scope.get_unique_variable_name('converted_float_tensor_%d' % i) + encoded_output_name = scope.get_unique_variable_name("indexed_tensor_%d" % i) + container.add_node( + "LabelEncoder", + converted_output, + encoded_output_name, + op_domain="ai.onnx.ml", + **attrs + ) + converted_float_output = scope.get_unique_variable_name( + "converted_float_tensor_%d" % i + ) if feature_count == 1: converted_float_output = operator.output_full_names[0] - container.add_node('Cast', encoded_output_name, converted_float_output, - name=scope.get_unique_operator_name('Cast'), - op_version=9, - to=onnx_proto.TensorProto.FLOAT) + container.add_node( + "Cast", + encoded_output_name, + converted_float_output, + name=scope.get_unique_operator_name("Cast"), + op_version=9, + to=onnx_proto.TensorProto.FLOAT, + ) concat_inputs[i] = converted_float_output # add the final Concat if feature_count > 1: - container.add_node('Concat', concat_inputs, operator.output_full_names[0], - name=scope.get_unique_operator_name('Concat'), - op_version=4, - axis=1) + container.add_node( + "Concat", + concat_inputs, + operator.output_full_names[0], + name=scope.get_unique_operator_name("Concat"), + op_version=4, + axis=1, + ) -register_converter('pyspark.ml.feature.VectorIndexerModel', convert_sparkml_vector_indexer) +register_converter( + "pyspark.ml.feature.VectorIndexerModel", convert_sparkml_vector_indexer +) def calculate_vector_indexer_shapes(operator): @@ -59,4 +86,6 @@ def calculate_vector_indexer_shapes(operator): operator.outputs[0].type = FloatTensorType([N, operator.raw_operator.numFeatures]) -register_shape_calculator('pyspark.ml.feature.VectorIndexerModel', calculate_vector_indexer_shapes) +register_shape_calculator( + "pyspark.ml.feature.VectorIndexerModel", calculate_vector_indexer_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/vector_slicer.py b/onnxmltools/convert/sparkml/operator_converters/vector_slicer.py index 0454d72c..abd28a52 100644 --- a/onnxmltools/convert/sparkml/operator_converters/vector_slicer.py +++ b/onnxmltools/convert/sparkml/operator_converters/vector_slicer.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 import copy +import onnx from ...common._registration import register_converter, register_shape_calculator from ...common.utils import check_input_and_output_numbers, check_input_and_output_types -from ...common.data_types import * +from ...common.data_types import StringTensorType, Int64TensorType, FloatTensorType def convert_sparkml_vector_slicer(scope, operator, container): @@ -12,23 +13,31 @@ def convert_sparkml_vector_slicer(scope, operator, container): if not indices: raise ValueError("Indices are needed for conversion of VectorSlicer") - indices_tensor = 'indices_tensor' - container.add_initializer(indices_tensor, onnx_proto.TensorProto.INT64, [len(indices)], indices) - container.add_node('ArrayFeatureExtractor', - [operator.input_full_names[0], indices_tensor], operator.output_full_names, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('ArrayFeatureExtractor')) + indices_tensor = "indices_tensor" + container.add_initializer( + indices_tensor, onnx.TensorProto.INT64, [len(indices)], indices + ) + container.add_node( + "ArrayFeatureExtractor", + [operator.input_full_names[0], indices_tensor], + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("ArrayFeatureExtractor"), + ) -register_converter('pyspark.ml.feature.VectorSlicer', convert_sparkml_vector_slicer) +register_converter("pyspark.ml.feature.VectorSlicer", convert_sparkml_vector_slicer) def calculate_vector_slicer_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=1) - check_input_and_output_types(operator, - good_input_types=[FloatTensorType, Int64TensorType, StringTensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType, StringTensorType] + ) operator.outputs[0].type = copy.deepcopy(operator.inputs[0].type) operator.outputs[0].type.shape[1] = len(operator.raw_operator.getIndices()) -register_shape_calculator('pyspark.ml.feature.VectorSlicer', calculate_vector_slicer_shapes) +register_shape_calculator( + "pyspark.ml.feature.VectorSlicer", calculate_vector_slicer_shapes +) diff --git a/onnxmltools/convert/sparkml/operator_converters/word2vec.py b/onnxmltools/convert/sparkml/operator_converters/word2vec.py index 67c3a710..5f59acf0 100644 --- a/onnxmltools/convert/sparkml/operator_converters/word2vec.py +++ b/onnxmltools/convert/sparkml/operator_converters/word2vec.py @@ -11,49 +11,74 @@ def convert_word2vec(scope, operator, container): op = operator.raw_operator - vectors = op.getVectors().toPandas().vector.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + vectors = ( + op.getVectors() + .toPandas() + .vector.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) cats_strings = op.getVectors().toPandas().word.values.tolist() cats_int64s = [x for x in range(0, len(cats_strings))] word_count = operator.inputs[0].type.shape[1] - vectors_tensor = scope.get_unique_variable_name('vectors_tensor') - container.add_initializer(vectors_tensor, onnx_proto.TensorProto.FLOAT, vectors.shape, vectors.flatten()) - word_indices = scope.get_unique_variable_name('word_indices_tensor') - container.add_node('CategoryMapper', operator.input_full_names, word_indices, - op_domain='ai.onnx.ml', - cats_int64s=cats_int64s, - cats_strings=cats_strings, - default_int64=-1) - one = scope.get_unique_variable_name('one_tensor') + vectors_tensor = scope.get_unique_variable_name("vectors_tensor") + container.add_initializer( + vectors_tensor, onnx_proto.TensorProto.FLOAT, vectors.shape, vectors.flatten() + ) + word_indices = scope.get_unique_variable_name("word_indices_tensor") + container.add_node( + "CategoryMapper", + operator.input_full_names, + word_indices, + op_domain="ai.onnx.ml", + cats_int64s=cats_int64s, + cats_strings=cats_strings, + default_int64=-1, + ) + one = scope.get_unique_variable_name("one_tensor") container.add_initializer(one, onnx_proto.TensorProto.INT64, [1], [1]) - zero = scope.get_unique_variable_name('zero_tensor') + zero = scope.get_unique_variable_name("zero_tensor") container.add_initializer(zero, onnx_proto.TensorProto.INT64, [1], [0]) sliced_outputs = [] for i in range(0, word_count): - index = scope.get_unique_variable_name('index_tensor') + index = scope.get_unique_variable_name("index_tensor") container.add_initializer(index, onnx_proto.TensorProto.INT64, [1], [i]) - selected_index = scope.get_unique_variable_name('selected_index_tensor') - container.add_node('ArrayFeatureExtractor', [word_indices, index], selected_index, - op_domain='ai.onnx.ml') - reshaped_index = scope.get_unique_variable_name('reshaped_tensor') - container.add_node('Reshape', [selected_index, one], reshaped_index, - op_version=5) - end_index = scope.get_unique_variable_name('end_index_tensor') + selected_index = scope.get_unique_variable_name("selected_index_tensor") + container.add_node( + "ArrayFeatureExtractor", + [word_indices, index], + selected_index, + op_domain="ai.onnx.ml", + ) + reshaped_index = scope.get_unique_variable_name("reshaped_tensor") + container.add_node( + "Reshape", [selected_index, one], reshaped_index, op_version=5 + ) + end_index = scope.get_unique_variable_name("end_index_tensor") apply_add(scope, [one, reshaped_index], end_index, container, axis=0) - sliced_output = scope.get_unique_variable_name('sliced_tensor') - container.add_node('DynamicSlice', [vectors_tensor, reshaped_index, end_index, zero], sliced_output) + sliced_output = scope.get_unique_variable_name("sliced_tensor") + container.add_node( + "DynamicSlice", + [vectors_tensor, reshaped_index, end_index, zero], + sliced_output, + ) sliced_outputs.append(sliced_output) - sum_vector = scope.get_unique_variable_name('sum_tensor') + sum_vector = scope.get_unique_variable_name("sum_tensor") apply_sum(scope, sliced_outputs, sum_vector, container) - factor = scope.get_unique_variable_name('factor_tensor') - container.add_initializer(factor, onnx_proto.TensorProto.FLOAT, [1], [1/operator.inputs[0].type.shape[1]]) + factor = scope.get_unique_variable_name("factor_tensor") + container.add_initializer( + factor, + onnx_proto.TensorProto.FLOAT, + [1], + [1 / operator.inputs[0].type.shape[1]], + ) apply_mul(scope, [factor, sum_vector], operator.output_full_names, container) -register_converter('pyspark.ml.feature.Word2VecModel', convert_word2vec) +register_converter("pyspark.ml.feature.Word2VecModel", convert_word2vec) def calculate_word2vec_output_shapes(operator): @@ -61,8 +86,10 @@ def calculate_word2vec_output_shapes(operator): check_input_and_output_types(operator, good_input_types=[StringTensorType]) N = operator.inputs[0].type.shape[0] - C = operator.raw_operator.getOrDefault('vectorSize') + C = operator.raw_operator.getOrDefault("vectorSize") operator.outputs[0].type = FloatTensorType([N, C]) -register_shape_calculator('pyspark.ml.feature.Word2VecModel', calculate_word2vec_output_shapes) +register_shape_calculator( + "pyspark.ml.feature.Word2VecModel", calculate_word2vec_output_shapes +) diff --git a/onnxmltools/convert/sparkml/ops_input_output.py b/onnxmltools/convert/sparkml/ops_input_output.py index 988ac8cc..ae241667 100644 --- a/onnxmltools/convert/sparkml/ops_input_output.py +++ b/onnxmltools/convert/sparkml/ops_input_output.py @@ -71,7 +71,10 @@ def build_io_name_map(): ), "pyspark.ml.classification.NaiveBayesModel": ( lambda model: [model.getOrDefault("featuresCol")], - lambda model: [model.getOrDefault("predictionCol"), model.getOrDefault("probabilityCol")], + lambda model: [ + model.getOrDefault("predictionCol"), + model.getOrDefault("probabilityCol"), + ], ), "pyspark.ml.feature.VectorSlicer": ( lambda model: [model.getOrDefault("inputCol")], @@ -95,11 +98,17 @@ def build_io_name_map(): ), "pyspark.ml.classification.RandomForestClassificationModel": ( lambda model: [model.getOrDefault("featuresCol")], - lambda model: [model.getOrDefault("predictionCol"), model.getOrDefault("probabilityCol")], + lambda model: [ + model.getOrDefault("predictionCol"), + model.getOrDefault("probabilityCol"), + ], ), "pyspark.ml.classification.MultilayerPerceptronClassificationModel": ( lambda model: [model.getOrDefault("featuresCol")], - lambda model: [model.getOrDefault("predictionCol"), model.getOrDefault("probabilityCol")], + lambda model: [ + model.getOrDefault("predictionCol"), + model.getOrDefault("probabilityCol"), + ], ), "pyspark.ml.regression.DecisionTreeRegressionModel": ( lambda model: [model.getOrDefault("featuresCol")], @@ -107,7 +116,10 @@ def build_io_name_map(): ), "pyspark.ml.classification.DecisionTreeClassificationModel": ( lambda model: [model.getOrDefault("featuresCol")], - lambda model: [model.getOrDefault("predictionCol"), model.getOrDefault("probabilityCol")], + lambda model: [ + model.getOrDefault("predictionCol"), + model.getOrDefault("probabilityCol"), + ], ), "pyspark.ml.feature.VectorIndexerModel": ( lambda model: [model.getOrDefault("inputCol")], @@ -159,7 +171,10 @@ def build_io_name_map(): ), "pyspark.ml.classification.LogisticRegressionModel": ( lambda model: [model.getOrDefault("featuresCol")], - lambda model: [model.getOrDefault("predictionCol"), model.getOrDefault("probabilityCol")], + lambda model: [ + model.getOrDefault("predictionCol"), + model.getOrDefault("probabilityCol"), + ], ), "pyspark.ml.feature.OneHotEncoderModel": ( lambda model: model.getOrDefault("inputCols") diff --git a/onnxmltools/convert/sparkml/ops_names.py b/onnxmltools/convert/sparkml/ops_names.py index 0ed8cc3f..6a46614d 100644 --- a/onnxmltools/convert/sparkml/ops_names.py +++ b/onnxmltools/convert/sparkml/ops_names.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -''' +""" Mapping and utility functions for Name to Spark ML operators -''' +""" from pyspark.ml import Transformer, Estimator from pyspark.ml.feature import Binarizer @@ -35,45 +35,96 @@ from pyspark.ml.feature import VectorSlicer from pyspark.ml.feature import Word2VecModel -from pyspark.ml.classification import LinearSVCModel, RandomForestClassificationModel, GBTClassificationModel, \ - MultilayerPerceptronClassificationModel +from pyspark.ml.classification import ( + LinearSVCModel, + RandomForestClassificationModel, + GBTClassificationModel, + MultilayerPerceptronClassificationModel, +) from pyspark.ml.classification import LogisticRegressionModel from pyspark.ml.classification import DecisionTreeClassificationModel from pyspark.ml.classification import NaiveBayesModel from pyspark.ml.classification import OneVsRestModel -from pyspark.ml.regression import AFTSurvivalRegressionModel, DecisionTreeRegressionModel, RandomForestRegressionModel +from pyspark.ml.regression import ( + AFTSurvivalRegressionModel, + DecisionTreeRegressionModel, + RandomForestRegressionModel, +) from pyspark.ml.regression import GBTRegressionModel from pyspark.ml.regression import GeneralizedLinearRegressionModel from pyspark.ml.regression import IsotonicRegressionModel from pyspark.ml.regression import LinearRegressionModel -from pyspark.ml.clustering import BisectingKMeans from pyspark.ml.clustering import KMeansModel -from pyspark.ml.clustering import GaussianMixture -from pyspark.ml.clustering import LDA def build_sparkml_operator_name_map(): - res = {k: "pyspark.ml.feature." + k.__name__ for k in [ - Binarizer, BucketedRandomProjectionLSHModel, Bucketizer, - ChiSqSelectorModel, CountVectorizerModel, DCT, ElementwiseProduct, HashingTF, IDFModel, ImputerModel, - IndexToString, MaxAbsScalerModel, MinHashLSHModel, MinMaxScalerModel, NGram, Normalizer, OneHotEncoderModel, - PCAModel, PolynomialExpansion, QuantileDiscretizer, RegexTokenizer, - StandardScalerModel, StopWordsRemover, StringIndexerModel, Tokenizer, VectorAssembler, VectorIndexerModel, - VectorSlicer, Word2VecModel - ]} - res.update({k: "pyspark.ml.classification." + k.__name__ for k in [ - LinearSVCModel, LogisticRegressionModel, DecisionTreeClassificationModel, GBTClassificationModel, - RandomForestClassificationModel, NaiveBayesModel, MultilayerPerceptronClassificationModel, OneVsRestModel - ]}) - res.update({k: "pyspark.ml.regression." + k.__name__ for k in [ - AFTSurvivalRegressionModel, DecisionTreeRegressionModel, GBTRegressionModel, GBTRegressionModel, - GeneralizedLinearRegressionModel, IsotonicRegressionModel, LinearRegressionModel, RandomForestRegressionModel - ]}) - res.update({k: "pyspark.ml.clustering." + k.__name__ for k in [ - KMeansModel - ]}) + res = { + k: "pyspark.ml.feature." + k.__name__ + for k in [ + Binarizer, + BucketedRandomProjectionLSHModel, + Bucketizer, + ChiSqSelectorModel, + CountVectorizerModel, + DCT, + ElementwiseProduct, + HashingTF, + IDFModel, + ImputerModel, + IndexToString, + MaxAbsScalerModel, + MinHashLSHModel, + MinMaxScalerModel, + NGram, + Normalizer, + OneHotEncoderModel, + PCAModel, + PolynomialExpansion, + QuantileDiscretizer, + RegexTokenizer, + StandardScalerModel, + StopWordsRemover, + StringIndexerModel, + Tokenizer, + VectorAssembler, + VectorIndexerModel, + VectorSlicer, + Word2VecModel, + ] + } + res.update( + { + k: "pyspark.ml.classification." + k.__name__ + for k in [ + LinearSVCModel, + LogisticRegressionModel, + DecisionTreeClassificationModel, + GBTClassificationModel, + RandomForestClassificationModel, + NaiveBayesModel, + MultilayerPerceptronClassificationModel, + OneVsRestModel, + ] + } + ) + res.update( + { + k: "pyspark.ml.regression." + k.__name__ + for k in [ + AFTSurvivalRegressionModel, + DecisionTreeRegressionModel, + GBTRegressionModel, + GBTRegressionModel, + GeneralizedLinearRegressionModel, + IsotonicRegressionModel, + LinearRegressionModel, + RandomForestRegressionModel, + ] + } + ) + res.update({k: "pyspark.ml.clustering." + k.__name__ for k in [KMeansModel]}) return res @@ -81,19 +132,18 @@ def build_sparkml_operator_name_map(): def get_sparkml_operator_name(model_type): - ''' + """ Get operator name of the input argument :param model_type: A spark-ml object (LinearRegression, StringIndexer, ...) :return: A string which stands for the type of the input model in our conversion framework - ''' + """ if not issubclass(model_type, Transformer): if issubclass(model_type, Estimator): raise ValueError("Estimator must be fitted before being converted to ONNX") else: raise ValueError("Unknown model type: {}".format(model_type)) - + if model_type not in sparkml_operator_name_map: raise ValueError("No proper operator name found for '%s'" % model_type) return sparkml_operator_name_map[model_type] - diff --git a/onnxmltools/convert/sparkml/utils.py b/onnxmltools/convert/sparkml/utils.py index c64986bb..d1589ba0 100644 --- a/onnxmltools/convert/sparkml/utils.py +++ b/onnxmltools/convert/sparkml/utils.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -''' +""" Utility functions for Spark ML to Onnx conversion intended for the end user mainly -''' +""" from ..common.data_types import StringTensorType, FloatTensorType @@ -14,16 +14,26 @@ def buildInitialTypesSimple(dataframe): def getTensorTypeFromSpark(sparktype): - if sparktype == 'StringType' or sparktype == 'StringType()': + if sparktype == "StringType" or sparktype == "StringType()": return StringTensorType([1, 1]) - elif sparktype == 'DecimalType' or sparktype == 'DecimalType()' \ - or sparktype == 'DoubleType' or sparktype == 'DoubleType()' \ - or sparktype == 'FloatType' or sparktype == 'FloatType()' \ - or sparktype == 'LongType' or sparktype == 'LongType()' \ - or sparktype == 'IntegerType' or sparktype == 'IntegerType()' \ - or sparktype == 'ShortType' or sparktype == 'ShortType()' \ - or sparktype == 'ByteType' or sparktype == 'ByteType()' \ - or sparktype == 'BooleanType' or sparktype == 'BooleanType()': + elif ( + sparktype == "DecimalType" + or sparktype == "DecimalType()" + or sparktype == "DoubleType" + or sparktype == "DoubleType()" + or sparktype == "FloatType" + or sparktype == "FloatType()" + or sparktype == "LongType" + or sparktype == "LongType()" + or sparktype == "IntegerType" + or sparktype == "IntegerType()" + or sparktype == "ShortType" + or sparktype == "ShortType()" + or sparktype == "ByteType" + or sparktype == "ByteType()" + or sparktype == "BooleanType" + or sparktype == "BooleanType()" + ): return FloatTensorType([1, 1]) else: raise TypeError("Cannot map this type to Onnx types: " + sparktype) @@ -31,12 +41,15 @@ def getTensorTypeFromSpark(sparktype): def buildInputDictSimple(dataframe): import numpy + result = {} for field in dataframe.schema.fields: - if str(field.dataType) == 'StringType': + if str(field.dataType) == "StringType": result[field.name] = dataframe.select(field.name).toPandas().values else: - result[field.name] = dataframe.select(field.name).toPandas().values.astype(numpy.float32) + result[field.name] = ( + dataframe.select(field.name).toPandas().values.astype(numpy.float32) + ) return result diff --git a/onnxmltools/convert/xgboost/__init__.py b/onnxmltools/convert/xgboost/__init__.py index 11db0855..017852fc 100644 --- a/onnxmltools/convert/xgboost/__init__.py +++ b/onnxmltools/convert/xgboost/__init__.py @@ -1,3 +1,4 @@ # SPDX-License-Identifier: Apache-2.0 +from . import operator_converters, shape_calculators from .convert import convert diff --git a/onnxmltools/convert/xgboost/_parse.py b/onnxmltools/convert/xgboost/_parse.py index d5228806..605d8ddc 100644 --- a/onnxmltools/convert/xgboost/_parse.py +++ b/onnxmltools/convert/xgboost/_parse.py @@ -14,36 +14,39 @@ xgboost_classifier_list = [XGBClassifier] # Associate types with our operator names. -xgboost_operator_name_map = {XGBClassifier: 'XGBClassifier', - XGBRegressor: 'XGBRegressor'} +xgboost_operator_name_map = { + XGBClassifier: "XGBClassifier", + XGBRegressor: "XGBRegressor", +} def _append_covers(node): res = [] - if 'cover' in node: - res.append(node['cover']) - if 'children' in node: - for ch in node['children']: + if "cover" in node: + res.append(node["cover"]) + if "children" in node: + for ch in node["children"]: res.extend(_append_covers(ch)) return res def _get_attributes(booster): atts = booster.attributes() - dp = booster.get_dump(dump_format='json', with_stats=True) + dp = booster.get_dump(dump_format="json", with_stats=True) res = [json.loads(d) for d in dp] # num_class - if Version(__version__) < Version('1.5'): + if Version(__version__) < Version("1.5"): state = booster.__getstate__() - bstate = bytes(state['handle']) + bstate = bytes(state["handle"]) reg = re.compile(b'("tree_info":\\[[0-9,]*\\])') objs = list(set(reg.findall(bstate))) if len(objs) != 1: raise RuntimeError( "Unable to retrieve the tree coefficients from\n%s" - "" % bstate.decode("ascii", errors="ignore")) - tree_info = json.loads("{{{}}}".format(objs[0].decode('ascii')))['tree_info'] + "" % bstate.decode("ascii", errors="ignore") + ) + tree_info = json.loads("{{{}}}".format(objs[0].decode("ascii")))["tree_info"] num_class = len(set(tree_info)) trees = len(res) try: @@ -56,12 +59,13 @@ def _get_attributes(booster): num_class = trees // ntrees if num_class == 0: raise RuntimeError( - "Unable to retrieve the number of classes, trees=%d, ntrees=%d." % ( - trees, ntrees)) + "Unable to retrieve the number of classes, trees=%d, ntrees=%d." + % (trees, ntrees) + ) kwargs = atts.copy() - kwargs['feature_names'] = booster.feature_names - kwargs['n_estimators'] = ntrees + kwargs["feature_names"] = booster.feature_names + kwargs["n_estimators"] = ntrees # covers covs = [] @@ -70,49 +74,50 @@ def _get_attributes(booster): if all(map(lambda x: int(x) == x, set(covs))): # regression - kwargs['num_target'] = num_class - kwargs['num_class'] = 0 + kwargs["num_target"] = num_class + kwargs["num_class"] = 0 kwargs["objective"] = "reg:squarederror" else: # classification - kwargs['num_class'] = num_class + kwargs["num_class"] = num_class if num_class != 1: - if Version(__version__) < Version('1.5'): - reg = re.compile(b'(multi:[a-z]{1,15})') + if Version(__version__) < Version("1.5"): + reg = re.compile(b"(multi:[a-z]{1,15})") objs = list(set(reg.findall(bstate))) if len(objs) == 1: - kwargs["objective"] = objs[0].decode('ascii') + kwargs["objective"] = objs[0].decode("ascii") else: raise RuntimeError( "Unable to guess objective in %r (trees=%r, ntrees=%r, num_class=%r)" - "." % (objs, trees, ntrees, kwargs['num_class'])) + "." % (objs, trees, ntrees, kwargs["num_class"]) + ) else: att = json.loads(booster.save_config()) - kwargs["objective"] = att['learner']['objective']['name'] - nc = int(att['learner']['learner_model_param']['num_class']) + kwargs["objective"] = att["learner"]["objective"]["name"] + nc = int(att["learner"]["learner_model_param"]["num_class"]) if nc != num_class: raise RuntimeError( - "Mismatched value %r != %r from\n%s" % ( - nc, num_class, pprint.pformat(att))) + "Mismatched value %r != %r from\n%s" + % (nc, num_class, pprint.pformat(att)) + ) else: kwargs["objective"] = "binary:logistic" - if 'base_score' not in kwargs: - kwargs['base_score'] = 0.5 + if "base_score" not in kwargs: + kwargs["base_score"] = 0.5 return kwargs class WrappedBooster: - def __init__(self, booster): self.booster_ = booster self.kwargs = _get_attributes(booster) - if self.kwargs['num_class'] > 0: + if self.kwargs["num_class"] > 0: self.classes_ = self._generate_classes(self.kwargs) - self.operator_name = 'XGBClassifier' + self.operator_name = "XGBClassifier" else: - self.operator_name = 'XGBRegressor' + self.operator_name = "XGBRegressor" def get_xgb_params(self): return self.kwargs @@ -121,18 +126,18 @@ def get_booster(self): return self.booster_ def _generate_classes(self, model_dict): - if model_dict['num_class'] == 1: + if model_dict["num_class"] == 1: return np.asarray([0, 1]) - return np.arange(model_dict['num_class']) + return np.arange(model_dict["num_class"]) def _get_xgboost_operator_name(model): - ''' + """ Get operator name of the input argument :param model_type: A xgboost object. :return: A string which stands for the type of the input model in our conversion framework - ''' + """ if isinstance(model, WrappedBooster): return model.operator_name if type(model) not in xgboost_operator_name_map: @@ -141,53 +146,70 @@ def _get_xgboost_operator_name(model): def _parse_xgboost_simple_model(scope, model, inputs): - ''' + """ This function handles all non-pipeline models. :param scope: Scope object :param model: A xgboost object :param inputs: A list of variables :return: A list of output variables which will be passed to next stage - ''' - this_operator = scope.declare_local_operator(_get_xgboost_operator_name(model), model) + """ + this_operator = scope.declare_local_operator( + _get_xgboost_operator_name(model), model + ) this_operator.inputs = inputs - if (type(model) in xgboost_classifier_list or - getattr(model, 'operator_name', None) == 'XGBClassifier'): - # For classifiers, we may have two outputs, one for label and the other one for probabilities of all classes. - # Notice that their types here are not necessarily correct and they will be fixed in shape inference phase - label_variable = scope.declare_local_variable('label', FloatTensorType()) - probability_map_variable = scope.declare_local_variable('probabilities', FloatTensorType()) + if ( + type(model) in xgboost_classifier_list + or getattr(model, "operator_name", None) == "XGBClassifier" + ): + # For classifiers, we may have two outputs, one for label and + # the other one for probabilities of all classes. + # Notice that their types here are not necessarily correct + # and they will be fixed in shape inference phase + label_variable = scope.declare_local_variable("label", FloatTensorType()) + probability_map_variable = scope.declare_local_variable( + "probabilities", FloatTensorType() + ) this_operator.outputs.append(label_variable) this_operator.outputs.append(probability_map_variable) else: # We assume that all scikit-learn operator can only produce a single float tensor. - variable = scope.declare_local_variable('variable', FloatTensorType()) + variable = scope.declare_local_variable("variable", FloatTensorType()) this_operator.outputs.append(variable) return this_operator.outputs def _parse_xgboost(scope, model, inputs): - ''' - This is a delegate function. It doesn't nothing but invoke the correct parsing function according to the input + """ + This is a delegate function. It doesn't nothing but invoke + the correct parsing function according to the input model's type. :param scope: Scope object :param model: A xgboost object :param inputs: A list of variables :return: The output variables produced by the input model - ''' + """ return _parse_xgboost_simple_model(scope, model, inputs) -def parse_xgboost(model, initial_types=None, target_opset=None, - custom_conversion_functions=None, custom_shape_calculators=None): - +def parse_xgboost( + model, + initial_types=None, + target_opset=None, + custom_conversion_functions=None, + custom_shape_calculators=None, +): raw_model_container = XGBoostModelContainer(model) - topology = Topology(raw_model_container, default_batch_size='None', - initial_types=initial_types, target_opset=target_opset, - custom_conversion_functions=custom_conversion_functions, - custom_shape_calculators=custom_shape_calculators) - scope = topology.declare_scope('__root__') + topology = Topology( + raw_model_container, + default_batch_size="None", + initial_types=initial_types, + target_opset=target_opset, + custom_conversion_functions=custom_conversion_functions, + custom_shape_calculators=custom_shape_calculators, + ) + scope = topology.declare_scope("__root__") inputs = [] for var_name, initial_type in initial_types: diff --git a/onnxmltools/convert/xgboost/common.py b/onnxmltools/convert/xgboost/common.py index 42da2bd0..d2eacdbb 100644 --- a/onnxmltools/convert/xgboost/common.py +++ b/onnxmltools/convert/xgboost/common.py @@ -4,19 +4,19 @@ Common function to converters and shape calculators. """ + def get_xgb_params(xgb_node): """ Retrieves parameters of a model. """ - if hasattr(xgb_node, 'kwargs'): + if hasattr(xgb_node, "kwargs"): # XGBoost >= 0.7 params = xgb_node.get_xgb_params() else: # XGBoost < 0.7 params = xgb_node.__dict__ - if ('n_estimators' not in params and - hasattr(xgb_node, 'n_estimators')): + if "n_estimators" not in params and hasattr(xgb_node, "n_estimators"): # xgboost >= 1.0.2 - params['n_estimators'] = xgb_node.n_estimators + params["n_estimators"] = xgb_node.n_estimators return params diff --git a/onnxmltools/convert/xgboost/convert.py b/onnxmltools/convert/xgboost/convert.py index debc05c2..462d84e1 100644 --- a/onnxmltools/convert/xgboost/convert.py +++ b/onnxmltools/convert/xgboost/convert.py @@ -9,42 +9,64 @@ # Invoke the registration of all our converters and shape calculators # from . import shape_calculators -from . import operator_converters, shape_calculators -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): - ''' +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 xgboost model. :param model: A xgboost model - :param initial_types: a python list. Each element is a tuple of a variable name and a type defined in data_types.py - :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 name: The name of the graph (type: GraphProto) + in the produced ONNX model (type: ModelProto) :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 + :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 - ''' + """ if initial_types is None: - raise ValueError('Initial types are required. See usage of convert(...) in \ - onnxmltools.convert.xgboost.convert for details') + raise ValueError( + "Initial types are required. See usage of convert(...) in \ + onnxmltools.convert.xgboost.convert for details" + ) if name is None: name = str(uuid4().hex) if isinstance(model, xgboost.Booster): model = WrappedBooster(model) target_opset = target_opset if target_opset else get_maximum_opset_supported() - topology = parse_xgboost(model, initial_types, target_opset, custom_conversion_functions, custom_shape_calculators) + topology = parse_xgboost( + 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) + onnx_model = convert_topology( + topology, name, doc_string, target_opset, targeted_onnx + ) opsets = {d.domain: d.version for d in onnx_model.opset_import} - if '' in opsets and opsets[''] < 9 and target_opset >= 9: + if "" in opsets and opsets[""] < 9 and target_opset >= 9: # a bug? - opsets[''] = 9 + opsets[""] = 9 del onnx_model.opset_import[:] for k, v in opsets.items(): opset = onnx_model.opset_import.add() diff --git a/onnxmltools/convert/xgboost/operator_converters/XGBoost.py b/onnxmltools/convert/xgboost/operator_converters/XGBoost.py index 4ae246ab..b88be168 100644 --- a/onnxmltools/convert/xgboost/operator_converters/XGBoost.py +++ b/onnxmltools/convert/xgboost/operator_converters/XGBoost.py @@ -9,7 +9,6 @@ class XGBConverter: - @staticmethod def get_xgb_params(xgb_node): """ @@ -22,12 +21,14 @@ def validate(xgb_node): params = XGBConverter.get_xgb_params(xgb_node) try: if "objective" not in params: - raise AttributeError('ojective') + raise AttributeError("ojective") except AttributeError as e: - raise RuntimeError('Missing attribute in XGBoost model ' + str(e)) - if hasattr(xgb_node, 'missing') and not np.isnan(xgb_node.missing): - raise RuntimeError("Cannot convert a XGBoost model where missing values are not " - "nan but {}.".format(xgb_node.missing)) + raise RuntimeError("Missing attribute in XGBoost model " + str(e)) + if hasattr(xgb_node, "missing") and not np.isnan(xgb_node.missing): + raise RuntimeError( + "Cannot convert a XGBoost model where missing values are not " + "nan but {}.".format(xgb_node.missing) + ) @staticmethod def common_members(xgb_node, inputs): @@ -39,29 +40,54 @@ def common_members(xgb_node, inputs): booster = xgb_node.get_booster() # The json format was available in October 2017. # XGBoost 0.7 was the first version released with it. - js_tree_list = booster.get_dump(with_stats=True, dump_format='json') + js_tree_list = booster.get_dump(with_stats=True, dump_format="json") js_trees = [json.loads(s) for s in js_tree_list] return objective, base_score, js_trees @staticmethod def _get_default_tree_attribute_pairs(is_classifier): attrs = {} - for k in {'nodes_treeids', 'nodes_nodeids', - 'nodes_featureids', 'nodes_modes', 'nodes_values', - 'nodes_truenodeids', 'nodes_falsenodeids', 'nodes_missing_value_tracks_true'}: + for k in { + "nodes_treeids", + "nodes_nodeids", + "nodes_featureids", + "nodes_modes", + "nodes_values", + "nodes_truenodeids", + "nodes_falsenodeids", + "nodes_missing_value_tracks_true", + }: attrs[k] = [] if is_classifier: - for k in {'class_treeids', 'class_nodeids', 'class_ids', 'class_weights'}: + for k in {"class_treeids", "class_nodeids", "class_ids", "class_weights"}: attrs[k] = [] else: - for k in {'target_treeids', 'target_nodeids', 'target_ids', 'target_weights'}: + for k in { + "target_treeids", + "target_nodeids", + "target_ids", + "target_weights", + }: attrs[k] = [] return attrs @staticmethod - def _add_node(attr_pairs, is_classifier, tree_id, tree_weight, node_id, - feature_id, mode, value, true_child_id, false_child_id, weights, weight_id_bias, - missing, hitrate): + def _add_node( + attr_pairs, + is_classifier, + tree_id, + tree_weight, + node_id, + feature_id, + mode, + value, + true_child_id, + false_child_id, + weights, + weight_id_bias, + missing, + hitrate, + ): if isinstance(feature_id, str): # Something like f0, f1... if feature_id[0] == "f": @@ -70,16 +96,16 @@ def _add_node(attr_pairs, is_classifier, tree_id, tree_weight, node_id, except ValueError: raise RuntimeError( "Unable to interpret '{0}', feature " - "names should follow pattern 'f%d'.".format( - feature_id)) + "names should follow pattern 'f%d'.".format(feature_id) + ) else: try: feature_id = int(float(feature_id)) except ValueError: raise RuntimeError( "Unable to interpret '{0}', feature " - "names should follow pattern 'f%d'.".format( - feature_id)) + "names should follow pattern 'f%d'.".format(feature_id) + ) # Split condition for sklearn # * if X_ptr[X_sample_stride * i + X_fx_stride * node.feature] <= node.threshold: @@ -88,69 +114,89 @@ def _add_node(attr_pairs, is_classifier, tree_id, tree_weight, node_id, # * if (fvalue < split_value) # * https://github.com/dmlc/xgboost/blob/master/include/xgboost/tree_model.h#L804 - attr_pairs['nodes_treeids'].append(tree_id) - attr_pairs['nodes_nodeids'].append(node_id) - attr_pairs['nodes_featureids'].append(feature_id) - attr_pairs['nodes_modes'].append(mode) - attr_pairs['nodes_values'].append(float(value)) - attr_pairs['nodes_truenodeids'].append(true_child_id) - attr_pairs['nodes_falsenodeids'].append(false_child_id) - attr_pairs['nodes_missing_value_tracks_true'].append(missing) - if 'nodes_hitrates' in attr_pairs: - attr_pairs['nodes_hitrates'].append(hitrate) - if mode == 'LEAF': + attr_pairs["nodes_treeids"].append(tree_id) + attr_pairs["nodes_nodeids"].append(node_id) + attr_pairs["nodes_featureids"].append(feature_id) + attr_pairs["nodes_modes"].append(mode) + attr_pairs["nodes_values"].append(float(value)) + attr_pairs["nodes_truenodeids"].append(true_child_id) + attr_pairs["nodes_falsenodeids"].append(false_child_id) + attr_pairs["nodes_missing_value_tracks_true"].append(missing) + if "nodes_hitrates" in attr_pairs: + attr_pairs["nodes_hitrates"].append(hitrate) + if mode == "LEAF": if is_classifier: for i, w in enumerate(weights): - attr_pairs['class_treeids'].append(tree_id) - attr_pairs['class_nodeids'].append(node_id) - attr_pairs['class_ids'].append(i + weight_id_bias) - attr_pairs['class_weights'].append(float(tree_weight * w)) + attr_pairs["class_treeids"].append(tree_id) + attr_pairs["class_nodeids"].append(node_id) + attr_pairs["class_ids"].append(i + weight_id_bias) + attr_pairs["class_weights"].append(float(tree_weight * w)) else: for i, w in enumerate(weights): - attr_pairs['target_treeids'].append(tree_id) - attr_pairs['target_nodeids'].append(node_id) - attr_pairs['target_ids'].append(i + weight_id_bias) - attr_pairs['target_weights'].append(float(tree_weight * w)) + attr_pairs["target_treeids"].append(tree_id) + attr_pairs["target_nodeids"].append(node_id) + attr_pairs["target_ids"].append(i + weight_id_bias) + attr_pairs["target_weights"].append(float(tree_weight * w)) @staticmethod - def _fill_node_attributes(treeid, tree_weight, jsnode, attr_pairs, is_classifier, remap): - if 'children' in jsnode: - XGBConverter._add_node(attr_pairs=attr_pairs, is_classifier=is_classifier, - tree_id=treeid, tree_weight=tree_weight, - value=jsnode['split_condition'], node_id=remap[jsnode['nodeid']], - feature_id=jsnode['split'], - mode='BRANCH_LT', # 'BRANCH_LEQ' --> is for sklearn - true_child_id=remap[jsnode['yes']], # ['children'][0]['nodeid'], - false_child_id=remap[jsnode['no']], # ['children'][1]['nodeid'], - weights=None, weight_id_bias=None, - missing=jsnode.get('missing', -1) == jsnode['yes'], # ['children'][0]['nodeid'], - hitrate=jsnode.get('cover', 0)) - - for ch in jsnode['children']: - if 'children' in ch or 'leaf' in ch: - XGBConverter._fill_node_attributes(treeid, tree_weight, ch, attr_pairs, is_classifier, remap) + def _fill_node_attributes( + treeid, tree_weight, jsnode, attr_pairs, is_classifier, remap + ): + if "children" in jsnode: + XGBConverter._add_node( + attr_pairs=attr_pairs, + is_classifier=is_classifier, + tree_id=treeid, + tree_weight=tree_weight, + value=jsnode["split_condition"], + node_id=remap[jsnode["nodeid"]], + feature_id=jsnode["split"], + mode="BRANCH_LT", # 'BRANCH_LEQ' --> is for sklearn + true_child_id=remap[jsnode["yes"]], # ['children'][0]['nodeid'], + false_child_id=remap[jsnode["no"]], # ['children'][1]['nodeid'], + weights=None, + weight_id_bias=None, + missing=jsnode.get("missing", -1) + == jsnode["yes"], # ['children'][0]['nodeid'], + hitrate=jsnode.get("cover", 0), + ) + + for ch in jsnode["children"]: + if "children" in ch or "leaf" in ch: + XGBConverter._fill_node_attributes( + treeid, tree_weight, ch, attr_pairs, is_classifier, remap + ) else: raise RuntimeError("Unable to convert this node {0}".format(ch)) else: - weights = [jsnode['leaf']] + weights = [jsnode["leaf"]] weights_id_bias = 0 - XGBConverter._add_node(attr_pairs=attr_pairs, is_classifier=is_classifier, - tree_id=treeid, tree_weight=tree_weight, - value=0., node_id=remap[jsnode['nodeid']], - feature_id=0, mode='LEAF', - true_child_id=0, false_child_id=0, - weights=weights, weight_id_bias=weights_id_bias, - missing=False, hitrate=jsnode.get('cover', 0)) + XGBConverter._add_node( + attr_pairs=attr_pairs, + is_classifier=is_classifier, + tree_id=treeid, + tree_weight=tree_weight, + value=0.0, + node_id=remap[jsnode["nodeid"]], + feature_id=0, + mode="LEAF", + true_child_id=0, + false_child_id=0, + weights=weights, + weight_id_bias=weights_id_bias, + missing=False, + hitrate=jsnode.get("cover", 0), + ) @staticmethod def _remap_nodeid(jsnode, remap=None): if remap is None: remap = {} - nid = jsnode['nodeid'] + nid = jsnode["nodeid"] remap[nid] = len(remap) - if 'children' in jsnode: - for ch in jsnode['children']: + if "children" in jsnode: + for ch in jsnode["children"]: XGBConverter._remap_nodeid(ch, remap) return remap @@ -160,11 +206,12 @@ def fill_tree_attributes(js_xgb_node, attr_pairs, tree_weights, is_classifier): raise TypeError("js_xgb_node must be a list") for treeid, (jstree, w) in enumerate(zip(js_xgb_node, tree_weights)): remap = XGBConverter._remap_nodeid(jstree) - XGBConverter._fill_node_attributes(treeid, w, jstree, attr_pairs, is_classifier, remap) + XGBConverter._fill_node_attributes( + treeid, w, jstree, attr_pairs, is_classifier, remap + ) class XGBRegressorConverter(XGBConverter): - @staticmethod def validate(xgb_node): return XGBConverter.validate(xgb_node) @@ -172,8 +219,8 @@ def validate(xgb_node): @staticmethod def _get_default_tree_attribute_pairs(): attrs = XGBConverter._get_default_tree_attribute_pairs(False) - attrs['post_transform'] = 'NONE' - attrs['n_targets'] = 1 + attrs["post_transform"] = "NONE" + attrs["n_targets"] = 1 return attrs @staticmethod @@ -186,28 +233,34 @@ def convert(scope, operator, container): raise RuntimeError("Objective '{}' not supported.".format(objective)) attr_pairs = XGBRegressorConverter._get_default_tree_attribute_pairs() - attr_pairs['base_values'] = [base_score] + attr_pairs["base_values"] = [base_score] bst = xgb_node.get_booster() - best_ntree_limit = getattr(bst, 'best_ntree_limit', len(js_trees)) + best_ntree_limit = getattr(bst, "best_ntree_limit", len(js_trees)) if best_ntree_limit < len(js_trees): js_trees = js_trees[:best_ntree_limit] - XGBConverter.fill_tree_attributes(js_trees, attr_pairs, [1 for _ in js_trees], False) + XGBConverter.fill_tree_attributes( + js_trees, attr_pairs, [1 for _ in js_trees], False + ) # add nodes - container.add_node('TreeEnsembleRegressor', operator.input_full_names, - operator.output_full_names, op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('TreeEnsembleRegressor'), **attr_pairs) - #try: + container.add_node( + "TreeEnsembleRegressor", + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("TreeEnsembleRegressor"), + **attr_pairs, + ) + # try: # if len(inputs[0].type.tensor_type.shape.dim) > 0: # output_dim = [inputs[0].type.tensor_type.shape.dim[0].dim_value, 1] - #except Exception: + # except Exception: # raise ValueError('Invalid/missing input dimension.') class XGBClassifierConverter(XGBConverter): - @staticmethod def validate(xgb_node): return XGBConverter.validate(xgb_node) @@ -215,8 +268,9 @@ def validate(xgb_node): @staticmethod def _get_default_tree_attribute_pairs(): attrs = XGBConverter._get_default_tree_attribute_pairs(True) - # TODO: check it is implemented. The model cannot be loaded when they are present. - #attrs['nodes_hitrates'] = [] + # TODO: check it is implemented. The model cannot + # be loaded when they are present. + # attrs['nodes_hitrates'] = [] return attrs @staticmethod @@ -228,99 +282,118 @@ def convert(scope, operator, container): params = XGBConverter.get_xgb_params(xgb_node) attr_pairs = XGBClassifierConverter._get_default_tree_attribute_pairs() - XGBConverter.fill_tree_attributes(js_trees, attr_pairs, [1 for _ in js_trees], True) - ncl = (max(attr_pairs['class_treeids']) + 1) // params['n_estimators'] + XGBConverter.fill_tree_attributes( + js_trees, attr_pairs, [1 for _ in js_trees], True + ) + ncl = (max(attr_pairs["class_treeids"]) + 1) // params["n_estimators"] bst = xgb_node.get_booster() - best_ntree_limit = getattr(bst, 'best_ntree_limit', len(js_trees)) * ncl + best_ntree_limit = getattr(bst, "best_ntree_limit", len(js_trees)) * ncl if 0 < best_ntree_limit < len(js_trees): js_trees = js_trees[:best_ntree_limit] attr_pairs = XGBClassifierConverter._get_default_tree_attribute_pairs() - XGBConverter.fill_tree_attributes(js_trees, attr_pairs, [1 for _ in js_trees], True) + XGBConverter.fill_tree_attributes( + js_trees, attr_pairs, [1 for _ in js_trees], True + ) - if len(attr_pairs['class_treeids']) == 0: + if len(attr_pairs["class_treeids"]) == 0: raise RuntimeError("XGBoost model is empty.") if ncl <= 1: ncl = 2 - if objective != 'binary:hinge': + if objective != "binary:hinge": # See https://github.com/dmlc/xgboost/blob/master/src/common/math.h#L23. - attr_pairs['post_transform'] = "LOGISTIC" - attr_pairs['class_ids'] = [0 for v in attr_pairs['class_treeids']] - if js_trees[0].get('leaf', None) == 0: - attr_pairs['base_values'] = [0.5] + attr_pairs["post_transform"] = "LOGISTIC" + attr_pairs["class_ids"] = [0 for v in attr_pairs["class_treeids"]] + if js_trees[0].get("leaf", None) == 0: + attr_pairs["base_values"] = [0.5] elif base_score != 0.5: - cst = - np.log(1 / np.float32(base_score) - 1.) - attr_pairs['base_values'] = [cst] + cst = -np.log(1 / np.float32(base_score) - 1.0) + attr_pairs["base_values"] = [cst] else: - attr_pairs['base_values'] = [base_score] + attr_pairs["base_values"] = [base_score] else: # See https://github.com/dmlc/xgboost/blob/master/src/common/math.h#L35. - attr_pairs['post_transform'] = "SOFTMAX" - attr_pairs['base_values'] = [base_score for n in range(ncl)] - attr_pairs['class_ids'] = [v % ncl for v in attr_pairs['class_treeids']] + attr_pairs["post_transform"] = "SOFTMAX" + attr_pairs["base_values"] = [base_score for n in range(ncl)] + attr_pairs["class_ids"] = [v % ncl for v in attr_pairs["class_treeids"]] classes = xgb_node.classes_ - if (np.issubdtype(classes.dtype, np.floating) or - np.issubdtype(classes.dtype, np.integer)): - attr_pairs['classlabels_int64s'] = classes.astype('int') + if np.issubdtype(classes.dtype, np.floating) or np.issubdtype( + classes.dtype, np.integer + ): + attr_pairs["classlabels_int64s"] = classes.astype("int") else: - classes = np.array([s.encode('utf-8') for s in classes]) - attr_pairs['classlabels_strings'] = classes + classes = np.array([s.encode("utf-8") for s in classes]) + attr_pairs["classlabels_strings"] = classes # add nodes if objective in ("binary:logistic", "binary:hinge"): ncl = 2 if objective == "binary:hinge": - attr_pairs['post_transform'] = 'NONE' - output_names = [operator.output_full_names[0], - scope.get_unique_variable_name("output_prob")] + attr_pairs["post_transform"] = "NONE" + output_names = [ + operator.output_full_names[0], + scope.get_unique_variable_name("output_prob"), + ] else: output_names = operator.output_full_names - container.add_node('TreeEnsembleClassifier', - operator.input_full_names, - output_names, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('TreeEnsembleClassifier'), - **attr_pairs) + container.add_node( + "TreeEnsembleClassifier", + operator.input_full_names, + output_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("TreeEnsembleClassifier"), + **attr_pairs, + ) if objective == "binary:hinge": if container.target_opset < 9: raise RuntimeError( f"hinge function cannot be implemented because " - f"opset={container.target_opset}<9.") + f"opset={container.target_opset}<9." + ) zero = scope.get_unique_variable_name("zero") one = scope.get_unique_variable_name("one") - container.add_initializer(zero, TensorProto.FLOAT, [1], [0.]) - container.add_initializer(one, TensorProto.FLOAT, [1], [1.]) + container.add_initializer(zero, TensorProto.FLOAT, [1], [0.0]) + container.add_initializer(one, TensorProto.FLOAT, [1], [1.0]) greater = scope.get_unique_variable_name("output_prob") container.add_node("Greater", [output_names[1], zero], [greater]) - container.add_node('Where', [greater, one, zero], - operator.output_full_names[1]) + container.add_node( + "Where", [greater, one, zero], operator.output_full_names[1] + ) elif objective in ("multi:softprob", "multi:softmax"): - ncl = len(js_trees) // params['n_estimators'] - if objective == 'multi:softmax': - attr_pairs['post_transform'] = 'NONE' - container.add_node('TreeEnsembleClassifier', operator.input_full_names, - operator.output_full_names, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('TreeEnsembleClassifier'), - **attr_pairs) + ncl = len(js_trees) // params["n_estimators"] + if objective == "multi:softmax": + attr_pairs["post_transform"] = "NONE" + container.add_node( + "TreeEnsembleClassifier", + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("TreeEnsembleClassifier"), + **attr_pairs, + ) elif objective == "reg:logistic": - ncl = len(js_trees) // params['n_estimators'] + ncl = len(js_trees) // params["n_estimators"] if ncl == 1: ncl = 2 - container.add_node('TreeEnsembleClassifier', operator.input_full_names, - operator.output_full_names, - op_domain='ai.onnx.ml', - name=scope.get_unique_operator_name('TreeEnsembleClassifier'), - **attr_pairs) + container.add_node( + "TreeEnsembleClassifier", + operator.input_full_names, + operator.output_full_names, + op_domain="ai.onnx.ml", + name=scope.get_unique_operator_name("TreeEnsembleClassifier"), + **attr_pairs, + ) else: raise RuntimeError("Unexpected objective: {0}".format(objective)) def convert_xgboost(scope, operator, container): xgb_node = operator.raw_operator - if (isinstance(xgb_node, XGBClassifier) or - getattr(xgb_node, 'operator_name', None) == 'XGBClassifier'): + if ( + isinstance(xgb_node, XGBClassifier) + or getattr(xgb_node, "operator_name", None) == "XGBClassifier" + ): cls = XGBClassifierConverter else: cls = XGBRegressorConverter @@ -328,5 +401,5 @@ def convert_xgboost(scope, operator, container): cls.convert(scope, operator, container) -register_converter('XGBClassifier', convert_xgboost) -register_converter('XGBRegressor', convert_xgboost) +register_converter("XGBClassifier", convert_xgboost) +register_converter("XGBRegressor", convert_xgboost) diff --git a/onnxmltools/convert/xgboost/shape_calculators/Classifier.py b/onnxmltools/convert/xgboost/shape_calculators/Classifier.py index f3bb88d1..55d1f430 100644 --- a/onnxmltools/convert/xgboost/shape_calculators/Classifier.py +++ b/onnxmltools/convert/xgboost/shape_calculators/Classifier.py @@ -4,37 +4,41 @@ from ...common._registration import register_shape_calculator from ...common.utils import check_input_and_output_numbers, check_input_and_output_types from ...common.data_types import ( - DictionaryType, FloatTensorType, Int64TensorType, - SequenceType, StringTensorType, + FloatTensorType, + Int64TensorType, + StringTensorType, ) from ..common import get_xgb_params def calculate_xgboost_classifier_output_shapes(operator): check_input_and_output_numbers(operator, input_count_range=1, output_count_range=2) - check_input_and_output_types(operator, good_input_types=[FloatTensorType, Int64TensorType]) + check_input_and_output_types( + operator, good_input_types=[FloatTensorType, Int64TensorType] + ) N = operator.inputs[0].type.shape[0] xgb_node = operator.raw_operator params = get_xgb_params(xgb_node) booster = xgb_node.get_booster() - atts = booster.attributes() - ntrees = len(booster.get_dump(with_stats=True, dump_format = 'json')) + booster.attributes() + ntrees = len(booster.get_dump(with_stats=True, dump_format="json")) objective = params["objective"] if objective == "binary:logistic": ncl = 2 else: - ncl = ntrees // params['n_estimators'] + ncl = ntrees // params["n_estimators"] if objective == "reg:logistic" and ncl == 1: ncl = 2 classes = xgb_node.classes_ - if (np.issubdtype(classes.dtype, np.floating) or - np.issubdtype(classes.dtype, np.integer)): + if np.issubdtype(classes.dtype, np.floating) or np.issubdtype( + classes.dtype, np.integer + ): operator.outputs[0].type = Int64TensorType(shape=[N]) else: operator.outputs[0].type = StringTensorType(shape=[N]) operator.outputs[1].type = operator.outputs[1].type = FloatTensorType([N, ncl]) -register_shape_calculator('XGBClassifier', calculate_xgboost_classifier_output_shapes) +register_shape_calculator("XGBClassifier", calculate_xgboost_classifier_output_shapes) diff --git a/onnxmltools/convert/xgboost/shape_calculators/Regressor.py b/onnxmltools/convert/xgboost/shape_calculators/Regressor.py index 31c77f30..31743ddc 100644 --- a/onnxmltools/convert/xgboost/shape_calculators/Regressor.py +++ b/onnxmltools/convert/xgboost/shape_calculators/Regressor.py @@ -3,4 +3,4 @@ from ...common._registration import register_shape_calculator from ...common.shape_calculator import calculate_linear_regressor_output_shapes -register_shape_calculator('XGBRegressor', calculate_linear_regressor_output_shapes) +register_shape_calculator("XGBRegressor", calculate_linear_regressor_output_shapes) diff --git a/onnxmltools/proto/__init__.py b/onnxmltools/proto/__init__.py index ff24c446..a6bd5faf 100644 --- a/onnxmltools/proto/__init__.py +++ b/onnxmltools/proto/__init__.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 -# Rather than using ONNX protobuf definition throughout our codebase, we import ONNX protobuf definition here so that -# we can conduct quick fixes by overwriting ONNX functions without changing any lines elsewhere. +# Rather than using ONNX protobuf definition throughout our codebase, +# we import ONNX protobuf definition here so that +# we can conduct quick fixes by overwriting ONNX functions without +# changing any lines elsewhere. from onnx import onnx_pb as onnx_proto # noqa from onnx import helper -# Overwrite the make_tensor defined in onnx.helper because of a bug (string tensor get assigned twice) +# Overwrite the make_tensor defined in onnx.helper because of +# a bug (string tensor get assigned twice) from onnx import mapping from onnx.onnx_pb import TensorProto from onnx.helper import split_complex_to_pairs @@ -13,34 +16,37 @@ def _check_onnx_version(): import pkg_resources - min_required_version = pkg_resources.parse_version('1.0.1') - current_version = pkg_resources.get_distribution('onnx').parsed_version - assert current_version >= min_required_version, 'ONNXMLTools requires ONNX version 1.0.1 or a newer one' + + min_required_version = pkg_resources.parse_version("1.0.1") + current_version = pkg_resources.get_distribution("onnx").parsed_version + assert ( + current_version >= min_required_version + ), "ONNXMLTools requires ONNX version 1.0.1 or a newer one" _check_onnx_version() def make_tensor_fixed(name, data_type, dims, vals, raw=False): - ''' + """ Make a TensorProto with specified arguments. If raw is False, this function will choose the corresponding proto field to store the values based on data_type. If raw is True, use "raw_data" proto field to store the values, and values should be of type bytes in this case. - ''' + """ tensor = TensorProto() tensor.data_type = data_type tensor.name = name - if (data_type == TensorProto.COMPLEX64 or - data_type == TensorProto.COMPLEX128): + if data_type == TensorProto.COMPLEX64 or data_type == TensorProto.COMPLEX128: vals = split_complex_to_pairs(vals) if raw: tensor.raw_data = vals else: field = mapping.STORAGE_TENSOR_TYPE_TO_FIELD[ - mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type]] + mapping.TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE[data_type] + ] getattr(tensor, field).extend(vals) tensor.dims.extend(dims) diff --git a/onnxmltools/utils/__init__.py b/onnxmltools/utils/__init__.py index 7346597f..47253267 100644 --- a/onnxmltools/utils/__init__.py +++ b/onnxmltools/utils/__init__.py @@ -9,6 +9,10 @@ from .visualize import visualize_model from .float16_converter import convert_float_to_float16 from .tests_helper import dump_data_and_model -from .tests_helper import dump_one_class_classification, dump_binary_classification, dump_multiple_classification +from .tests_helper import ( + dump_one_class_classification, + dump_binary_classification, + dump_multiple_classification, +) from .tests_helper import dump_multiple_regression, dump_single_regression from .tests_dl_helper import create_tensor diff --git a/onnxmltools/utils/main.py b/onnxmltools/utils/main.py index 00d8af3f..c7f54c3c 100644 --- a/onnxmltools/utils/main.py +++ b/onnxmltools/utils/main.py @@ -23,7 +23,7 @@ def load_model(file_path): if not path.exists(file_path): raise FileNotFoundError("{0} was not found.".format(file_path)) model = onnx_proto.ModelProto() - with open(file_path, 'rb') as f: + with open(file_path, "rb") as f: model.ParseFromString(f.read()) return model @@ -46,7 +46,7 @@ def save_model(model, file_path): directory = os.path.dirname(os.path.abspath(file_path)) if not path.exists(directory): raise FileNotFoundError("Directory does not exist {0}".format(directory)) - with open(file_path, 'wb') as f: + with open(file_path, "wb") as f: f.write(model.SerializeToString()) @@ -113,5 +113,7 @@ def set_model_doc_string(model, doc, override=False): raise ValueError("Doc must be a string type.") if model.doc_string and not doc and override is False: raise ValueError( - "Failing to overwrite the doc string with a blank string, set override to True if intentional.") + "Failing to overwrite the doc string with a " + "blank string, set override to True if intentional." + ) model.doc_string = doc diff --git a/onnxmltools/utils/tests_dl_helper.py b/onnxmltools/utils/tests_dl_helper.py index 5003c7c3..9e78305f 100644 --- a/onnxmltools/utils/tests_dl_helper.py +++ b/onnxmltools/utils/tests_dl_helper.py @@ -20,15 +20,18 @@ def find_inference_engine(): return _rt_installed try: - import onnxruntime + pass + _rt_installed = rt_onnxruntime except ImportError: try: - import cntk + pass + _rt_installed = rt_cntk except ImportError: try: - import caffe2 + pass + _rt_installed = rt_caffe2 except ImportError: pass @@ -50,16 +53,25 @@ def evaluate_deep_model(onnx_model, inputs, rt_type=None): elif rt_type == rt_caffe2: return _evaluate_caffe2(onnx_model, inputs) else: - raise ImportError('No runtime found. Need either CNTK or Caffe2') + raise ImportError("No runtime found. Need either CNTK or Caffe2") def _evaluate_onnxruntime(onnx_model, inputs): import onnxruntime + runtime = onnxruntime.InferenceSession(onnx_model.SerializeToString()) result = None inputs = inputs if isinstance(inputs, list) else [inputs] - for i_ in range(inputs[0].shape[0]): # TODO: onnxruntime can't support batch_size > 1 - out = runtime.run([], {x.name: inputs[n_][i_:i_ + 1] for n_, x in enumerate(runtime.get_inputs())})[0] + for i_ in range( + inputs[0].shape[0] + ): # TODO: onnxruntime can't support batch_size > 1 + out = runtime.run( + [], + { + x.name: inputs[n_][i_ : i_ + 1] + for n_, x in enumerate(runtime.get_inputs()) + }, + )[0] result = out if result is None else np.concatenate((result, out)) return result[0] if isinstance(result, list) else result @@ -82,6 +94,7 @@ def _evaluate_caffe2(onnx_model, inputs): def _evaluate_cntk(onnx_model, inputs): import cntk + if not isinstance(inputs, list): inputs = [inputs] @@ -89,11 +102,15 @@ def _evaluate_cntk(onnx_model, inputs): for i, x in enumerate(inputs): onnx_name = onnx_model.graph.input[i].name - adjusted_inputs[onnx_name] = [np.ascontiguousarray(np.squeeze(_, axis=0)) for _ in np.split(x, x.shape[0])] + adjusted_inputs[onnx_name] = [ + np.ascontiguousarray(np.squeeze(_, axis=0)) for _ in np.split(x, x.shape[0]) + ] - temporary_onnx_model_file_name = 'temp_' + onnx_model.graph.name + '.onnx' + temporary_onnx_model_file_name = "temp_" + onnx_model.graph.name + ".onnx" save_model(onnx_model, temporary_onnx_model_file_name) - cntk_model = cntk.Function.load(temporary_onnx_model_file_name, format=cntk.ModelFormat.ONNX) + cntk_model = cntk.Function.load( + temporary_onnx_model_file_name, format=cntk.ModelFormat.ONNX + ) return cntk_model.eval(adjusted_inputs) @@ -104,4 +121,4 @@ def create_tensor(N, C, H=None, W=None): elif H is not None and W is not None: return np.random.rand(N, C, H, W).astype(np.float32, copy=False) else: - raise ValueError('This function only produce 2-D or 4-D tensor') + raise ValueError("This function only produce 2-D or 4-D tensor") diff --git a/onnxmltools/utils/tests_helper.py b/onnxmltools/utils/tests_helper.py index 2c33ff7e..7b5fd1d4 100644 --- a/onnxmltools/utils/tests_helper.py +++ b/onnxmltools/utils/tests_helper.py @@ -2,9 +2,7 @@ import pickle import os -import warnings import numpy -import packaging.version as pv from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from ..convert.common.data_types import FloatTensorType @@ -14,9 +12,18 @@ TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) -def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, - inputs=None, backend="onnxruntime", context=None, - allow_failure=None, verbose=False): +def dump_data_and_model( + data, + model, + onnx=None, + basename="model", + folder=None, + inputs=None, + backend="onnxruntime", + context=None, + allow_failure=None, + verbose=False, +): """ Saves data with pickle, saves the model with pickle and *onnx*, runs and saves the predictions for the given model. @@ -55,14 +62,17 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, * ``-CannotLoad``: the model can be converted but the runtime cannot load it * ``-Dec3``: compares expected and computed outputs up to 3 decimals (5 by default) * ``-Dec4``: compares expected and computed outputs up to 4 decimals (5 by default) - * ``-NoProb``: The original models computed probabilites for two classes *size=(N, 2)* - but the runtime produces a vector of size *N*, the test will compare the second column + * ``-NoProb``: The original models computed + probabilites for two classes *size=(N, 2)* + but the runtime produces a vector of size *N*, + the test will compare the second column to the column * ``-OneOff``: the ONNX runtime cannot computed the prediction for several inputs, it must be called for each of them and computed output. * ``-Out0``: only compares the first output on both sides - * ``-Reshape``: merges all outputs into one single vector and resizes it before comparing + * ``-Reshape``: merges all outputs into one single vector and resizes + it before comparing * ``-SkipDim1``: before comparing expected and computed output, arrays with a shape like *(2, 1, 2)* becomes *(2, 2)* @@ -73,23 +83,24 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, runtime_test = dict(model=model, data=data) if folder is None: - folder = os.environ.get('ONNXTESTDUMP', 'tests/temp') + folder = os.environ.get("ONNXTESTDUMP", "tests/temp") if not os.path.exists(folder): os.makedirs(folder) if hasattr(model, "predict"): import lightgbm import xgboost + if isinstance(model, lightgbm.Booster): # LightGBM Booster model_dict = model.dump_model() - if model_dict['objective'].startswith('binary'): + if model_dict["objective"].startswith("binary"): score = model.predict(data) if len(score.shape) < 2 or score.shape[1] == 1: score = score.ravel() - score = numpy.vstack([1-score, score]).T + score = numpy.vstack([1 - score, score]).T prediction = [score[:, 1] > 0.5, score] - elif model_dict['objective'].startswith('multiclass'): + elif model_dict["objective"].startswith("multiclass"): score = model.predict(data) prediction = [score.argmax(axis=1), score] else: @@ -98,25 +109,26 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, # XGBoost Booster from ..convert.xgboost._parse import _get_attributes from xgboost import DMatrix + datax = DMatrix(data) model_dict = _get_attributes(model) - if model_dict['objective'].startswith('binary'): + if model_dict["objective"].startswith("binary"): score = model.predict(datax) - prediction = [score > 0.5, numpy.vstack([1-score, score]).T] - elif model_dict['objective'].startswith('multi:softprob'): + prediction = [score > 0.5, numpy.vstack([1 - score, score]).T] + elif model_dict["objective"].startswith("multi:softprob"): score = model.predict(datax) prediction = [score.argmax(axis=1), score] - elif model_dict['objective'].startswith('multi:softmax'): + elif model_dict["objective"].startswith("multi:softmax"): score = model.predict(datax, output_margin=True) prediction = [score.argmax(axis=1), score] else: prediction = [model.predict(datax)] elif hasattr(model, "predict_proba"): # Classifier - if hasattr(model, 'get_params'): + if hasattr(model, "get_params"): params = model.get_params() - if 'objective' in params: - objective = params['objective'] + if "objective" in params: + objective = params["objective"] if objective == "multi:softmax": prediction = [model.predict(data)] else: @@ -134,7 +146,9 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, elif hasattr(model, "layers"): # Keras if len(model.input_names) != 1: - raise NotImplemented("Only neural network with one input are supported") + raise NotImplementedError( + "Only neural network with one input are supported" + ) prediction = [model.predict(data)] else: # Regressor @@ -142,9 +156,11 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, elif hasattr(model, "transform"): prediction = model.transform(data) else: - raise TypeError("Model has not predict or transform method: {0}".format(type(model))) + raise TypeError( + "Model has not predict or transform method: {0}".format(type(model)) + ) - runtime_test['expected'] = prediction + runtime_test["expected"] = prediction names = [] dest = os.path.join(folder, basename + ".expected.pkl") @@ -165,7 +181,7 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, if onnx is None: array = numpy.array(data) if inputs is None: - inputs = [('input', FloatTensorType(list(array.shape)))] + inputs = [("input", FloatTensorType(list(array.shape)))] onnx, _ = convert_model(model, basename, inputs) dest = os.path.join(folder, basename + ".model.onnx") @@ -184,8 +200,13 @@ def dump_data_and_model(data, model, onnx=None, basename="model", folder=None, continue if isinstance(allow_failure, str): raise NotImplementedError("allow_failure is deprecated.") - output = compare_backend(b, runtime_test, options=extract_options(basename), - context=context, verbose=verbose) + output = compare_backend( + b, + runtime_test, + options=extract_options(basename), + context=context, + verbose=verbose, + ) if output is not None: dest = os.path.join(folder, basename + ".backend.{0}.pkl".format(b)) names.append(dest) @@ -203,27 +224,47 @@ def convert_model(model, name, input_types, without_onnx_ml=False, **kwargs): :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, without_onnx_ml=without_onnx_ml, **kwargs), "LightGbm" + + model, prefix = ( + convert_lightgbm( + model, name, input_types, without_onnx_ml=without_onnx_ml, **kwargs + ), + "LightGbm", + ) elif model.__class__.__name__.startswith("XGB"): from onnxmltools.convert import convert_xgboost + model, prefix = convert_xgboost(model, name, input_types, **kwargs), "XGB" - elif model.__class__.__name__ == 'Booster': + 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, without_onnx_ml=without_onnx_ml, **kwargs), "LightGbm" + + model, prefix = ( + convert_lightgbm( + model, name, input_types, without_onnx_ml=without_onnx_ml, **kwargs + ), + "LightGbm", + ) else: - raise RuntimeError("Unable to convert model of type '{0}'.".format(type(model))) + raise RuntimeError( + "Unable to convert model of type '{0}'.".format(type(model)) + ) elif model.__class__.__name__.startswith("CatBoost"): from onnxmltools.convert import convert_catboost + model, prefix = convert_catboost(model, name, input_types, **kwargs), "CatBoost" elif isinstance(model, BaseEstimator): from onnxmltools.convert import convert_sklearn + model, prefix = convert_sklearn(model, name, input_types, **kwargs), "Sklearn" else: from onnxmltools.convert import convert_coreml + model, prefix = convert_coreml(model, name, input_types, **kwargs), "Cml" if model is None: raise RuntimeError("Unable to convert model of type '{0}'.".format(type(model))) @@ -248,17 +289,29 @@ def dump_one_class_classification(model, suffix="", folder=None, allow_failure=N Every created filename will follow the pattern: ``/..``. """ - X = [[0., 1.], [1., 1.], [2., 0.]] + X = [[0.0, 1.0], [1.0, 1.0], [2.0, 0.0]] X = numpy.array(X, dtype=numpy.float32) y = [1, 1, 1] model.fit(X, y) - model_onnx, prefix = convert_model(model, 'one_class', [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - return dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure, - basename=prefix + "One" + model.__class__.__name__ + suffix) - - -def dump_binary_classification(model, suffix="", folder=None, allow_failure=None, verbose=False): + model_onnx, prefix = convert_model( + model, + "one_class", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + return dump_data_and_model( + X, + model, + model_onnx, + folder=folder, + allow_failure=allow_failure, + basename=prefix + "One" + model.__class__.__name__ + suffix, + ) + + +def dump_binary_classification( + model, suffix="", folder=None, allow_failure=None, verbose=False +): """ Trains and dumps a model for a binary classification problem. @@ -279,11 +332,22 @@ def dump_binary_classification(model, suffix="", folder=None, allow_failure=None X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0] model.fit(X, y) - model_onnx, prefix = convert_model(model, 'tree-based binary classifier', [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure, - basename=prefix + "Bin" + model.__class__.__name__ + suffix, - verbose=verbose) + model_onnx, prefix = convert_model( + model, + "tree-based binary classifier", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + folder=folder, + allow_failure=allow_failure, + basename=prefix + "Bin" + model.__class__.__name__ + suffix, + verbose=verbose, + ) + def dump_multiple_classification(model, suffix="", folder=None, allow_failure=None): """ @@ -305,10 +369,20 @@ def dump_multiple_classification(model, suffix="", folder=None, allow_failure=No X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 2, 1, 1, 2] model.fit(X, y) - model_onnx, prefix = convert_model(model, 'tree-based multi-output regressor', [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure, - basename=prefix + "Mcl" + model.__class__.__name__ + suffix) + model_onnx, prefix = convert_model( + model, + "tree-based multi-output regressor", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + folder=folder, + allow_failure=allow_failure, + basename=prefix + "Mcl" + model.__class__.__name__ + suffix, + ) def dump_multiple_regression(model, suffix="", folder=None, allow_failure=None): @@ -331,10 +405,20 @@ def dump_multiple_regression(model, suffix="", folder=None, allow_failure=None): X = numpy.array(X, dtype=numpy.float32) y = numpy.array([[100, 50], [100, 49], [100, 99]], dtype=numpy.float32) model.fit(X, y) - model_onnx, prefix = convert_model(model, 'tree-based multi-output regressor', [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure, - basename=prefix + "MRg" + model.__class__.__name__ + suffix) + model_onnx, prefix = convert_model( + model, + "tree-based multi-output regressor", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + folder=folder, + allow_failure=allow_failure, + basename=prefix + "MRg" + model.__class__.__name__ + suffix, + ) def dump_single_regression(model, suffix="", folder=None, allow_failure=None): @@ -358,10 +442,20 @@ def dump_single_regression(model, suffix="", folder=None, allow_failure=None): X = numpy.array(X, dtype=numpy.float32) y = numpy.array([100, -10, 50], dtype=numpy.float32) model.fit(X, y) - model_onnx, prefix = convert_model(model, 'tree-based regressor', [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, folder=folder, allow_failure=allow_failure, - basename=prefix + "Reg" + model.__class__.__name__ + suffix) + model_onnx, prefix = convert_model( + model, + "tree-based regressor", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + folder=folder, + allow_failure=allow_failure, + basename=prefix + "Reg" + model.__class__.__name__ + suffix, + ) def make_report_backend(folder): @@ -377,7 +471,7 @@ def make_report_backend(folder): if model not in res: res[model] = {} res[model]["_tested"] = True - elif '.backend.' in name: + elif ".backend." in name: bk = name.split(".backend.")[-1].split(".")[0] model = name.split(".")[0] if model not in res: diff --git a/onnxmltools/utils/utils_backend.py b/onnxmltools/utils/utils_backend.py index 70ef0fd0..84dbe146 100644 --- a/onnxmltools/utils/utils_backend.py +++ b/onnxmltools/utils/utils_backend.py @@ -6,7 +6,6 @@ import os import glob import pickle -import packaging.version as pv import numpy from numpy.testing import assert_array_almost_equal, assert_array_equal @@ -15,14 +14,12 @@ class ExpectedAssertionError(Exception): """ Expected failure. """ - pass class OnnxRuntimeAssertionError(AssertionError): """ Expected failure. """ - pass def is_backend_enabled(backend): @@ -31,7 +28,8 @@ def is_backend_enabled(backend): """ if backend == "onnxruntime": try: - import onnxruntime + pass + return True except ImportError: return False @@ -39,7 +37,9 @@ def is_backend_enabled(backend): raise NotImplementedError("Not implemented for backend '{0}'".format(backend)) -def compare_backend(backend, test, decimal=5, options=None, verbose=False, context=None): +def compare_backend( + backend, test, decimal=5, options=None, verbose=False, context=None +): """ The function compares the expected output (computed with the model before being converted to ONNX) and the ONNX output @@ -62,6 +62,7 @@ def compare_backend(backend, test, decimal=5, options=None, verbose=False, conte """ if backend == "onnxruntime": from .utils_backend_onnxruntime import compare_runtime + return compare_runtime(test, decimal, options, verbose) raise ValueError("Does not support backend '{0}'.".format(backend)) @@ -82,7 +83,7 @@ def search_converted_models(root=None): keep = [] for found in founds: onnx = found - basename = onnx[:-len(".model.onnx")] + basename = onnx[: -len(".model.onnx")] data = basename + ".data.pkl" expected = basename + ".expected.pkl" res = dict(onnx=onnx, data=data, expected=expected) @@ -94,9 +95,9 @@ def search_converted_models(root=None): models = [basename + ".model.pkl", basename + ".model.keras"] for model in models: if os.path.exists(model): - res['model'] = model + res["model"] = model break - if 'model' in res: + if "model" in res: keep.append((basename, res)) keep.sort() return [_[1] for _ in keep] @@ -117,13 +118,16 @@ def load_data_and_model(items_as_dict, **context): try: bin = pickle.load(f) except ImportError as e: - if '.model.' in v: + if ".model." in v: continue else: - raise ImportError("Unable to load '{0}' due to {1}".format(v, e)) + raise ImportError( + "Unable to load '{0}' due to {1}".format(v, e) + ) res[k] = bin elif os.path.splitext(v)[-1] == ".keras": import keras.models + res[k] = keras.models.load_model(v, custom_objects=context) else: res[k] = v @@ -141,19 +145,30 @@ def extract_options(name): Available options: * `'SkipDim1'`: reshape arrays by skipping 1-dimension: ``(1, 2)`` --> ``(2,)`` - * `'OneOff'`: inputs comes in a list for the predictions are computed with a call for each of them, - not with one call + * `'OneOff'`: inputs comes in a list for the predictions + are computed with a call for each of them, + not with one call * ... See function *dump_data_and_model* to get the full list. """ - opts = name.replace("\\", "/").split("/")[-1].split('.')[0].split('-') + opts = name.replace("\\", "/").split("/")[-1].split(".")[0].split("-") if len(opts) == 1: return {} else: res = {} for opt in opts[1:]: - if opt in ("SkipDim1", "OneOff", "NoProb", "Dec4", "Dec3", 'Out0', 'Dec2', 'Reshape', 'Opp'): + if opt in ( + "SkipDim1", + "OneOff", + "NoProb", + "Dec4", + "Dec3", + "Out0", + "Dec2", + "Reshape", + "Opp", + ): res[opt] = True else: raise NameError("Unable to parse option '{}'".format(opts[1:])) @@ -192,20 +207,34 @@ def compare_outputs(expected, output, **kwargs): # One vector is (N,) with scores, negative for class 0 # positive for class 1 # The other vector is (N, 2) score in two columns. - if len(output.shape) == 2 and output.shape[1] == 2 and len(expected.shape) == 1: + if ( + len(output.shape) == 2 + and output.shape[1] == 2 + and len(expected.shape) == 1 + ): output = output[:, 1] elif len(output.shape) == 1 and len(expected.shape) == 1: pass - elif len(expected.shape) == 1 and len(output.shape) == 2 and \ - expected.shape[0] == output.shape[0] and output.shape[1] == 1: + elif ( + len(expected.shape) == 1 + and len(output.shape) == 2 + and expected.shape[0] == output.shape[0] + and output.shape[1] == 1 + ): output = output[:, 0] elif expected.shape != output.shape: - raise NotImplementedError("No good shape: {0} != {1}".format(expected.shape, output.shape)) + raise NotImplementedError( + "No good shape: {0} != {1}".format(expected.shape, output.shape) + ) if Opp: output = -output if len(expected.shape) == 1 and len(output.shape) == 2 and output.shape[1] == 1: output = output.ravel() - if len(expected.shape) == 2 and len(output.shape) == 1 and expected.shape[1] == 1: + if ( + len(expected.shape) == 2 + and len(output.shape) == 1 + and expected.shape[1] == 1 + ): expected = expected.ravel() if not numpy.issubdtype(expected.dtype, numpy.number): try: @@ -225,14 +254,28 @@ def compare_outputs(expected, output, **kwargs): if len(expected_) == len(output_): diff = numpy.abs(expected_ - output_).max() elif Mism: - return ExpectedAssertionError("dimension mismatch={0}, {1}\n{2}".format(expected.shape, output.shape, e)) + return ExpectedAssertionError( + "dimension mismatch={0}, {1}\n{2}".format( + expected.shape, output.shape, e + ) + ) else: - return OnnxRuntimeAssertionError("dimension mismatch={0}, {1}\n{2}".format(expected.shape, output.shape, e)) + return OnnxRuntimeAssertionError( + "dimension mismatch={0}, {1}\n{2}".format( + expected.shape, output.shape, e + ) + ) if Disc: # Bug to be fixed later. - return ExpectedAssertionError("max diff(expected, output)={0}\n{1}".format(diff, e)) + return ExpectedAssertionError( + "max diff(expected, output)={0}\n{1}".format(diff, e) + ) else: - return OnnxRuntimeAssertionError("max diff(expected, output)={0}\n{1}".format(diff, e)) + return OnnxRuntimeAssertionError( + "max diff(expected, output)={0}\n{1}".format(diff, e) + ) else: - return OnnxRuntimeAssertionError("Unexpected types {0} != {1}".format(type(expected), type(output))) + return OnnxRuntimeAssertionError( + "Unexpected types {0} != {1}".format(type(expected), type(output)) + ) return None diff --git a/onnxmltools/utils/utils_backend_onnxruntime.py b/onnxmltools/utils/utils_backend_onnxruntime.py index bb755a78..edb6fea7 100644 --- a/onnxmltools/utils/utils_backend_onnxruntime.py +++ b/onnxmltools/utils/utils_backend_onnxruntime.py @@ -3,14 +3,16 @@ """ Helpers to test runtimes. """ -import os -import glob -import pickle import warnings import numpy -from numpy.testing import assert_array_almost_equal, assert_array_equal -from .utils_backend import load_data_and_model, extract_options, ExpectedAssertionError, OnnxRuntimeAssertionError, compare_outputs +from .utils_backend import ( + load_data_and_model, + extract_options, + ExpectedAssertionError, + OnnxRuntimeAssertionError, + compare_outputs, +) def compare_runtime(test, decimal=5, options=None, verbose=False, context=None): @@ -36,7 +38,7 @@ def compare_runtime(test, decimal=5, options=None, verbose=False, context=None): context = {} load = load_data_and_model(test, **context) - onx = test['onnx'] + onx = test["onnx"] if options is None: if isinstance(onx, str): options = extract_options(onx) @@ -47,7 +49,7 @@ def compare_runtime(test, decimal=5, options=None, verbose=False, context=None): try: import onnxruntime - except ImportError as e: + except ImportError: warnings.warn("Unable to import onnxruntime.") return @@ -57,17 +59,20 @@ def compare_runtime(test, decimal=5, options=None, verbose=False, context=None): raise expe except Exception as e: if "CannotLoad" in options: - raise ExpectedAssertionError("Unable to load onnx '{0}' due to\n{1}".format(onx, e)) + raise ExpectedAssertionError( + "Unable to load onnx '{0}' due to\n{1}".format(onx, e) + ) else: if verbose: import onnx + model = onnx.load(onx) smodel = "\nJSON ONNX\n" + str(model) else: smodel = "" raise OnnxRuntimeAssertionError( - "Unable to load onnx '{0}' due to {1}\nONNX\n{2}".format( - onx, e, smodel)) + "Unable to load onnx '{0}' due to {1}\nONNX\n{2}".format(onx, e, smodel) + ) input = load["data"] if isinstance(input, dict): @@ -83,27 +88,47 @@ def compare_runtime(test, decimal=5, options=None, verbose=False, context=None): if shape == input.shape[1]: inputs = {n.name: input[:, i] for i, n in enumerate(inp)} else: - raise OnnxRuntimeAssertionError("Wrong number of inputs onnx {0} != original shape {1}, onnx='{2}'".format(len(inp), input.shape, onnx)) + raise OnnxRuntimeAssertionError( + "Wrong number of inputs onnx {0} " + "!= original shape {1}, onnx='{2}'".format( + len(inp), input.shape, onnx + ) + ) elif isinstance(input, list): try: array_input = numpy.array(input) - except Exception as e: - raise OnnxRuntimeAssertionError("Wrong number of inputs onnx {0} != original {1}, onnx='{2}'".format(len(inp), len(input), onnx)) + except Exception: + raise OnnxRuntimeAssertionError( + "Wrong number of inputs onnx {0} != original {1}, " + "onnx='{2}'".format(len(inp), len(input), onnx) + ) shape = sum(i.shape[1] for i in inp) if shape == array_input.shape[1]: - inputs = {n.name: _create_column([row[i] for row in input], n.type) for i, n in enumerate(inp)} + inputs = { + n.name: _create_column([row[i] for row in input], n.type) + for i, n in enumerate(inp) + } else: - raise OnnxRuntimeAssertionError("Wrong number of inputs onnx {0} != original shape {1}, onnx='{2}'*".format(len(inp), array_input.shape, onnx)) + raise OnnxRuntimeAssertionError( + "Wrong number of inputs onnx {0} != original " + "shape {1}, onnx='{2}'*".format(len(inp), array_input.shape, onnx) + ) else: - raise OnnxRuntimeAssertionError("Wrong number of inputs onnx {0} != original {1}, onnx='{2}'".format(len(inp), len(input), onnx)) + raise OnnxRuntimeAssertionError( + "Wrong number of inputs onnx {0} != original {1}, onnx='{2}'".format( + len(inp), len(input), onnx + ) + ) else: - raise OnnxRuntimeAssertionError("Dict or list is expected, not {0}".format(type(input))) + raise OnnxRuntimeAssertionError( + "Dict or list is expected, not {0}".format(type(input)) + ) for k in inputs: if isinstance(inputs[k], list): inputs[k] = numpy.array(inputs[k]) - OneOff = options.pop('OneOff', False) + OneOff = options.pop("OneOff", False) if OneOff: if len(inputs) == 1: name, values = list(inputs.items())[0] @@ -114,15 +139,19 @@ def compare_runtime(test, decimal=5, options=None, verbose=False, context=None): except ExpectedAssertionError as expe: raise expe except Exception as e: - raise OnnxRuntimeAssertionError("Unable to run onnx '{0}' due to {1}".format(onnx, e)) + raise OnnxRuntimeAssertionError( + "Unable to run onnx '{0}' due to {1}".format(onnx, e) + ) res.append(one) output = _post_process_output(res) else: + def to_array(vv): if isinstance(vv, (numpy.ndarray, numpy.int64, numpy.float32)): return numpy.array([vv]) else: return numpy.array([vv], dtype=numpy.float32) + t = list(inputs.items())[0] res = [] for i in range(0, len(t[1])): @@ -132,7 +161,9 @@ def to_array(vv): except ExpectedAssertionError as expe: raise expe except Exception as e: - raise OnnxRuntimeAssertionError("Unable to run onnx '{0}' due to {1}".format(onx, e)) + raise OnnxRuntimeAssertionError( + "Unable to run onnx '{0}' due to {1}".format(onx, e) + ) res.append(one) output = _post_process_output(res) else: @@ -142,26 +173,41 @@ def to_array(vv): raise expe except RuntimeError as e: if "-Fail" in onx: - raise ExpectedAssertionError("onnxruntime cannot compute the prediction for '{0}'".format(onx)) + raise ExpectedAssertionError( + "onnxruntime cannot compute the prediction for '{0}'".format(onx) + ) else: - raise OnnxRuntimeAssertionError("onnxruntime cannot compute the prediction for '{0}' due to {1}".format(onx, e)) + raise OnnxRuntimeAssertionError( + "onnxruntime cannot compute the prediction for '{0}' due to {1}".format( + onx, e + ) + ) except Exception as e: - raise OnnxRuntimeAssertionError("Unable to run onnx '{0}' due to {1}".format(onx, e)) + raise OnnxRuntimeAssertionError( + "Unable to run onnx '{0}' due to {1}".format(onx, e) + ) output0 = output.copy() try: - _compare_expected(load["expected"], output, sess, onx, decimal=decimal, **options) + _compare_expected( + load["expected"], output, sess, onx, decimal=decimal, **options + ) except ExpectedAssertionError as expe: raise expe except Exception as e: if verbose: import onnx + model = onnx.load(onx) smodel = "\nJSON ONNX\n" + str(model) else: smodel = "" - raise OnnxRuntimeAssertionError("Model '{0}' has discrepancies.\n{1}: {2}{3}".format(onx, type(e), e, smodel)) + raise OnnxRuntimeAssertionError( + "Model '{0}' has discrepancies.\n{1}: {2}{3}".format( + onx, type(e), e, smodel + ) + ) return output0 @@ -180,12 +226,17 @@ def _post_process_output(res): return numpy.array(res) elif isinstance(res[0], dict): import pandas + return pandas.DataFrame(res).values else: ls = [len(r) for r in res] mi = min(ls) if mi != max(ls): - raise NotImplementedError("Unable to postprocess various number of outputs in [{0}, {1}]".format(min(ls), max(ls))) + raise NotImplementedError( + "Unable to postprocess various number of outputs in [{0}, {1}]".format( + min(ls), max(ls) + ) + ) if mi > 1: output = [] for i in range(mi): @@ -201,7 +252,9 @@ def _post_process_output(res): return res else: if len(res[0]) != 1: - raise NotImplementedError("Not conversion implemented for {0}".format(res)) + raise NotImplementedError( + "Not conversion implemented for {0}".format(res) + ) st = [r[0] for r in res] return numpy.vstack(st) else: @@ -209,6 +262,7 @@ def _post_process_output(res): else: return res + def _create_column(values, dtype): "Creates a column from values with dtype" if str(dtype) == "tensor(int64)": @@ -216,10 +270,14 @@ def _create_column(values, dtype): elif str(dtype) == "tensor(float)": return numpy.array(values, dtype=numpy.float32) else: - raise OnnxRuntimeAssertionError("Unable to create one column from dtype '{0}'".format(dtype)) + raise OnnxRuntimeAssertionError( + "Unable to create one column from dtype '{0}'".format(dtype) + ) -def _compare_expected(expected, output, sess, onnx, decimal=5, onnx_shape=None, **kwargs): +def _compare_expected( + expected, output, sess, onnx, decimal=5, onnx_shape=None, **kwargs +): """ Compares the expected output against the runtime outputs. This is specific to *onnxruntime* due to variable *sess* @@ -228,27 +286,36 @@ def _compare_expected(expected, output, sess, onnx, decimal=5, onnx_shape=None, tested = 0 if isinstance(expected, list): if isinstance(output, list): - if 'Out0' in kwargs: + if "Out0" in kwargs: expected = expected[:1] output = output[:1] - del kwargs['Out0'] - if 'Reshape' in kwargs: - del kwargs['Reshape'] + del kwargs["Out0"] + if "Reshape" in kwargs: + del kwargs["Reshape"] output = numpy.hstack(output).ravel() - output = output.reshape((len(expected), - len(output.ravel()) // len(expected))) + output = output.reshape( + (len(expected), len(output.ravel()) // len(expected)) + ) if len(expected) != len(output): - raise OnnxRuntimeAssertionError("Unexpected number of outputs '{0}', expected={1}, got={2}".format(onnx, len(expected), len(output))) + raise OnnxRuntimeAssertionError( + "Unexpected number of outputs '{0}', expected={1}, got={2}".format( + onnx, len(expected), len(output) + ) + ) try: onnx_shapes = [_.shape for _ in sess.get_outputs()] except TypeError: # Unable to convert function return value to a Python type! onnx_shapes = [None for o in output] for exp, out, osh in zip(expected, output, onnx_shapes): - _compare_expected(exp, out, sess, onnx, decimal=decimal, onnx_shape=osh, **kwargs) + _compare_expected( + exp, out, sess, onnx, decimal=decimal, onnx_shape=osh, **kwargs + ) tested += 1 else: - raise OnnxRuntimeAssertionError("Type mismatch for '{0}', output type is {1}".format(onnx, type(output))) + raise OnnxRuntimeAssertionError( + "Type mismatch for '{0}', output type is {1}".format(onnx, type(output)) + ) elif isinstance(expected, dict): if not isinstance(output, dict): raise OnnxRuntimeAssertionError("Type mismatch for '{0}'".format(onnx)) @@ -257,12 +324,15 @@ def _compare_expected(expected, output, sess, onnx, decimal=5, onnx_shape=None, continue msg = compare_outputs(expected[k], v, decimal=decimal, **kwargs) if msg: - raise OnnxRuntimeAssertionError("Unexpected output '{0}' in model '{1}'\n{2}".format(k, onnx, msg)) + raise OnnxRuntimeAssertionError( + "Unexpected output '{0}' in model '{1}'\n{2}".format(k, onnx, msg) + ) tested += 1 elif isinstance(expected, numpy.ndarray): if isinstance(output, list): if expected.shape[0] == len(output) and isinstance(output[0], dict): import pandas + output = pandas.DataFrame(output) output = output[list(sorted(output.columns))] output = output.values @@ -271,49 +341,75 @@ def _compare_expected(expected, output, sess, onnx, decimal=5, onnx_shape=None, ex = str(output) if len(ex) > 70: ex = ex[:70] + "..." - raise OnnxRuntimeAssertionError("More than one output when 1 is expected for onnx '{0}'\n{1}".format(onnx, ex)) + raise OnnxRuntimeAssertionError( + "More than one output when 1 is expected for onnx '{0}'" + "\n{1}".format(onnx, ex) + ) output = output[-1] if not isinstance(output, numpy.ndarray): - raise OnnxRuntimeAssertionError("output must be an array for onnx '{0}' not {1}".format(onnx, type(output))) + raise OnnxRuntimeAssertionError( + "output must be an array for onnx '{0}' not {1}".format( + onnx, type(output) + ) + ) if onnx_shape is not None: if len(onnx_shape) == 2 and onnx_shape[0] is not None: cols = onnx_shape[1] ecols = output.shape[1] if len(output.shape) == 2 else 1 if cols != ecols: - raise OnnxRuntimeAssertionError("Unexpected onnx shape {0} != {1} for onnx '{2}'".format( - onnx_shape, output.shape, onnx)) + raise OnnxRuntimeAssertionError( + "Unexpected onnx shape {0} != {1} for onnx '{2}'".format( + onnx_shape, output.shape, onnx + ) + ) msg = compare_outputs(expected, output, decimal=decimal, **kwargs) if isinstance(msg, ExpectedAssertionError): raise msg if msg: - raise OnnxRuntimeAssertionError("Unexpected output in model '{0}'\n{1}".format(onnx, msg)) + raise OnnxRuntimeAssertionError( + "Unexpected output in model '{0}'\n{1}".format(onnx, msg) + ) tested += 1 else: from scipy.sparse.csr import csr_matrix + if isinstance(expected, csr_matrix): # DictVectorizer one_array = numpy.array(output) - msg = compare_outputs(expected.todense(), one_array, decimal=decimal, **kwargs) + msg = compare_outputs( + expected.todense(), one_array, decimal=decimal, **kwargs + ) if msg: - raise OnnxRuntimeAssertionError("Unexpected output in model '{0}'\n{1}".format(onnx, msg)) + raise OnnxRuntimeAssertionError( + "Unexpected output in model '{0}'\n{1}".format(onnx, msg) + ) tested += 1 else: - raise OnnxRuntimeAssertionError("Unexpected type for expected output ({1}) and onnx '{0}'".format(onnx, type(expected))) - if tested ==0: + raise OnnxRuntimeAssertionError( + "Unexpected type for expected output ({1}) and onnx '{0}'".format( + onnx, type(expected) + ) + ) + if tested == 0: raise OnnxRuntimeAssertionError("No test for onnx '{0}'".format(onnx)) def run_with_runtime(inputs, model_path): - ''' + """ :param inputs: inputs to the model :param model_path: onnx model file path :return: (output,session) - ''' + """ try: import onnxruntime - session = onnxruntime.InferenceSession(model_path, providers=["CPUExecutionProvider"]) + + session = onnxruntime.InferenceSession( + model_path, providers=["CPUExecutionProvider"] + ) output = session.run(None, inputs) return (output, session) except Exception as e: - raise OnnxRuntimeAssertionError("The runtime does either not exists of fails to load model") + raise OnnxRuntimeAssertionError( + "The runtime does either not exists of fails to load model" + ) from e diff --git a/onnxmltools/utils/visualize.py b/onnxmltools/utils/visualize.py index 765af110..2fcf07d8 100644 --- a/onnxmltools/utils/visualize.py +++ b/onnxmltools/utils/visualize.py @@ -7,7 +7,15 @@ def get_set_node(node, i="0"): - return "g.setNode(" + str(i) + ", { label: '" + node + "', class: 'type-" + node + "' });" + return ( + "g.setNode(" + + str(i) + + ", { label: '" + + node + + "', class: 'type-" + + node + + "' });" + ) def get_set_edge(start, end): @@ -16,10 +24,12 @@ def get_set_edge(start, end): def get_nodes(graph): graph_nodes = [(i, node.op_type) for i, node in enumerate(graph.node, 0)] - graph_nodes.extend([(i, node.name) - for i, node in enumerate(graph.input, len(graph_nodes))]) - graph_nodes.extend([(i, node.name) - for i, node in enumerate(graph.output, len(graph_nodes) + 1)]) + graph_nodes.extend( + [(i, node.name) for i, node in enumerate(graph.input, len(graph_nodes))] + ) + graph_nodes.extend( + [(i, node.name) for i, node in enumerate(graph.output, len(graph_nodes) + 1)] + ) return graph_nodes @@ -44,16 +54,15 @@ def get_edges(graph): for i, node in enumerate(nodes, 0): for input in node.input: if input in output_node_hash.keys(): - edge_list.extend([(node_id, i) - for node_id in output_node_hash[input]]) + edge_list.extend([(node_id, i) for node_id in output_node_hash[input]]) else: - if not input in initializer_names: - print( - "No corresponding output found for {0}.".format(input)) + if input not in initializer_names: + print("No corresponding output found for {0}.".format(input)) for i, output in enumerate(graph.output, len(nodes) + len(graph.input) + 1): if output.name in output_node_hash.keys(): - edge_list.extend([(node_id, i) - for node_id in output_node_hash[output.name]]) + edge_list.extend( + [(node_id, i) for node_id in output_node_hash[output.name]] + ) else: pass return edge_list @@ -76,8 +85,13 @@ def visualize_model(onnx_model, open_browser=True, dest="index.html"): visualize_model(model) """ graph = onnx_model.graph - model_info = "Model produced by: " + onnx_model.producer_name + \ - " version(" + onnx_model.producer_version + ")" + model_info = ( + "Model produced by: " + + onnx_model.producer_name + + " version(" + + onnx_model.producer_version + + ")" + ) html_str = """ @@ -123,11 +137,14 @@ def visualize_model(onnx_model, open_browser=True, dest="index.html"): """ - html_str = html_str.replace("[nodes_html]", "\n".join( - get_nodes_builder(get_nodes(graph)))) + html_str = html_str.replace( + "[nodes_html]", "\n".join(get_nodes_builder(get_nodes(graph))) + ) - html_str = html_str.replace("[edges_html]", "\n".join( - [get_set_edge(edge[0], edge[1]) for edge in get_edges(graph)])) + html_str = html_str.replace( + "[edges_html]", + "\n".join([get_set_edge(edge[0], edge[1]) for edge in get_edges(graph)]), + ) html_str = html_str.replace("[model_info]", model_info) @@ -135,7 +152,7 @@ def visualize_model(onnx_model, open_browser=True, dest="index.html"): Html_file.write(html_str) Html_file.close() - pkgdir = sys.modules['onnxmltools'].__path__[0] + pkgdir = sys.modules["onnxmltools"].__path__[0] fullpath = os.path.join(pkgdir, "utils", "styles.css") shutil.copy(fullpath, os.getcwd()) fullpath = os.path.join(pkgdir, "utils", "dagre-d3.min.js") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..79f48d06 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,24 @@ +[tool.ruff] + +exclude = [ + ".eggs", + ".git", + "build", + "dist", +] + +# Same as Black. +line-length = 95 + +[tool.ruff.mccabe] +max-complexity = 10 + +[tool.ruff.per-file-ignores] +"**/__init__.py" = ["F401"] +"onnxmltools/convert/coreml/operator_converters/GLMClassifier.py" = ["E501"] +"onnxmltools/convert/coreml/operator_converters/SVC.py" = ["E501"] +"onnxmltools/convert/coreml/operator_converters/TensorToLabel.py" = ["E501"] +"onnxmltools/convert/coreml/operator_converters/TreeEnsemble.py" = ["E501"] +"onnxmltools/convert/coreml/operator_converters/neural_network/BidirectionalLSTM.py" = ["E501"] +"onnxmltools/convert/coreml/operator_converters/neural_network/GRU.py" = ["E501"] +"onnxmltools/convert/coreml/operator_converters/neural_network/SimpleRNN.py" = ["E501"] diff --git a/requirements-dev.txt b/requirements-dev.txt index d4b6e564..80d2026b 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,19 +1,18 @@ +black catboost cython dill -flake8 -flatbuffers libsvm lightgbm mleap numpy openpyxl pandas -psutil pyspark pytest pytest-cov pytest-spark +ruff scikit-learn>=1.2.0 scipy wheel diff --git a/setup.py b/setup.py index 126758b2..196b60d0 100644 --- a/setup.py +++ b/setup.py @@ -5,49 +5,55 @@ from distutils.core import setup from setuptools import find_packages import os + this = os.path.dirname(__file__) with open(os.path.join(this, "requirements.txt"), "r") as f: - requirements = [_ for _ in [_.strip("\r\n ") - for _ in f.readlines()] if _ is not None] + requirements = [ + _ for _ in [_.strip("\r\n ") for _ in f.readlines()] if _ is not None + ] packages = find_packages() assert packages # read version from the package file. -version_str = '1.0.0' -with (open(os.path.join(this, 'onnxmltools/__init__.py'), "r")) as f: - line = [_ for _ in [_.strip("\r\n ") - for _ in f.readlines()] if _.startswith("__version__")] +version_str = "1.0.0" +with open(os.path.join(this, "onnxmltools/__init__.py"), "r") as f: + line = [ + _ + for _ in [_.strip("\r\n ") for _ in f.readlines()] + if _.startswith("__version__") + ] if len(line) > 0: - version_str = line[0].split('=')[1].strip('" ') + version_str = line[0].split("=")[1].strip('" ') README = os.path.join(os.getcwd(), "README.md") with open(README) as f: long_description = f.read() setup( - name='onnxmltools', + name="onnxmltools", version=version_str, description="Converts Machine Learning models to ONNX", long_description=long_description, - long_description_content_type='text/markdown', - license='Apache License v2.0', - author='ONNX', - author_email='onnx-technical-discuss@lists.lfaidata.foundation', - url='https://github.com/onnx/onnxmltools', + long_description_content_type="text/markdown", + license="Apache License v2.0", + author="ONNX", + author_email="onnx-technical-discuss@lists.lfaidata.foundation", + url="https://github.com/onnx/onnxmltools", packages=packages, include_package_data=True, install_requires=requirements, classifiers=[ - 'Development Status :: 4 - Beta', - 'Environment :: Console', - 'Intended Audience :: Developers', - 'Operating System :: MacOS :: MacOS X', - 'Operating System :: Microsoft :: Windows', - 'Programming Language :: Python', - 'Programming Language :: Python :: 3.5', - 'Programming Language :: Python :: 3.6', - 'Programming Language :: Python :: 3.7', - 'License :: OSI Approved :: Apache Software License'], + "Development Status :: 4 - Beta", + "Environment :: Console", + "Intended Audience :: Developers", + "Operating System :: MacOS :: MacOS X", + "Operating System :: Microsoft :: Windows", + "Programming Language :: Python", + "Programming Language :: Python :: 3.5", + "Programming Language :: Python :: 3.6", + "Programming Language :: Python :: 3.7", + "License :: OSI Approved :: Apache Software License", + ], ) diff --git a/tests/baseline/test_convert_baseline.py b/tests/baseline/test_convert_baseline.py index ddabacac..3b56bbba 100644 --- a/tests/baseline/test_convert_baseline.py +++ b/tests/baseline/test_convert_baseline.py @@ -7,8 +7,8 @@ import re import unittest from onnx.defs import onnx_opset_version -from onnxmltools.convert import convert_coreml from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER + try: import coremltools except ImportError: @@ -19,36 +19,37 @@ class TestBaseLine(unittest.TestCase): - def check_baseline(self, input_file, ref_file): diff = self.get_diff(input_file, ref_file) return self.normalize_diff(diff) def get_diff(self, input_file, ref_file): + from onnxmltools.convert import convert_coreml + this = os.path.dirname(__file__) coreml_file = os.path.join(this, "models", input_file) - if coremltools is None: - return [] cml = coremltools.utils.load_spec(coreml_file) onnx_model = convert_coreml(cml, target_opset=TARGET_OPSET) 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) - with open(output_file, 'w') as f: + with open(output_file, "w") as f: f.write(str(onnx_model)) 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: + 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 def normalize_diff(self, diff): invalid_comparisons = [] - invalid_comparisons.append(re.compile('producer_version: \"\d+\.\d+\.\d+\.\d+.*')) - invalid_comparisons.append(re.compile('\s+name: \".*')) - invalid_comparisons.append(re.compile('ir_version: \d+')) - invalid_comparisons.append(re.compile('\s+')) + invalid_comparisons.append( + re.compile('producer_version: "\d+\.\d+\.\d+\.\d+.*') + ) + invalid_comparisons.append(re.compile('\s+name: ".*')) + invalid_comparisons.append(re.compile("ir_version: \d+")) + invalid_comparisons.append(re.compile("\s+")) valid_diff = set() for line in diff: if any(comparison.match(line) for comparison in invalid_comparisons): @@ -56,12 +57,18 @@ def normalize_diff(self, diff): valid_diff.add(line) return valid_diff + @unittest.skipIf(coremltools is None, reason="not installed") def test_keras2coreml_Dense_ImageNet_small(self): """ - Converting keras2coreml_Dense_ImageNet_small using onnxmltools and comparing with last known good result + Converting keras2coreml_Dense_ImageNet_small using + onnxmltools and comparing with last known good result """ - self.assertFalse(self.check_baseline( - "keras2coreml_Dense_ImageNet_small.mlmodel", "keras2coreml_Dense_ImageNet_small.json")) + self.assertFalse( + self.check_baseline( + "keras2coreml_Dense_ImageNet_small.mlmodel", + "keras2coreml_Dense_ImageNet_small.json", + ) + ) if __name__ == "__main__": diff --git a/tests/catboost/test_CatBoost_converter.py b/tests/catboost/test_CatBoost_converter.py index 5b729343..daaa28b8 100644 --- a/tests/catboost/test_CatBoost_converter.py +++ b/tests/catboost/test_CatBoost_converter.py @@ -7,6 +7,7 @@ import warnings import packaging.version as pv import numpy + try: import sklearn from sklearn.datasets import make_regression, make_classification @@ -17,56 +18,100 @@ except (ImportError, FileNotFoundError): catboost = None from onnxmltools.convert import convert_catboost -from onnxmltools.utils import dump_data_and_model, dump_single_regression, dump_multiple_classification +from onnxmltools.utils import ( + dump_data_and_model, + dump_single_regression, + dump_multiple_classification, +) class TestCatBoost(unittest.TestCase): - - @unittest.skipIf(catboost is None or sklearn is None, reason="catboost not imported") + @unittest.skipIf( + catboost is None or sklearn is None, reason="catboost not imported" + ) def test_catboost_regressor(self): X, y = make_regression(n_samples=100, n_features=4, random_state=0) - catboost_model = catboost.CatBoostRegressor(task_type='CPU', loss_function='RMSE', - n_estimators=10, verbose=0) + catboost_model = catboost.CatBoostRegressor( + task_type="CPU", loss_function="RMSE", n_estimators=10, verbose=0 + ) dump_single_regression(catboost_model) catboost_model.fit(X.astype(numpy.float32), y) - catboost_onnx = convert_catboost(catboost_model, name='CatBoostRegression', - doc_string='test regression') + catboost_onnx = convert_catboost( + catboost_model, name="CatBoostRegression", doc_string="test regression" + ) self.assertTrue(catboost_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), catboost_model, catboost_onnx, basename="CatBoostReg-Dec4") + dump_data_and_model( + X.astype(numpy.float32), + catboost_model, + catboost_onnx, + basename="CatBoostReg-Dec4", + ) - @unittest.skipIf(catboost is None or sklearn is None, reason="catboost not imported") + @unittest.skipIf( + catboost is None or sklearn is None, reason="catboost not imported" + ) def test_catboost_bin_classifier(self): import onnxruntime - if pv.Version('.'.join(onnxruntime.__version__.split('.')[:2])) >= pv.Version('1.3.0'): + if pv.Version(".".join(onnxruntime.__version__.split(".")[:2])) >= pv.Version( + "1.3.0" + ): X, y = make_classification(n_samples=100, n_features=4, random_state=0) - catboost_model = catboost.CatBoostClassifier(task_type='CPU', loss_function='CrossEntropy', - n_estimators=10, verbose=0) + catboost_model = catboost.CatBoostClassifier( + task_type="CPU", + loss_function="CrossEntropy", + n_estimators=10, + verbose=0, + ) catboost_model.fit(X.astype(numpy.float32), y) - catboost_onnx = convert_catboost(catboost_model, name='CatBoostBinClassification', - doc_string='test binary classification') + catboost_onnx = convert_catboost( + catboost_model, + name="CatBoostBinClassification", + doc_string="test binary classification", + ) self.assertTrue(catboost_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), catboost_model, catboost_onnx, basename="CatBoostBinClass") + dump_data_and_model( + X.astype(numpy.float32), + catboost_model, + catboost_onnx, + basename="CatBoostBinClass", + ) else: - warnings.warn('Converted CatBoost models for binary classification work with onnxruntime version 1.3.0 or ' - 'a newer one') + warnings.warn( + "Converted CatBoost models for binary classification " + "work with onnxruntime version 1.3.0 or " + "a newer one" + ) - @unittest.skipIf(catboost is None or sklearn is None, reason="catboost not imported") + @unittest.skipIf( + catboost is None or sklearn is None, reason="catboost not imported" + ) def test_catboost_multi_classifier(self): - X, y = make_classification(n_samples=10, n_informative=8, n_classes=3, random_state=0) - catboost_model = catboost.CatBoostClassifier(task_type='CPU', loss_function='MultiClass', - n_estimators=100, verbose=0) + X, y = make_classification( + n_samples=10, n_informative=8, n_classes=3, random_state=0 + ) + catboost_model = catboost.CatBoostClassifier( + task_type="CPU", loss_function="MultiClass", n_estimators=100, verbose=0 + ) dump_multiple_classification(catboost_model) catboost_model.fit(X.astype(numpy.float32), y) - catboost_onnx = convert_catboost(catboost_model, name='CatBoostMultiClassification', - doc_string='test multiclass classification') + catboost_onnx = convert_catboost( + catboost_model, + name="CatBoostMultiClassification", + doc_string="test multiclass classification", + ) self.assertTrue(catboost_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), catboost_model, catboost_onnx, basename="CatBoostMultiClass") + dump_data_and_model( + X.astype(numpy.float32), + catboost_model, + catboost_onnx, + basename="CatBoostMultiClass", + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_AllNeuralNetworkConverters.py b/tests/coreml/test_cml_AllNeuralNetworkConverters.py index 425bfd07..c209b3f1 100644 --- a/tests/coreml/test_cml_AllNeuralNetworkConverters.py +++ b/tests/coreml/test_cml_AllNeuralNetworkConverters.py @@ -3,12 +3,14 @@ import unittest import packaging.version as pv import numpy + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import onnx @@ -24,97 +26,156 @@ class TestNeuralNetworkLayerConverter(unittest.TestCase): - def test_inner_product_converter(self): input_dim = (3,) output_dim = (2,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] weights = numpy.zeros(shape=(3, 2)) weights[:] = [[1, 2], [3, 4], [5, 6]] bias = numpy.zeros(shape=(2)) bias[:] = [-100, 100] builder = NeuralNetworkBuilder(input, output) - builder.add_inner_product(name='FC', W=weights, b=bias, input_channels=3, output_channels=2, has_bias=True, - input_name='input', output_name='output') + builder.add_inner_product( + name="FC", + W=weights, + b=bias, + input_channels=3, + output_channels=2, + has_bias=True, + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_unary_function_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_unary(name='Unary1', input_name='input', output_name='mid1', mode='abs') - builder.add_unary(name='Unary2', input_name='mid1', output_name='mid2', mode='sqrt') - builder.add_unary(name='Unary3', input_name='mid2', output_name='mid3', mode='rsqrt') - builder.add_unary(name='Unary4', input_name='mid3', output_name='mid4', mode='inverse') - builder.add_unary(name='Unary5', input_name='mid4', output_name='mid5', mode='power', alpha=2) - builder.add_unary(name='Unary6', input_name='mid5', output_name='mid6', mode='exp') - builder.add_unary(name='Unary7', input_name='mid6', output_name='mid7', mode='log') - builder.add_unary(name='Unary8', input_name='mid7', output_name='output', mode='threshold') + builder.add_unary( + name="Unary1", input_name="input", output_name="mid1", mode="abs" + ) + builder.add_unary( + name="Unary2", input_name="mid1", output_name="mid2", mode="sqrt" + ) + builder.add_unary( + name="Unary3", input_name="mid2", output_name="mid3", mode="rsqrt" + ) + builder.add_unary( + name="Unary4", input_name="mid3", output_name="mid4", mode="inverse" + ) + builder.add_unary( + name="Unary5", input_name="mid4", output_name="mid5", mode="power", alpha=2 + ) + builder.add_unary( + name="Unary6", input_name="mid5", output_name="mid6", mode="exp" + ) + builder.add_unary( + name="Unary7", input_name="mid6", output_name="mid7", mode="log" + ) + builder.add_unary( + name="Unary8", input_name="mid7", output_name="output", mode="threshold" + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_convolution_converter(self): input_dim = (1, 1, 4, 2) output_dim = (1, 1, 4, 2) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) weights = numpy.zeros(shape=(1, 1, 2, 2)) weights[:] = [[1, 1], [-1, -1]] bias = numpy.zeros(shape=(1,)) bias[:] = 100 - builder.add_convolution(name='Conv', kernel_channels=1, output_channels=1, height=2, width=2, stride_height=1, - stride_width=1, border_mode='same', groups=1, - W=weights, b=bias, has_bias=True, input_name='input', output_name='output', - is_deconv=True, output_shape=(1, 1, 4, 2)) + builder.add_convolution( + name="Conv", + kernel_channels=1, + output_channels=1, + height=2, + width=2, + stride_height=1, + stride_width=1, + border_mode="same", + groups=1, + W=weights, + b=bias, + has_bias=True, + input_name="input", + output_name="output", + is_deconv=True, + output_shape=(1, 1, 4, 2), + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_pooling_converter(self): input_dim = (1, 1, 4, 2) output_dim = (1, 1, 4, 2) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_pooling(name='Pool', height=2, width=2, stride_height=1, stride_width=1, layer_type='MAX', - padding_type='SAME', input_name='input', output_name='output') + builder.add_pooling( + name="Pool", + height=2, + width=2, + stride_height=1, + stride_width=1, + layer_type="MAX", + padding_type="SAME", + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_activation_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_activation(name='Activation', non_linearity='RELU', input_name='input', output_name='output') + builder.add_activation( + name="Activation", + non_linearity="RELU", + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_embedding_converter(self): input_dim = (1, 1, 1, 1) output_dim = (1, 2, 1, 1) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) weights = numpy.zeros(shape=(2)) weights[:] = [-1, 1] bias = numpy.zeros(shape=(2)) bias[:] = [-100, 100] - builder.add_embedding(name='Embed', input_dim=1, W=weights, b=bias, output_channels=2, has_bias=True, - input_name='input', output_name='output') + builder.add_embedding( + name="Embed", + input_dim=1, + W=weights, + b=bias, + output_channels=2, + has_bias=True, + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_batchnorm_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) gamma = numpy.ndarray(shape=(3,)) gamma[:] = [-1, 0, 1] @@ -124,48 +185,64 @@ def test_batchnorm_converter(self): mean[:] = [0, 0, 0] variance = numpy.ndarray(shape=(3,)) variance[:] = [1, 1, 1] - builder.add_batchnorm(name='BatchNormalize', channels=3, gamma=gamma, beta=beta, mean=mean, variance=variance, - input_name='input', output_name='output') + builder.add_batchnorm( + name="BatchNormalize", + channels=3, + gamma=gamma, + beta=beta, + mean=mean, + variance=variance, + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_mean_variance_normalize_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_mvn(name='MVN', input_name='input', output_name='output', epsilon=0) + builder.add_mvn(name="MVN", input_name="input", output_name="output", epsilon=0) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_l2_normalize_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_l2_normalize(name='L2', input_name='input', output_name='output') + builder.add_l2_normalize(name="L2", input_name="input", output_name="output") model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_softmax_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_softmax(name='Softmax', input_name='input', output_name='output') + builder.add_softmax(name="Softmax", input_name="input", output_name="output") model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_lrn_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_lrn(name='LRN', input_name='input', output_name='output', alpha=0.5, beta=2, k=1, local_size=2) + builder.add_lrn( + name="LRN", + input_name="input", + output_name="output", + alpha=0.5, + beta=2, + k=1, + local_size=2, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) @@ -176,129 +253,217 @@ def test_crop_converter(self): right_crop = 1 input_dim = (8, 6, 4) output_dim = (8, 6 - top_crop - bottom_crop, 4 - left_crop - right_crop) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_crop(name='Crop', left=left_crop, right=right_crop, top=top_crop, bottom=bottom_crop, offset=[0, 0], - input_names=['input'], output_name='output') + builder.add_crop( + name="Crop", + left=left_crop, + right=right_crop, + top=top_crop, + bottom=bottom_crop, + offset=[0, 0], + input_names=["input"], + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_padding_converter(self): input_dim = (1, 3, 4) output_dim = (1, 5, 6) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_padding(name='Pad', left=2, right=0, top=2, bottom=0, input_name='input', output_name='output', - padding_type='constant') + builder.add_padding( + name="Pad", + left=2, + right=0, + top=2, + bottom=0, + input_name="input", + output_name="output", + padding_type="constant", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_upsample_converter(self): input_dim = (1, 1, 1) output_dim = (1, 2, 2) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) - builder.add_upsample(name='Upsample', scaling_factor_h=2, scaling_factor_w=2, input_name='input', - output_name='output') + builder.add_upsample( + name="Upsample", + scaling_factor_h=2, + scaling_factor_w=2, + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_add_converter(self): input_dim = (1, 2, 2) output_dim = (1, 2, 2) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_elementwise(name='Add', input_names=['input1', 'input2'], output_name='output', mode='ADD') + builder.add_elementwise( + name="Add", + input_names=["input1", "input2"], + output_name="output", + mode="ADD", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_multiply_converter(self): input_dim = (1, 2, 2) output_dim = (1, 2, 2) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_elementwise(name='Mul', input_names=['input1', 'input2'], output_name='output', mode='MULTIPLY') + builder.add_elementwise( + name="Mul", + input_names=["input1", "input2"], + output_name="output", + mode="MULTIPLY", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_average_converter(self): input_dim = (1, 2, 2) output_dim = (1, 2, 2) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_elementwise(name='MEAN', input_names=['input1', 'input2'], output_name='output', mode='AVE') + builder.add_elementwise( + name="MEAN", + input_names=["input1", "input2"], + output_name="output", + mode="AVE", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_scale_converter(self): input_dim = (3,) output_dim = (3,) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) scale = numpy.ndarray(shape=(1,)) scale[:] = 10 bias = numpy.ndarray(shape=(1,)) bias[:] = -100 - builder.add_scale(name='ImageScaler', W=scale, b=bias, has_bias=True, input_name='input', output_name='output') + builder.add_scale( + name="ImageScaler", + W=scale, + b=bias, + has_bias=True, + input_name="input", + output_name="output", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_bias_converter(self): input_dim = (2, 1, 1) output_dim = (2, 1, 1) - input = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + input = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(input, output) bias = numpy.ndarray(shape=(2,)) bias[:] = [1, 2] - builder.add_bias(name='Bias', b=bias, input_name='input', output_name='output', shape_bias=[2]) + builder.add_bias( + name="Bias", + b=bias, + input_name="input", + output_name="output", + shape_bias=[2], + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_max_converter(self): input_dim = (1, 2, 2) output_dim = (1, 2, 2) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_elementwise(name='Max', input_names=['input1', 'input2'], output_name='output', mode='MAX') + builder.add_elementwise( + name="Max", + input_names=["input1", "input2"], + output_name="output", + mode="MAX", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_min_converter(self): input_dim = (1, 2, 2) output_dim = (1, 2, 2) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_elementwise(name='Min', input_names=['input1', 'input2'], output_name='output', mode='MIN') + builder.add_elementwise( + name="Min", + input_names=["input1", "input2"], + output_name="output", + mode="MIN", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_dot_product_converter(self): input_dim = (3,) output_dim = (1,) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_elementwise(name='Dot', input_names=['input1', 'input2'], output_name='output', mode='DOT') + builder.add_elementwise( + name="Dot", + input_names=["input1", "input2"], + output_name="output", + mode="DOT", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_reduce_converter(self): input_dim = (1, 2, 2) output_dim = (1,) - inputs = [('input', datatypes.Array(*input_dim))] - output = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + output = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, output) - builder.add_reduce(name='Reduce', input_name='input', output_name='output', axis='CHW', mode='sum') + builder.add_reduce( + name="Reduce", + input_name="input", + output_name="output", + axis="CHW", + mode="sum", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) @@ -307,71 +472,103 @@ def test_load_constant_converter(self): value[:] = [[[-95, 95]]] shape = value.shape input_dim = (1, 2, 3, 4) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('const', datatypes.Array(*shape)), ('output', datatypes.Array(*input_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [ + ("const", datatypes.Array(*shape)), + ("output", datatypes.Array(*input_dim)), + ] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_load_constant(name='LoadConstant', output_name='const', constant_value=value, shape=shape) - builder.add_permute(name='Permute', input_name='input', output_name='output', dim=(0, 1, 2, 3)) + builder.add_load_constant( + name="LoadConstant", output_name="const", constant_value=value, shape=shape + ) + builder.add_permute( + name="Permute", input_name="input", output_name="output", dim=(0, 1, 2, 3) + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_reshape_converter(self): input_dim = (1, 1, 2) output_dim = (1, 2, 1) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_reshape(name='Reshape', input_name='input', output_name='output', target_shape=output_dim, mode=1) + builder.add_reshape( + name="Reshape", + input_name="input", + output_name="output", + target_shape=output_dim, + mode=1, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_flatten_converter(self): input_dim = (1, 2, 3) output_dim = (6, 1, 1) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_flatten(name='Flatten', input_name='input', output_name='output', mode=1) + builder.add_flatten( + name="Flatten", input_name="input", output_name="output", mode=1 + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_permute_converter(self): input_dim = (4, 1, 2, 3) output_dim = (4, 3, 1, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_permute(name='Permute', input_name='input', output_name='output', dim=(0, 2, 3, 1)) + builder.add_permute( + name="Permute", input_name="input", output_name="output", dim=(0, 2, 3, 1) + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_concat_converter(self): input_dim = (5, 1, 1) output_dim = (10, 1, 1) - inputs = [('input1', datatypes.Array(*input_dim)), ('input2', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [ + ("input1", datatypes.Array(*input_dim)), + ("input2", datatypes.Array(*input_dim)), + ] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_elementwise(name='Concate', input_names=['input1', 'input2'], output_name='output', mode='CONCAT') + builder.add_elementwise( + name="Concate", + input_names=["input1", "input2"], + output_name="output", + mode="CONCAT", + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_split_converter(self): input_dim = (8, 1, 1) output_dim = (4, 1, 1) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output1', datatypes.Array(*output_dim)), ('output2', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [ + ("output1", datatypes.Array(*output_dim)), + ("output2", datatypes.Array(*output_dim)), + ] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_split(name='Split', input_name='input', output_names=['output1', 'output2']) + builder.add_split( + name="Split", input_name="input", output_names=["output1", "output2"] + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_sequence_repeat_converter(self): input_dim = (3, 1, 1) output_dim = (9, 1, 1) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_sequence_repeat(name='Repeat', input_name='input', output_name='output', nrep=3) + builder.add_sequence_repeat( + name="Repeat", input_name="input", output_name="output", nrep=3 + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) @@ -379,110 +576,243 @@ def test_reorganize_data_converter(self): block_size = 2 input_dim = (3, 4 * block_size, 2 * block_size) output_dim = (3 * block_size * block_size, 4, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_reorganize_data(name='Reorg', input_name='input', output_name='output', mode='SPACE_TO_DEPTH', - block_size=2) + builder.add_reorganize_data( + name="Reorg", + input_name="input", + output_name="output", + mode="SPACE_TO_DEPTH", + block_size=2, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_slice_converter(self): input_dim = (1, 4, 2) output_dim = (1, 2, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_slice(name='Slice', input_name='input', output_name='output', axis='height', start_index=0, - end_index=-1, stride=1) + builder.add_slice( + name="Slice", + input_name="input", + output_name="output", + axis="height", + start_index=0, + end_index=-1, + stride=1, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_gru_converter(self): input_dim = (1, 8) output_dim = (1, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - W_h = [numpy.random.rand(2, 2), numpy.random.rand(2, 2), numpy.random.rand(2, 2)] - W_x = [numpy.random.rand(2, 8), numpy.random.rand(2, 8), numpy.random.rand(2, 8)] + W_h = [ + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + ] + W_x = [ + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + ] b = [numpy.random.rand(2, 1), numpy.random.rand(2, 1), numpy.random.rand(2, 1)] - builder.add_gru(name='GRU', W_h=W_h, W_x=W_x, b=b, hidden_size=2, input_size=8, input_names=['input'], - output_names=['output'], activation='TANH', inner_activation='SIGMOID_HARD', output_all=False, - reverse_input=False) + builder.add_gru( + name="GRU", + W_h=W_h, + W_x=W_x, + b=b, + hidden_size=2, + input_size=8, + input_names=["input"], + output_names=["output"], + activation="TANH", + inner_activation="SIGMOID_HARD", + output_all=False, + reverse_input=False, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_simple_recurrent_converter(self): input_dim = (1, 8) output_dim = (1, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) W_h = numpy.random.rand(2, 2) W_x = numpy.random.rand(2, 8) b = numpy.random.rand(2, 1) - builder.add_simple_rnn(name='RNN', W_h=W_h, W_x=W_x, b=b, hidden_size=2, input_size=8, - input_names=['input', 'h_init'], output_names=['output', 'h'], activation='TANH', - output_all=False, reverse_input=False) + builder.add_simple_rnn( + name="RNN", + W_h=W_h, + W_x=W_x, + b=b, + hidden_size=2, + input_size=8, + input_names=["input", "h_init"], + output_names=["output", "h"], + activation="TANH", + output_all=False, + reverse_input=False, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_unidirectional_lstm_converter(self): input_dim = (1, 8) output_dim = (1, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - W_h = [numpy.random.rand(2, 2), numpy.random.rand(2, 2), numpy.random.rand(2, 2), numpy.random.rand(2, 2)] - W_x = [numpy.random.rand(2, 8), numpy.random.rand(2, 8), numpy.random.rand(2, 8), numpy.random.rand(2, 8)] - b = [numpy.random.rand(2, 1), numpy.random.rand(2, 1), numpy.random.rand(2, 1), numpy.random.rand(2, 1)] - p = [numpy.zeros(shape=(2, 1)), numpy.zeros(shape=(2, 1)), numpy.zeros(shape=(2, 1))] - builder.add_unilstm(name='LSTM', W_h=W_h, W_x=W_x, b=b, hidden_size=2, input_size=8, input_names=['input'], - output_names=['output'], inner_activation='SIGMOID', cell_state_update_activation='TANH', - output_activation='TANH', peep=p, output_all=False, forget_bias=False, - coupled_input_forget_gate=False, cell_clip_threshold=10000, reverse_input=False) + W_h = [ + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + ] + W_x = [ + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + ] + b = [ + numpy.random.rand(2, 1), + numpy.random.rand(2, 1), + numpy.random.rand(2, 1), + numpy.random.rand(2, 1), + ] + p = [ + numpy.zeros(shape=(2, 1)), + numpy.zeros(shape=(2, 1)), + numpy.zeros(shape=(2, 1)), + ] + builder.add_unilstm( + name="LSTM", + W_h=W_h, + W_x=W_x, + b=b, + hidden_size=2, + input_size=8, + input_names=["input"], + output_names=["output"], + inner_activation="SIGMOID", + cell_state_update_activation="TANH", + output_activation="TANH", + peep=p, + output_all=False, + forget_bias=False, + coupled_input_forget_gate=False, + cell_clip_threshold=10000, + reverse_input=False, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_bidirectional_lstm_converter(self): input_dim = (1, 8) output_dim = (1, 2) - inputs = [('input', datatypes.Array(*input_dim))] - outputs = [('output', datatypes.Array(*output_dim))] + inputs = [("input", datatypes.Array(*input_dim))] + outputs = [("output", datatypes.Array(*output_dim))] builder = NeuralNetworkBuilder(inputs, outputs) - W_h = [numpy.random.rand(2, 2), numpy.random.rand(2, 2), numpy.random.rand(2, 2), numpy.random.rand(2, 2)] - W_x = [numpy.random.rand(2, 8), numpy.random.rand(2, 8), numpy.random.rand(2, 8), numpy.random.rand(2, 8)] - b = [numpy.random.rand(2, 1), numpy.random.rand(2, 1), numpy.random.rand(2, 1), numpy.random.rand(2, 1)] - p = [numpy.zeros(shape=(2, 1)), numpy.zeros(shape=(2, 1)), numpy.zeros(shape=(2, 1))] - builder.add_bidirlstm(name='LSTM', W_h=W_h, W_x=W_x, W_h_back=W_h, b=b, W_x_back=W_x, b_back=b, hidden_size=2, - input_size=8, input_names=['input'], output_names=['output'], inner_activation='SIGMOID', - cell_state_update_activation='TANH', output_activation='TANH', peep=p, peep_back=p, - output_all=False, forget_bias=False, coupled_input_forget_gate=False, - cell_clip_threshold=10000) + W_h = [ + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + numpy.random.rand(2, 2), + ] + W_x = [ + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + numpy.random.rand(2, 8), + ] + b = [ + numpy.random.rand(2, 1), + numpy.random.rand(2, 1), + numpy.random.rand(2, 1), + numpy.random.rand(2, 1), + ] + p = [ + numpy.zeros(shape=(2, 1)), + numpy.zeros(shape=(2, 1)), + numpy.zeros(shape=(2, 1)), + ] + builder.add_bidirlstm( + name="LSTM", + W_h=W_h, + W_x=W_x, + W_h_back=W_h, + b=b, + W_x_back=W_x, + b_back=b, + hidden_size=2, + input_size=8, + input_names=["input"], + output_names=["output"], + inner_activation="SIGMOID", + cell_state_update_activation="TANH", + output_activation="TANH", + peep=p, + peep_back=p, + output_all=False, + forget_bias=False, + coupled_input_forget_gate=False, + cell_clip_threshold=10000, + ) model_onnx = convert_coreml(builder.spec, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) def test_image_input_type_converter(self): dim = (3, 15, 25) - inputs = [('input', datatypes.Array(*dim))] - outputs = [('output', datatypes.Array(*dim))] + inputs = [("input", datatypes.Array(*dim))] + outputs = [("output", datatypes.Array(*dim))] builder = NeuralNetworkBuilder(inputs, outputs) - builder.add_elementwise(name='Identity', input_names=['input'], - output_name='output', mode='ADD', alpha=0.0) + builder.add_elementwise( + name="Identity", + input_names=["input"], + output_name="output", + mode="ADD", + alpha=0.0, + ) spec = builder.spec input = spec.description.input[0] input.type.imageType.height = dim[1] input.type.imageType.width = dim[2] - for coreml_colorspace, onnx_colorspace in (('RGB', 'Rgb8'), ('BGR', 'Bgr8'), ('GRAYSCALE', 'Gray8')): - input.type.imageType.colorSpace = ImageFeatureType.ColorSpace.Value(coreml_colorspace) + for coreml_colorspace, onnx_colorspace in ( + ("RGB", "Rgb8"), + ("BGR", "Bgr8"), + ("GRAYSCALE", "Gray8"), + ): + input.type.imageType.colorSpace = ImageFeatureType.ColorSpace.Value( + coreml_colorspace + ) model_onnx = convert_coreml(spec, target_opset=TARGET_OPSET) - dims = [(d.dim_param or d.dim_value) for d in model_onnx.graph.input[0].type.tensor_type.shape.dim] - self.assertEqual(dims, ['None', 1 if onnx_colorspace == 'Gray8' else 3, 15, 25]) - - if pv.Version(onnx.__version__) >= pv.Version('1.2.1'): + dims = [ + (d.dim_param or d.dim_value) + for d in model_onnx.graph.input[0].type.tensor_type.shape.dim + ] + self.assertEqual( + dims, ["None", 1 if onnx_colorspace == "Gray8" else 3, 15, 25] + ) + + if pv.Version(onnx.__version__) >= pv.Version("1.2.1"): metadata = {prop.key: prop.value for prop in model_onnx.metadata_props} - self.assertEqual(metadata, { 'Image.BitmapPixelFormat': onnx_colorspace }) - self.assertEqual(model_onnx.graph.input[0].type.denotation, 'IMAGE') - channel_denotations = [d.denotation for d in model_onnx.graph.input[0].type.tensor_type.shape.dim] - self.assertEqual(channel_denotations, ['DATA_BATCH', 'DATA_CHANNEL', 'DATA_FEATURE', 'DATA_FEATURE']) + self.assertEqual(metadata, {"Image.BitmapPixelFormat": onnx_colorspace}) + self.assertEqual(model_onnx.graph.input[0].type.denotation, "IMAGE") + channel_denotations = [ + d.denotation + for d in model_onnx.graph.input[0].type.tensor_type.shape.dim + ] + self.assertEqual( + channel_denotations, + ["DATA_BATCH", "DATA_CHANNEL", "DATA_FEATURE", "DATA_FEATURE"], + ) diff --git a/tests/coreml/test_cml_DictVectorizerConverter.py b/tests/coreml/test_cml_DictVectorizerConverter.py index 8818f1ae..2cbd8386 100644 --- a/tests/coreml/test_cml_DictVectorizerConverter.py +++ b/tests/coreml/test_cml_DictVectorizerConverter.py @@ -8,12 +8,14 @@ import unittest import onnx import sklearn + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -28,13 +30,12 @@ class TestCoreMLDictVectorizerConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_dict_vectorizer(self): model = DictVectorizer() - data = [{'amy': 1., 'chin': 200.}, {'nice': 3., 'amy': 1.}] + data = [{"amy": 1.0, "chin": 200.0}, {"nice": 3.0, "amy": 1.0}] model.fit_transform(data) try: model_coreml = coremltools.converters.sklearn.convert(model) @@ -42,12 +43,19 @@ def test_dict_vectorizer(self): raise AssertionError( "Unable to use coremltools, coremltools.__version__=%r, " "onnx.__version__=%r, sklearn.__version__=%r, " - "sys.platform=%r." % ( - coremltools.__version__, onnx.__version__, - sklearn.__version__, sys.platform)) from e + "sys.platform=%r." + % ( + coremltools.__version__, + onnx.__version__, + sklearn.__version__, + sys.platform, + ) + ) from e model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - dump_data_and_model(data, model, model_onnx, basename="CmlDictVectorizer-OneOff-SkipDim1") + dump_data_and_model( + data, model, model_onnx, basename="CmlDictVectorizer-OneOff-SkipDim1" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_GLMClassifierConverter.py b/tests/coreml/test_cml_GLMClassifierConverter.py index 66d654f9..ce5469c3 100644 --- a/tests/coreml/test_cml_GLMClassifierConverter.py +++ b/tests/coreml/test_cml_GLMClassifierConverter.py @@ -9,12 +9,14 @@ from sklearn.datasets import load_iris from sklearn.linear_model import LogisticRegression from sklearn.svm import LinearSVC + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -28,18 +30,17 @@ class TestCoreMLGLMClassifierConverter(unittest.TestCase): - def validate_zipmap(self, model): # Validate that it contains a ZipMap nodes = model.graph.node - node = next((n for n in nodes if n.op_type == 'ZipMap'), None) + node = next((n for n in nodes if n.op_type == "ZipMap"), None) self.assertIsNotNone(node) self.assertEqual(len(node.output), 1) - self.assertTrue('classProbability' in node.output) + self.assertTrue("classProbability" in node.output) @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_glm_classifier(self): iris = load_iris() X = iris.data[:, :2] @@ -47,13 +48,15 @@ def test_glm_classifier(self): y[y == 2] = 1 # scikit-learn has changed the default value for multi_class. - lr = LogisticRegression(multi_class='ovr') + lr = LogisticRegression(multi_class="ovr") lr.fit(X, y) lr_coreml = coremltools.converters.sklearn.convert(lr) lr_onnx = convert(lr_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(lr_onnx is not None) self.validate_zipmap(lr_onnx) - dump_data_and_model(X.astype(numpy.float32), lr, lr_onnx, basename="CmlbinLogitisticRegression") + dump_data_and_model( + X.astype(numpy.float32), lr, lr_onnx, basename="CmlbinLogitisticRegression" + ) # Ensure there is a probability output svm = LinearSVC() @@ -62,7 +65,9 @@ def test_glm_classifier(self): svm_onnx = convert(svm_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(svm_onnx is not None) self.validate_zipmap(svm_onnx) - dump_data_and_model(X.astype(numpy.float32), svm, svm_onnx, basename="CmlBinLinearSVC-NoProb") + dump_data_and_model( + X.astype(numpy.float32), svm, svm_onnx, basename="CmlBinLinearSVC-NoProb" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_GLMRegressorConverter.py b/tests/coreml/test_cml_GLMRegressorConverter.py index 9fb9ed82..1af23ea1 100644 --- a/tests/coreml/test_cml_GLMRegressorConverter.py +++ b/tests/coreml/test_cml_GLMRegressorConverter.py @@ -6,12 +6,14 @@ import unittest import packaging.version as pv import numpy + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -28,10 +30,9 @@ class TestCoreMLGLMRegressorConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_glm_regressor(self): X, y = make_regression(n_features=4, random_state=0) @@ -40,14 +41,18 @@ def test_glm_regressor(self): lr_coreml = coremltools.converters.sklearn.convert(lr) lr_onnx = convert(lr_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(lr_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), lr, lr_onnx, basename="CmlLinearRegression-Dec4") + dump_data_and_model( + X.astype(numpy.float32), lr, lr_onnx, basename="CmlLinearRegression-Dec4" + ) svr = LinearSVR() svr.fit(X, y) svr_coreml = coremltools.converters.sklearn.convert(svr) svr_onnx = convert(svr_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(svr_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), svr, svr_onnx, basename="CmlLinearSvr-Dec4") + dump_data_and_model( + X.astype(numpy.float32), svr, svr_onnx, basename="CmlLinearSvr-Dec4" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_ImputerConverter.py b/tests/coreml/test_cml_ImputerConverter.py index 893c52ce..0dfdd42a 100644 --- a/tests/coreml/test_cml_ImputerConverter.py +++ b/tests/coreml/test_cml_ImputerConverter.py @@ -6,12 +6,14 @@ import unittest import packaging.version as pv import numpy as np + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import sklearn.preprocessing @@ -25,30 +27,34 @@ class TestCoreMLImputerConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_imputer(self): try: - model = Imputer(missing_values='NaN', strategy='mean', axis=0) + model = Imputer(missing_values="NaN", strategy="mean", axis=0) except TypeError: - model = Imputer(missing_values=np.nan, strategy='mean') + model = Imputer(missing_values=np.nan, strategy="mean") model.axis = 0 data = [[1, 2], [np.nan, 3], [7, 6]] model.fit(data) from onnxmltools.convert.coreml.convert import convert import coremltools # noqa + try: model_coreml = coremltools.converters.sklearn.convert(model) except ValueError as e: - if 'not supported' in str(e): + if "not supported" in str(e): # Python 2.7 + scikit-learn 0.22 return model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - dump_data_and_model(np.array(data, dtype=np.float32), - model, model_onnx, basename="CmlImputerMeanFloat32") + dump_data_and_model( + np.array(data, dtype=np.float32), + model, + model_onnx, + basename="CmlImputerMeanFloat32", + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_OneHotEncoderConverter.py b/tests/coreml/test_cml_OneHotEncoderConverter.py index f8e243cb..c4bed52f 100644 --- a/tests/coreml/test_cml_OneHotEncoderConverter.py +++ b/tests/coreml/test_cml_OneHotEncoderConverter.py @@ -3,19 +3,19 @@ """ Main functions to convert machine learned model from *Core ML* model to *ONNX*. """ -import sys import os import unittest import warnings import packaging.version as pv import numpy -import onnx + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -30,10 +30,9 @@ class TestCoremlOneHotEncoderConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_one_hot_encoder(self): script_dir = os.path.dirname(__file__) relative_path = "../data/onehot_simple.mlmodel" @@ -42,27 +41,31 @@ def test_one_hot_encoder(self): model_onnx = convert(model_coreml, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - @unittest.skip('broken with the dump_data_and_model change.') + @unittest.skip("broken with the dump_data_and_model change.") def test_conversion_one_column(self): scikit_data = [[0], [1], [2], [4], [3], [2], [4], [5], [6], [7]] - scikit_data = numpy.asarray(scikit_data, dtype='d') - scikit_data_multiple_cols = [[0, 1], [1, 0], [2, 2], [3, 3], [4, 4]] - scikit_data_multiple_cols = numpy.asarray(scikit_data_multiple_cols, dtype='d') + scikit_data = numpy.asarray(scikit_data, dtype="d") + scikit_data_multiple_cols = [[0, 1], [1, 0], [2, 2], [3, 3], [4, 4]] + scikit_data_multiple_cols = numpy.asarray(scikit_data_multiple_cols, dtype="d") scikit_model = OneHotEncoder() # scikit_model.fit(scikit_data) - # model_coreml = coremltools.converters.sklearn.convert(scikit_model, 'single_feature', 'out') + # model_coreml = coremltools.converters.sklearn.convert( + # scikit_model, 'single_feature', 'out') scikit_model.fit(scikit_data_multiple_cols) try: - model_coreml = coremltools.converters.sklearn.convert(scikit_model, ['feature_1', 'feature_2'], 'out') - except Exception as e: - warnings.warn("Unable to run convert OneHotEncoder with coreml.") + model_coreml = coremltools.converters.sklearn.convert( + scikit_model, ["feature_1", "feature_2"], "out" + ) + except Exception: + warnings.warn("Unable to run convert OneHotEncoder with coreml due to {e}.") return model_onnx = convert(model_coreml, target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - dump_data_and_model(scikit_data, scikit_model, model_onnx, basename="CmlOneHotEncoder-SkipDim1") - + dump_data_and_model( + scikit_data, scikit_model, model_onnx, basename="CmlOneHotEncoder-SkipDim1" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_ScalerConverter.py b/tests/coreml/test_cml_ScalerConverter.py index 11e846bc..8f515f23 100644 --- a/tests/coreml/test_cml_ScalerConverter.py +++ b/tests/coreml/test_cml_ScalerConverter.py @@ -6,12 +6,14 @@ import unittest import packaging.version as pv import numpy + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -26,18 +28,21 @@ class TestCoreMLScalerConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_scaler(self): model = StandardScaler() - data = numpy.array([[0, 0, 3], [1, 1, 0], [0, 2, 1], [1, 0, 2]], dtype=numpy.float32) + data = numpy.array( + [[0, 0, 3], [1, 1, 0], [0, 2, 1], [1, 0, 2]], dtype=numpy.float32 + ) model.fit(data) model_coreml = coremltools.converters.sklearn.convert(model) model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - dump_data_and_model(data, model, model_onnx, basename="CmlStandardScalerFloat32") + dump_data_and_model( + data, model, model_onnx, basename="CmlStandardScalerFloat32" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_SupportVectorClassifierConverter.py b/tests/coreml/test_cml_SupportVectorClassifierConverter.py index a3f884b2..b72cff5a 100644 --- a/tests/coreml/test_cml_SupportVectorClassifierConverter.py +++ b/tests/coreml/test_cml_SupportVectorClassifierConverter.py @@ -4,12 +4,14 @@ Tests CoreML SupportVectorClassifier converter. """ import packaging.version as pv + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -27,7 +29,6 @@ class TestCoreMLSupportVectorClassifierConverter(unittest.TestCase): - def _fit_binary_classification(self, model): iris = load_iris() @@ -56,18 +57,18 @@ def _check_model_outputs(self, model, output_names): self.assertTrue(name in output_map) def validate_zipmap(self, model): - ''' + """ Validate that it contains a ZipMap - ''' + """ nodes = model.graph.node - node = next((n for n in nodes if n.op_type == 'ZipMap'), None) + node = next((n for n in nodes if n.op_type == "ZipMap"), None) self.assertIsNotNone(node) self.assertEqual(len(node.output), 1) - self.assertTrue('classProbability' in node.output) + self.assertTrue("classProbability" in node.output) @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_support_vector_classifier_binary_no_prob(self): svm, X = self._fit_binary_classification(SVC(gamma=0.5)) svm_coreml = coremltools.converters.sklearn.convert(svm) @@ -76,23 +77,23 @@ def test_support_vector_classifier_binary_no_prob(self): # This should not have a probability output and will be a single node nodes = svm_onnx.graph.node self.assertEqual(len(nodes), 1) - self._check_model_outputs(svm_onnx, ['classLabel']) + self._check_model_outputs(svm_onnx, ["classLabel"]) dump_data_and_model(X, svm, svm_onnx, basename="CmlBinSVC-Out0") @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_support_vector_classifier_binary_with_prob(self): svm, X = self._fit_binary_classification(SVC(probability=True, gamma=0.5)) svm_coreml = coremltools.converters.sklearn.convert(svm) svm_onnx = convert(svm_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(svm_onnx is not None) self.validate_zipmap(svm_onnx) - self._check_model_outputs(svm_onnx, ['classLabel', 'classProbability']) + self._check_model_outputs(svm_onnx, ["classLabel", "classProbability"]) @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_support_vector_classifier_multiclass_no_prob(self): svm, X = self._fit_multi_classification(SVC(gamma=0.5)) svm_coreml = coremltools.converters.sklearn.convert(svm) @@ -100,18 +101,18 @@ def test_support_vector_classifier_multiclass_no_prob(self): self.assertTrue(svm_onnx is not None) nodes = svm_onnx.graph.node self.assertEqual(len(nodes), 1) - self._check_model_outputs(svm_onnx, ['classLabel']) + self._check_model_outputs(svm_onnx, ["classLabel"]) @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_support_vector_classifier_multiclass_with_prob(self): svm, X = self._fit_multi_classification(SVC(probability=True, gamma=0.5)) svm_coreml = coremltools.converters.sklearn.convert(svm) svm_onnx = convert(svm_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(svm_onnx is not None) self.validate_zipmap(svm_onnx) - self._check_model_outputs(svm_onnx, ['classLabel', 'classProbability']) + self._check_model_outputs(svm_onnx, ["classLabel", "classProbability"]) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_SupportVectorRegressorConverter.py b/tests/coreml/test_cml_SupportVectorRegressorConverter.py index ca274eca..67dbcd8a 100644 --- a/tests/coreml/test_cml_SupportVectorRegressorConverter.py +++ b/tests/coreml/test_cml_SupportVectorRegressorConverter.py @@ -4,12 +4,14 @@ Tests SupportVectorRegressor converter. """ import packaging.version as pv + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -27,19 +29,20 @@ class TestCoreMLSupportVectorRegressorConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_support_vector_regressor(self): X, y = make_regression(n_features=4, random_state=0) - svm = SVR(gamma=1./len(X)) + svm = SVR(gamma=1.0 / len(X)) svm.fit(X, y) svm_coreml = coremltools.converters.sklearn.convert(svm) svm_onnx = convert(svm_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(svm_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), svm, svm_onnx, basename="CmlRegSVR-Dec3") + dump_data_and_model( + X.astype(numpy.float32), svm, svm_onnx, basename="CmlRegSVR-Dec3" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py b/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py index 74a5cfdf..347806f0 100644 --- a/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py +++ b/tests/coreml/test_cml_TreeEnsembleClassifierConverter.py @@ -6,12 +6,14 @@ import unittest import packaging.version as pv import numpy + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -26,18 +28,17 @@ class TestCoreMLTreeEnsembleClassifierConverter(unittest.TestCase): - def validate_zipmap(self, model): # Validate that it contains a ZipMap nodes = model.graph.node - node = next((n for n in nodes if n.op_type == 'ZipMap'), None) + node = next((n for n in nodes if n.op_type == "ZipMap"), None) self.assertIsNotNone(node) self.assertEqual(len(node.output), 1) - self.assertTrue('classProbability' in node.output) + self.assertTrue("classProbability" in node.output) @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_tree_ensemble_classifier(self): X = numpy.array([[0, 1], [1, 1], [2, 0]], dtype=numpy.float32) y = [1, 0, 1] @@ -46,7 +47,9 @@ def test_tree_ensemble_classifier(self): model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) self.validate_zipmap(model_onnx) - dump_data_and_model(X, model, model_onnx, basename="CmlBinRandomForestClassifier") + dump_data_and_model( + X, model, model_onnx, basename="CmlBinRandomForestClassifier" + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_TreeEnsembleRegressorConverter.py b/tests/coreml/test_cml_TreeEnsembleRegressorConverter.py index cb0d186d..da2b296e 100644 --- a/tests/coreml/test_cml_TreeEnsembleRegressorConverter.py +++ b/tests/coreml/test_cml_TreeEnsembleRegressorConverter.py @@ -6,12 +6,14 @@ import unittest import packaging.version as pv import numpy + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer import coremltools @@ -27,18 +29,21 @@ class TestCoreMLTreeEnsembleRegressorConverter(unittest.TestCase): - @unittest.skipIf( - pv.Version(coremltools.__version__) > pv.Version("3.1"), - reason="untested") + pv.Version(coremltools.__version__) > pv.Version("3.1"), reason="untested" + ) def test_tree_ensemble_regressor(self): X, y = make_regression(n_features=4, random_state=0) model = RandomForestRegressor().fit(X, y) model_coreml = coremltools.converters.sklearn.convert(model) model_onnx = convert(model_coreml.get_spec(), target_opset=TARGET_OPSET) self.assertTrue(model_onnx is not None) - dump_data_and_model(X.astype(numpy.float32), model, model_onnx, - basename="CmlRegRandomForestRegressor-Dec3") + dump_data_and_model( + X.astype(numpy.float32), + model, + model_onnx, + basename="CmlRegRandomForestRegressor-Dec3", + ) if __name__ == "__main__": diff --git a/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py b/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py index 728e6361..983ec3d0 100644 --- a/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py +++ b/tests/coreml/test_cml_TreeEnsembleRegressorConverterXGBoost.py @@ -8,12 +8,14 @@ import unittest import numpy import pandas + try: from sklearn.impute import SimpleImputer as Imputer import sklearn.preprocessing - if not hasattr(sklearn.preprocessing, 'Imputer'): + + if not hasattr(sklearn.preprocessing, "Imputer"): # coremltools 3.1 does not work with scikit-learn 0.22 - setattr(sklearn.preprocessing, 'Imputer', Imputer) + setattr(sklearn.preprocessing, "Imputer", Imputer) except ImportError: from sklearn.preprocessing import Imputer from coremltools.converters.xgboost import convert as convert_xgb_to_coreml @@ -28,12 +30,12 @@ class TestCoreMLTreeEnsembleRegressorConverterXGBoost(unittest.TestCase): - @unittest.skipIf(True, reason="broken") def test_tree_ensemble_regressor_xgboost(self): - this = os.path.dirname(__file__) - data_train = pandas.read_csv(os.path.join(this, "xgboost.model.xgb.n4.d3.train.txt"), header=None) + data_train = pandas.read_csv( + os.path.join(this, "xgboost.model.xgb.n4.d3.train.txt"), header=None + ) X = data_train.iloc[:, 1:].values y = data_train.iloc[:, 0].values @@ -47,8 +49,12 @@ def test_tree_ensemble_regressor_xgboost(self): assert model_onnx is not None if sys.version_info[0] >= 3: # python 2.7 returns TypeError: can't pickle instancemethod objects - dump_data_and_model(X.astype(numpy.float32), model, model_onnx, - basename="CmlXGBoostRegressor-OneOff-Reshape") + dump_data_and_model( + X.astype(numpy.float32), + model, + model_onnx, + basename="CmlXGBoostRegressor-OneOff-Reshape", + ) if __name__ == "__main__": diff --git a/tests/h2o/test_h2o_converters.py b/tests/h2o/test_h2o_converters.py index a4d8a214..100ec5cc 100644 --- a/tests/h2o/test_h2o_converters.py +++ b/tests/h2o/test_h2o_converters.py @@ -1,11 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -""" -Tests h2o's tree-based methods' converters. -""" import unittest import os -import sys import numpy as np import pandas as pd from onnx.defs import onnx_opset_version @@ -16,7 +12,6 @@ import h2o from h2o import H2OFrame from h2o.estimators.gbm import H2OGradientBoostingEstimator -from h2o.exceptions import H2OConnectionError from h2o.estimators.random_forest import H2ORandomForestEstimator from onnxmltools.convert import convert_h2o from onnxmltools.utils import dump_data_and_model @@ -33,7 +28,7 @@ def _make_mojo(model, train, y=-1, force_y_numeric=False): x = list(range(0, train.ncol)) x.remove(y) model.train(x=x, y=y, training_frame=train) - folder = os.environ.get('ONNXTESTDUMP', 'tests/temp') + folder = os.environ.get("ONNXTESTDUMP", "tests/temp") if not os.path.exists(folder): os.makedirs(folder) return model.download_mojo(path=folder) @@ -65,7 +60,7 @@ def _prepare_one_hot(file, y, exclude_cols=None): train_frame = train.as_data_frame() train_encode = train_frame.loc[:, cols_to_encode] train_other = train_frame.loc[:, other_cols + [y]] - enc = OneHotEncoder(categories='auto', handle_unknown='ignore') + enc = OneHotEncoder(categories="auto", handle_unknown="ignore") enc.fit(train_encode) colnames = [] for cidx in range(len(cols_to_encode)): @@ -103,8 +98,11 @@ def _train_test_split_as_frames(x, y, is_str=False, is_classifier=False): def _train_classifier(model, n_classes, is_str=False, force_y_numeric=False): x, y = make_classification( - n_classes=n_classes, n_features=100, n_samples=1000, - random_state=42, n_informative=7 + n_classes=n_classes, + n_features=100, + n_samples=1000, + random_state=42, + n_informative=7, ) train, test = _train_test_split_as_frames(x, y, is_str, is_classifier=True) mojo_path = _make_mojo(model, train, force_y_numeric=force_y_numeric) @@ -112,16 +110,13 @@ def _train_classifier(model, n_classes, is_str=False, force_y_numeric=False): class H2OMojoWrapper: - def __init__(self, mojo_path, column_names=None): self._mojo_path = mojo_path self._mojo_model = h2o.upload_mojo(mojo_path) self._column_names = column_names def __getstate__(self): - return { - "path": self._mojo_path, - "colnames": self._column_names} + return {"path": self._mojo_path, "colnames": self._column_names} def __setstate__(self, state): self._mojo_path = state.path @@ -139,12 +134,11 @@ def predict_with_probabilities(self, data): else: return [ preds.iloc[:, 0].to_numpy().astype(np.str_), - preds.iloc[:, 1:].to_numpy() + preds.iloc[:, 1:].to_numpy(), ] class TestH2OModels(unittest.TestCase): - @classmethod def setUpClass(cls): h2o.init(port=54440) @@ -181,13 +175,15 @@ def test_h2o_regressor(self): onnx_model = _convert_mojo(mojo_path) self.assertIsNot(onnx_model, None) dump_data_and_model( - test, H2OMojoWrapper(mojo_path), - onnx_model, basename="H2OReg-Dec4") + test, H2OMojoWrapper(mojo_path), onnx_model, basename="H2OReg-Dec4" + ) @unittest.skipIf(True, reason="Failure with latest version of h2o") def test_h2o_regressor_cat(self): y = "IsDepDelayed" - train, test = _prepare_one_hot("airlines.csv", y, exclude_cols=["IsDepDelayed_REC"]) + train, test = _prepare_one_hot( + "airlines.csv", y, exclude_cols=["IsDepDelayed_REC"] + ) gbm = H2OGradientBoostingEstimator(ntrees=8, max_depth=5) mojo_path = _make_mojo(gbm, train, y=train.columns.index(y)) onnx_model = _convert_mojo(mojo_path) @@ -195,10 +191,14 @@ def test_h2o_regressor_cat(self): dump_data_and_model( test.values.astype(np.float32), H2OMojoWrapper(mojo_path, list(test.columns)), - onnx_model, basename="H2ORegCat-Dec4") + onnx_model, + basename="H2ORegCat-Dec4", + ) def test_h2o_classifier_multi_2class(self): - gbm = H2OGradientBoostingEstimator(ntrees=7, max_depth=5, distribution="multinomial") + gbm = H2OGradientBoostingEstimator( + ntrees=7, max_depth=5, distribution="multinomial" + ) mojo_path, test_data = _train_classifier(gbm, 2, is_str=True) with self.assertRaises(ValueError) as err: _convert_mojo(mojo_path) @@ -214,7 +214,9 @@ def test_h2o_classifier_bin_cat(self): dump_data_and_model( test.values.astype(np.float32), H2OMojoWrapper(mojo_path, list(test.columns)), - onnx_model, basename="H2OClassBinCat") + onnx_model, + basename="H2OClassBinCat", + ) def test_h2o_classifier_multi_cat(self): y = "fYear" @@ -227,7 +229,9 @@ def test_h2o_classifier_multi_cat(self): dump_data_and_model( test.values.astype(np.float32), H2OMojoWrapper(mojo_path, list(test.columns)), - onnx_model, basename="H2OClassMultiCat") + onnx_model, + basename="H2OClassMultiCat", + ) @unittest.skipIf(True, reason="Failure with latest version of h2o") def test_h2o_classifier_bin_str(self): @@ -236,17 +240,19 @@ def test_h2o_classifier_bin_str(self): onnx_model = _convert_mojo(mojo_path) self.assertIsNot(onnx_model, None) dump_data_and_model( - test_data, H2OMojoWrapper(mojo_path), onnx_model, - basename="H2OClassBinStr") + test_data, H2OMojoWrapper(mojo_path), onnx_model, basename="H2OClassBinStr" + ) def test_h2o_classifier_bin_int(self): gbm = H2OGradientBoostingEstimator(ntrees=8, max_depth=5) - mojo_path, test_data = _train_classifier(gbm, 2, is_str=False, force_y_numeric=True) + mojo_path, test_data = _train_classifier( + gbm, 2, is_str=False, force_y_numeric=True + ) onnx_model = _convert_mojo(mojo_path) self.assertIsNot(onnx_model, None) dump_data_and_model( - test_data, H2OMojoWrapper(mojo_path), onnx_model, - basename="H2OClassBinInt") + test_data, H2OMojoWrapper(mojo_path), onnx_model, basename="H2OClassBinInt" + ) def test_h2o_classifier_multi_str(self): gbm = H2OGradientBoostingEstimator(ntrees=10, max_depth=5) @@ -254,8 +260,11 @@ def test_h2o_classifier_multi_str(self): onnx_model = _convert_mojo(mojo_path) self.assertIsNot(onnx_model, None) dump_data_and_model( - test_data, H2OMojoWrapper(mojo_path), onnx_model, - basename="H2OClassMultiStr") + test_data, + H2OMojoWrapper(mojo_path), + onnx_model, + basename="H2OClassMultiStr", + ) def test_h2o_classifier_multi_int(self): gbm = H2OGradientBoostingEstimator(ntrees=9, max_depth=5) @@ -263,8 +272,11 @@ def test_h2o_classifier_multi_int(self): onnx_model = _convert_mojo(mojo_path) self.assertIsNot(onnx_model, None) dump_data_and_model( - test_data, H2OMojoWrapper(mojo_path), onnx_model, - basename="H2OClassMultiBin") + test_data, + H2OMojoWrapper(mojo_path), + onnx_model, + basename="H2OClassMultiBin", + ) def test_h2o_classifier_multi_discrete_int_labels(self): iris = load_iris() @@ -273,14 +285,16 @@ def test_h2o_classifier_multi_discrete_int_labels(self): y[y == 0] = 10 y[y == 1] = 20 y[y == 2] = -30 - train, test = _train_test_split_as_frames(x, y, is_str=False, is_classifier=True) + train, test = _train_test_split_as_frames( + x, y, is_str=False, is_classifier=True + ) gbm = H2OGradientBoostingEstimator(ntrees=7, max_depth=5) mojo_path = _make_mojo(gbm, train) onnx_model = _convert_mojo(mojo_path) self.assertIsNot(onnx_model, None) dump_data_and_model( - test, H2OMojoWrapper(mojo_path), onnx_model, - basename="H2OClassMultiDiscInt") + test, H2OMojoWrapper(mojo_path), onnx_model, basename="H2OClassMultiDiscInt" + ) if __name__ == "__main__": diff --git a/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py b/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py index 7565711c..80dd1d9b 100644 --- a/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py +++ b/tests/hummingbirdml/test_LightGbmTreeEnsembleConverters_hummingbird.py @@ -1,22 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 import unittest -import packaging.version as pv import lightgbm import numpy -from numpy.testing import assert_almost_equal from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from lightgbm import LGBMClassifier, LGBMRegressor from onnxruntime import InferenceSession from onnxmltools.convert.common.utils import hummingbird_installed from onnxmltools.convert.common.data_types import FloatTensorType -from onnxmltools.convert import convert_lightgbm -from onnxmltools.utils import dump_data_and_model -from onnxmltools.utils import dump_binary_classification, dump_multiple_classification -from onnxmltools.utils import dump_single_regression from onnxmltools.utils.tests_helper import convert_model +from onnxmltools.utils import dump_data_and_model TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) @@ -25,11 +20,11 @@ class TestLightGbmTreeEnsembleModelsHummingBird(unittest.TestCase): - @classmethod def setUpClass(cls): - print('BEGIN.') + print("BEGIN.") import torch + print(torch.device("cuda:0" if torch.cuda.is_available() else "cpu")) @classmethod @@ -43,15 +38,30 @@ def test_lightgbm_booster_classifier(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'binary', - 'n_estimators': 3, 'min_child_samples': 1, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, - target_opset=HUMMINGBIRD_TARGET_OPSET, - zipmap=False) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + without_onnx_ml=True, + target_opset=HUMMINGBIRD_TARGET_OPSET, + zipmap=False, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") def test_lightgbm_booster_classifier_zipmap(self): @@ -59,23 +69,47 @@ def test_lightgbm_booster_classifier_zipmap(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'binary', - 'n_estimators': 3, 'min_child_samples': 1, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], without_onnx_ml=False, - target_opset=HUMMINGBIRD_TARGET_OPSET) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + without_onnx_ml=False, + target_opset=HUMMINGBIRD_TARGET_OPSET, + ) assert "zipmap" in str(model_onnx).lower() with self.assertRaises(NotImplementedError): - convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, - target_opset=HUMMINGBIRD_TARGET_OPSET) - - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, - target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + without_onnx_ml=True, + target_opset=HUMMINGBIRD_TARGET_OPSET, + ) + + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + without_onnx_ml=True, + target_opset=HUMMINGBIRD_TARGET_OPSET, + zipmap=False, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") def test_lightgbm_booster_multi_classifier(self): @@ -83,18 +117,37 @@ def test_lightgbm_booster_multi_classifier(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1, 2, 2] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'multiclass', - 'n_estimators': 3, 'min_child_samples': 1, 'num_class': 3, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, - target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) - sess = InferenceSession(model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "multiclass", + "n_estimators": 3, + "min_child_samples": 1, + "num_class": 3, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + without_onnx_ml=True, + target_opset=HUMMINGBIRD_TARGET_OPSET, + zipmap=False, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) + sess = InferenceSession( + model_onnx.SerializeToString(), providers=["CPUExecutionProvider"] + ) out = sess.get_outputs() names = [o.name for o in out] - assert names == ['label', 'probabilities'] + assert names == ["label", "probabilities"] @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") def test_lightgbm_booster_regressor(self): @@ -102,30 +155,55 @@ def test_lightgbm_booster_regressor(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 1.1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'regression', - 'n_estimators': 3, 'min_child_samples': 1, 'max_depth': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based binary regressor', - [('input', FloatTensorType([None, 2]))], without_onnx_ml=True, - target_opset=HUMMINGBIRD_TARGET_OPSET, zipmap=False) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "regression", + "n_estimators": 3, + "min_child_samples": 1, + "max_depth": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based binary regressor", + [("input", FloatTensorType([None, 2]))], + without_onnx_ml=True, + target_opset=HUMMINGBIRD_TARGET_OPSET, + zipmap=False, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) # Base test implementation comparing ONNXML and ONNX models. def _test_lgbm(self, X, model, extra_config={}): # Create ONNX-ML model onnx_ml_model = convert_model( - model, 'lgbm-onnxml', [("input", FloatTensorType([None, X.shape[1]]))], - target_opset=TARGET_OPSET)[0] + model, + "lgbm-onnxml", + [("input", FloatTensorType([None, X.shape[1]]))], + target_opset=TARGET_OPSET, + )[0] # Create ONNX model onnx_model = convert_model( - model, 'lgbm-onnx', [("input", FloatTensorType([None, X.shape[1]]))], without_onnx_ml=True, - target_opset=TARGET_OPSET)[0] + model, + "lgbm-onnx", + [("input", FloatTensorType([None, X.shape[1]]))], + without_onnx_ml=True, + target_opset=TARGET_OPSET, + )[0] # Get the predictions for the ONNX-ML model session = InferenceSession(onnx_ml_model.SerializeToString()) - output_names = [session.get_outputs()[i].name for i in range(len(session.get_outputs()))] + output_names = [ + session.get_outputs()[i].name for i in range(len(session.get_outputs())) + ] onnx_ml_pred = [[] for i in range(len(output_names))] inputs = {session.get_inputs()[0].name: X} pred = session.run(output_names, inputs) @@ -152,15 +230,22 @@ def _test_regressor(self, X, model, rtol=1e-06, atol=1e-06, extra_config={}): onnx_ml_pred, onnx_pred, output_names = self._test_lgbm(X, model, extra_config) # Check that predicted values match - numpy.testing.assert_allclose(onnx_ml_pred[0], onnx_pred[0], rtol=rtol, atol=atol) + numpy.testing.assert_allclose( + onnx_ml_pred[0], onnx_pred[0], rtol=rtol, atol=atol + ) # Utility function for testing classification models. def _test_classifier(self, X, model, rtol=1e-06, atol=1e-06, extra_config={}): onnx_ml_pred, onnx_pred, output_names = self._test_lgbm(X, model, extra_config) - numpy.testing.assert_allclose(onnx_ml_pred[1], onnx_pred[1], rtol=rtol, atol=atol) # labels numpy.testing.assert_allclose( - list(map(lambda x: list(x.values()), onnx_ml_pred[0])), onnx_pred[0], rtol=rtol, atol=atol + onnx_ml_pred[1], onnx_pred[1], rtol=rtol, atol=atol + ) # labels + numpy.testing.assert_allclose( + list(map(lambda x: list(x.values()), onnx_ml_pred[0])), + onnx_pred[0], + rtol=rtol, + atol=atol, ) # probs # Regression test with 3 estimators. @@ -186,7 +271,9 @@ def _test_lightgbm_regressor1(self): # Regression test with 2 estimators. @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") def _test_lightgbm_regressor2(self): - model = LGBMRegressor(n_estimators=2, max_depth=1, min_child_samples=1, num_thread=1) + model = LGBMRegressor( + n_estimators=2, max_depth=1, min_child_samples=1, num_thread=1 + ) X = [[0, 1], [1, 1], [2, 0]] X = numpy.array(X, dtype=numpy.float32) y = numpy.array([100, -10, 50], dtype=numpy.float32) @@ -201,8 +288,14 @@ def _test_lightgbm_booster_regressor(self): y = [0, 1, 1.1] data = lightgbm.Dataset(X, label=y) model = lightgbm.train( - {"boosting_type": "gbdt", "objective": "regression", "n_estimators": 3, - "min_child_samples": 1, "max_depth": 1, 'num_thread': 1}, + { + "boosting_type": "gbdt", + "objective": "regression", + "n_estimators": 3, + "min_child_samples": 1, + "max_depth": 1, + "num_thread": 1, + }, data, ) self._test_regressor(X, model) @@ -234,7 +327,16 @@ def _test_lightgbm_booster_classifier(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({"boosting_type": "gbdt", "objective": "binary", "n_estimators": 3, "min_child_samples": 1, 'num_thread': 1}, data) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) self._test_classifier(X, model) # Binary classification test with 3 estimators and selecting boosting type zipmap. @@ -244,7 +346,16 @@ def _test_lightgbm_booster_classifier_zipmap(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({"boosting_type": "gbdt", "objective": "binary", "n_estimators": 3, "min_child_samples": 1, 'num_thread': 1}, data) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) self._test_classifier(X, model) # Multiclass classification test with 3 estimators. @@ -265,7 +376,14 @@ def _test_lightgbm_booster_multi_classifier(self): y = [0, 1, 0, 1, 2, 2] data = lightgbm.Dataset(X, label=y) model = lightgbm.train( - {"boosting_type": "gbdt", "objective": "multiclass", "n_estimators": 3, "min_child_samples": 1, "num_class": 3, 'num_thread': 1}, + { + "boosting_type": "gbdt", + "objective": "multiclass", + "n_estimators": 3, + "min_child_samples": 1, + "num_class": 3, + "num_thread": 1, + }, data, ) self._test_classifier(X, model) diff --git a/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py b/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py index 241642fe..a51c4873 100644 --- a/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py +++ b/tests/lightgbm/test_LightGbmTreeEnsembleConverters.py @@ -1,11 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 import unittest -import packaging.version as pv import lightgbm import numpy -import onnx from numpy.testing import assert_almost_equal from onnx.defs import onnx_opset_version from lightgbm import LGBMClassifier, LGBMRegressor @@ -23,7 +21,6 @@ class TestLightGbmTreeEnsembleModels(unittest.TestCase): - def test_lightgbm_classifier_binary(self): model = LGBMClassifier(n_estimators=3, min_child_samples=1, num_thread=1) dump_binary_classification(model) @@ -39,29 +36,40 @@ def test_lightgbm_classifier_zipmap(self): model = LGBMClassifier(n_estimators=3, min_child_samples=1, num_thread=1) model.fit(X, y) onx = convert_model( - model, 'dummy', input_types=[('X', FloatTensorType([None, X.shape[1]]))], - target_opset=TARGET_OPSET) + model, + "dummy", + input_types=[("X", FloatTensorType([None, X.shape[1]]))], + target_opset=TARGET_OPSET, + ) assert "zipmap" in str(onx).lower() def test_lightgbm_classifier_nozipmap(self): X = [[0, 1], [1, 1], [2, 0], [1, 2], [1, 5], [6, 2]] X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1, 1, 0] - model = LGBMClassifier(n_estimators=3, min_child_samples=1, max_depth=2, num_thread=1) + model = LGBMClassifier( + n_estimators=3, min_child_samples=1, max_depth=2, num_thread=1 + ) model.fit(X, y) onx = convert_model( - model, 'dummy', input_types=[('X', FloatTensorType([None, X.shape[1]]))], - zipmap=False, target_opset=TARGET_OPSET) + model, + "dummy", + input_types=[("X", FloatTensorType([None, X.shape[1]]))], + zipmap=False, + target_opset=TARGET_OPSET, + ) assert "zipmap" not in str(onx).lower() onxs = onx[0].SerializeToString() try: - sess = onnxruntime.InferenceSession(onxs, providers=["CPUExecutionProvider"]) + sess = onnxruntime.InferenceSession( + onxs, providers=["CPUExecutionProvider"] + ) except Exception as e: raise AssertionError( - "Model cannot be loaded by onnxruntime due to %r\n%s." % ( - e, onx[0])) + "Model cannot be loaded by onnxruntime due to %r\n%s." % (e, onx[0]) + ) exp = model.predict(X), model.predict_proba(X) - got = sess.run(None, {'X': X}) + got = sess.run(None, {"X": X}) assert_almost_equal(exp[0], got[0]) assert_almost_equal(exp[1], got[1]) @@ -69,21 +77,29 @@ def test_lightgbm_classifier_nozipmap2(self): X = [[0, 1], [1, 1], [2, 0], [1, 2], [1, 5], [6, 2]] X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1, 1, 0] - model = LGBMClassifier(n_estimators=3, min_child_samples=1, max_depth=2, num_thread=1) + model = LGBMClassifier( + n_estimators=3, min_child_samples=1, max_depth=2, num_thread=1 + ) model.fit(X, y) onx = convert_lightgbm( - model, 'dummy', initial_types=[('X', FloatTensorType([None, X.shape[1]]))], - zipmap=False, target_opset=TARGET_OPSET) + model, + "dummy", + initial_types=[("X", FloatTensorType([None, X.shape[1]]))], + zipmap=False, + target_opset=TARGET_OPSET, + ) assert "zipmap" not in str(onx).lower() onxs = onx.SerializeToString() try: - sess = onnxruntime.InferenceSession(onxs, providers=["CPUExecutionProvider"]) + sess = onnxruntime.InferenceSession( + onxs, providers=["CPUExecutionProvider"] + ) except Exception as e: raise AssertionError( - "Model cannot be loaded by onnxruntime due to %r\n%s." % ( - e, onx)) + "Model cannot be loaded by onnxruntime due to %r\n%s." % (e, onx) + ) exp = model.predict(X), model.predict_proba(X) - got = sess.run(None, {'X': X}) + got = sess.run(None, {"X": X}) assert_almost_equal(exp[0], got[0]) assert_almost_equal(exp[1], got[1]) @@ -96,7 +112,9 @@ def test_lightgbm_regressor1(self): dump_single_regression(model, suffix="1") def test_lightgbm_regressor2(self): - model = LGBMRegressor(n_estimators=2, max_depth=1, min_child_samples=1, num_thread=1) + model = LGBMRegressor( + n_estimators=2, max_depth=1, min_child_samples=1, num_thread=1 + ) dump_single_regression(model, suffix="2") def test_lightgbm_booster_classifier(self): @@ -104,81 +122,156 @@ def test_lightgbm_booster_classifier(self): X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'binary', - 'n_estimators': 3, 'min_child_samples': 1, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) def test_lightgbm_booster_classifier_nozipmap(self): X = [[0, 1], [1, 1], [2, 0], [1, 2]] X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'binary', - 'n_estimators': 3, 'min_child_samples': 1, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], - zipmap=False, target_opset=TARGET_OPSET) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + zipmap=False, + target_opset=TARGET_OPSET, + ) assert "zipmap" not in str(model_onnx).lower() - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) def test_lightgbm_booster_classifier_zipmap(self): X = [[0, 1], [1, 1], [2, 0], [1, 2]] X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'binary', - 'n_estimators': 3, 'min_child_samples': 1, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "binary", + "n_estimators": 3, + "min_child_samples": 1, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) assert "zipmap" in str(model_onnx).lower() - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) def test_lightgbm_booster_multi_classifier(self): X = [[0, 1], [1, 1], [2, 0], [1, 2], [-1, 2], [1, -2]] X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 0, 1, 2, 2] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'multiclass', - 'n_estimators': 3, 'min_child_samples': 1, 'num_class': 3, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based classifier', - [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "multiclass", + "n_estimators": 3, + "min_child_samples": 1, + "num_class": 3, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based classifier", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) try: from onnxruntime import InferenceSession except ImportError: # onnxruntime not installed (python 2.7) return - sess = InferenceSession(model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]) + sess = InferenceSession( + model_onnx.SerializeToString(), providers=["CPUExecutionProvider"] + ) out = sess.get_outputs() names = [o.name for o in out] - assert names == ['label', 'probabilities'] + assert names == ["label", "probabilities"] def test_lightgbm_booster_regressor(self): X = [[0, 1], [1, 1], [2, 0]] X = numpy.array(X, dtype=numpy.float32) y = [0, 1, 1.1] data = lightgbm.Dataset(X, label=y) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'regression', - 'n_estimators': 3, 'min_child_samples': 1, 'max_depth': 1, 'num_thread': 1}, - data) - model_onnx, prefix = convert_model(model, 'tree-based binary classifier', - [('input', FloatTensorType([None, 2]))], - target_opset=TARGET_OPSET) - dump_data_and_model(X, model, model_onnx, - basename=prefix + "BoosterBin" + model.__class__.__name__) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "regression", + "n_estimators": 3, + "min_child_samples": 1, + "max_depth": 1, + "num_thread": 1, + }, + data, + ) + model_onnx, prefix = convert_model( + model, + "tree-based binary classifier", + [("input", FloatTensorType([None, 2]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, + model, + model_onnx, + basename=prefix + "BoosterBin" + model.__class__.__name__, + ) if __name__ == "__main__": diff --git a/tests/lightgbm/test_LightGbmTreeEnsembleConvertersPkl.py b/tests/lightgbm/test_LightGbmTreeEnsembleConvertersPkl.py index c75ccd4b..4b9b4f42 100644 --- a/tests/lightgbm/test_LightGbmTreeEnsembleConvertersPkl.py +++ b/tests/lightgbm/test_LightGbmTreeEnsembleConvertersPkl.py @@ -17,40 +17,62 @@ TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) -ort_version = ".".join(onnxruntime.__version__.split('.')[:2]) +ort_version = ".".join(onnxruntime.__version__.split(".")[:2]) class TestLightGbmTreeEnsembleModelsPkl(unittest.TestCase): - - @unittest.skipIf(sys.platform.startswith('win'), reason="pickled on linux, may not work on windows") - @unittest.skipIf(sys.platform.startswith('lin'), reason="recover linux CI build, needs to be fixed") + @unittest.skipIf( + sys.platform.startswith("win"), + reason="pickled on linux, may not work on windows", + ) + @unittest.skipIf( + sys.platform.startswith("lin"), + reason="recover linux CI build, needs to be fixed", + ) def test_root_leave(self): this = os.path.abspath(os.path.dirname(__file__)) for name in ["example.pkl"]: with open(os.path.join(this, name), "rb") as f: model = pickle.load(f) - X = [[0., 1.], [1., 1.], [2., 0.]] + X = [[0.0, 1.0], [1.0, 1.0], [2.0, 0.0]] X = numpy.array(X, dtype=numpy.float32) - model_onnx = convert_lightgbm(model.steps[1][1], 'pkl1', [('input', FloatTensorType([1, X.shape[1]]))]) - dump_data_and_model(X, model.steps[1][1], model_onnx, basename="LightGbmPkl1") + model_onnx = convert_lightgbm( + model.steps[1][1], "pkl1", [("input", FloatTensorType([1, X.shape[1]]))] + ) + dump_data_and_model( + X, model.steps[1][1], model_onnx, basename="LightGbmPkl1" + ) - - @unittest.skipIf(sys.platform.startswith('win'), reason="pickled on linux, may not work on windows") - @unittest.skipIf(sys.platform.startswith('lin'), reason="recover linux CI build, needs to be fixed") + @unittest.skipIf( + sys.platform.startswith("win"), + reason="pickled on linux, may not work on windows", + ) + @unittest.skipIf( + sys.platform.startswith("lin"), + reason="recover linux CI build, needs to be fixed", + ) @unittest.skipIf(not hummingbird_installed(), reason="Hummingbird is not installed") @unittest.skipIf( - pv.Version(ort_version) < pv.Version('1.0.0'), reason="Hummingbird supports only latest versions of ORT" + pv.Version(ort_version) < pv.Version("1.0.0"), + reason="Hummingbird supports only latest versions of ORT", ) def test_root_leave_onnx_only(self): this = os.path.abspath(os.path.dirname(__file__)) for name in ["example.pkl"]: with open(os.path.join(this, name), "rb") as f: model = pickle.load(f) - X = [[0., 1.], [1., 1.], [2., 0.]] + X = [[0.0, 1.0], [1.0, 1.0], [2.0, 0.0]] X = numpy.array(X, dtype=numpy.float32) - model_onnx = convert_lightgbm(model.steps[1][1], 'pkl1', [('input', FloatTensorType([1, X.shape[1]]))], - without_onnx_ml=True, target_opset=TARGET_OPSET) - dump_data_and_model(X, model.steps[1][1], model_onnx, basename="LightGbmPkl1") + model_onnx = convert_lightgbm( + model.steps[1][1], + "pkl1", + [("input", FloatTensorType([1, X.shape[1]]))], + without_onnx_ml=True, + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + X, model.steps[1][1], model_onnx, basename="LightGbmPkl1" + ) if __name__ == "__main__": diff --git a/tests/lightgbm/test_LightGbmTreeEnsembleConverters_split.py b/tests/lightgbm/test_LightGbmTreeEnsembleConverters_split.py index 2bf441f4..1d7bb423 100644 --- a/tests/lightgbm/test_LightGbmTreeEnsembleConverters_split.py +++ b/tests/lightgbm/test_LightGbmTreeEnsembleConverters_split.py @@ -10,25 +10,20 @@ from sklearn.datasets import load_iris from sklearn.model_selection import train_test_split from onnxruntime import InferenceSession, __version__ as ort_version -from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER -from onnxmltools.convert.common.utils import hummingbird_installed from onnxmltools.convert.common.data_types import FloatTensorType from onnxmltools.convert import convert_lightgbm -from onnxmltools.utils import dump_data_and_model -from onnxmltools.utils import dump_binary_classification, dump_multiple_classification -from onnxmltools.utils import dump_single_regression -from onnxmltools.utils.tests_helper import convert_model TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) -ort_version = ".".join(ort_version.split('.')[:2]) +ort_version = ".".join(ort_version.split(".")[:2]) class TestLightGbmTreeEnsembleModelsSplit(unittest.TestCase): - - @unittest.skipIf(pv.Version(ort_version) < pv.Version('1.7.0'), - reason="Sum not implemented.") + @unittest.skipIf( + pv.Version(ort_version) < pv.Version("1.7.0"), + reason="Sum not implemented.", + ) def test_lgbm_regressor10(self): data = load_iris() X, y = data.data, data.target @@ -39,24 +34,26 @@ def test_lgbm_regressor10(self): expected = reg.predict(X_test) # float - init = [('X', FloatTensorType([None, X_train.shape[1]]))] + init = [("X", FloatTensorType([None, X_train.shape[1]]))] onx = convert_lightgbm(reg, None, init, target_opset=TARGET_OPSET) self.assertNotIn('op_type: "Sum"', str(onx)) oinf = InferenceSession(onx.SerializeToString()) - got1 = oinf.run(None, {'X': X_test})[0] + got1 = oinf.run(None, {"X": X_test})[0] # float split onx = convert_lightgbm(reg, None, init, split=2, target_opset=TARGET_OPSET) self.assertIn('op_type: "Sum"', str(onx)) oinf = InferenceSession(onx.SerializeToString()) - got2 = oinf.run(None, {'X': X_test})[0] + got2 = oinf.run(None, {"X": X_test})[0] # final check assert_almost_equal(expected, got1.ravel(), decimal=5) assert_almost_equal(expected, got2.ravel(), decimal=5) - @unittest.skipIf(pv.Version(ort_version) < pv.Version('1.7.0'), - reason="Sum not implemented.") + @unittest.skipIf( + pv.Version(ort_version) < pv.Version("1.7.0"), + reason="Sum not implemented.", + ) def test_lgbm_regressor(self): data = load_iris() X, y = data.data, data.target @@ -67,18 +64,18 @@ def test_lgbm_regressor(self): expected = reg.predict(X_test) # float - init = [('X', FloatTensorType([None, X_train.shape[1]]))] + init = [("X", FloatTensorType([None, X_train.shape[1]]))] onx = convert_lightgbm(reg, None, init, target_opset=TARGET_OPSET) self.assertNotIn('op_type: "Sum"', str(onx)) oinf = InferenceSession(onx.SerializeToString()) - got1 = oinf.run(None, {'X': X_test})[0] + got1 = oinf.run(None, {"X": X_test})[0] assert_almost_equal(expected, got1.ravel(), decimal=5) # float split onx = convert_lightgbm(reg, None, init, split=10, target_opset=TARGET_OPSET) self.assertIn('op_type: "Sum"', str(onx)) oinf = InferenceSession(onx.SerializeToString()) - got2 = oinf.run(None, {'X': X_test})[0] + got2 = oinf.run(None, {"X": X_test})[0] assert_almost_equal(expected, got2.ravel(), decimal=5) # final @@ -86,28 +83,45 @@ def test_lgbm_regressor(self): d2 = numpy.abs(expected.ravel() - got2.ravel()).mean() self.assertGreater(d1, d2) - @unittest.skipIf(pv.Version(ort_version) < pv.Version('1.7.0'), - reason="Sum not implemented.") + @unittest.skipIf( + pv.Version(ort_version) < pv.Version("1.7.0"), + reason="Sum not implemented.", + ) def test_lightgbm_booster_regressor(self): data = load_iris() X, y = data.data, data.target X_train, X_test, y_train, _ = train_test_split(X, y, random_state=0) data = lightgbm.Dataset(X_train, label=y_train) - model = lightgbm.train({'boosting_type': 'gbdt', 'objective': 'regression', - 'n_estimators': 100, 'max_depth': 2, 'num_thread': 1}, - data) + model = lightgbm.train( + { + "boosting_type": "gbdt", + "objective": "regression", + "n_estimators": 100, + "max_depth": 2, + "num_thread": 1, + }, + data, + ) expected = model.predict(X_test) - onx = convert_lightgbm(model, '', [('X', FloatTensorType([None, 4]))], target_opset=TARGET_OPSET) - onx10 = convert_lightgbm(model, '', [('X', FloatTensorType([None, 4]))], split=1, target_opset=TARGET_OPSET) + onx = convert_lightgbm( + model, "", [("X", FloatTensorType([None, 4]))], target_opset=TARGET_OPSET + ) + onx10 = convert_lightgbm( + model, + "", + [("X", FloatTensorType([None, 4]))], + split=1, + target_opset=TARGET_OPSET, + ) self.assertNotIn('op_type: "Sum"', str(onx)) oinf = InferenceSession(onx.SerializeToString()) - got1 = oinf.run(None, {'X': X_test.astype(numpy.float32)})[0] + got1 = oinf.run(None, {"X": X_test.astype(numpy.float32)})[0] assert_almost_equal(expected, got1.ravel(), decimal=5) self.assertIn('op_type: "Sum"', str(onx10)) oinf = InferenceSession(onx10.SerializeToString()) - got2 = oinf.run(None, {'X': X_test.astype(numpy.float32)})[0] + got2 = oinf.run(None, {"X": X_test.astype(numpy.float32)})[0] assert_almost_equal(expected, got2.ravel(), decimal=5) d1 = numpy.abs(expected.ravel() - got1.ravel()).mean() diff --git a/tests/lightgbm/test_lightgbm_missing_values.py b/tests/lightgbm/test_lightgbm_missing_values.py index 80f70d0b..d60af0bd 100644 --- a/tests/lightgbm/test_lightgbm_missing_values.py +++ b/tests/lightgbm/test_lightgbm_missing_values.py @@ -9,11 +9,13 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from lightgbm import LGBMRegressor -_N_DECIMALS=5 -_FRAC=0.9999 +_N_DECIMALS = 5 +_FRAC = 0.9999 _y = np.array([0, 0, 1, 1, 1]) -_X_train = np.array([[1.0, 0.0], [1.0, -1.0], [1.0, -1.0], [2.0, -1.0], [2.0, -1.0]], dtype=np.float32) +_X_train = np.array( + [[1.0, 0.0], [1.0, -1.0], [1.0, -1.0], [2.0, -1.0], [2.0, -1.0]], dtype=np.float32 +) _X_test = np.array([[1.0, np.nan]], dtype=np.float32) _INITIAL_TYPES = [("input", FloatTensorType([None, _X_train.shape[1]]))] @@ -21,33 +23,41 @@ class TestMissingValues(unittest.TestCase): - @staticmethod def _predict_with_onnx(model: ModelProto, X: np.array) -> np.array: session = InferenceSession(model.SerializeToString()) output_names = [s_output.name for s_output in session.get_outputs()] input_names = [s_input.name for s_input in session.get_inputs()] if len(input_names) > 1: - raise RuntimeError(f"Test expects one input. Found multiple inputs: {input_names}.") + raise RuntimeError( + f"Test expects one input. Found multiple inputs: {input_names}." + ) input_name = input_names[0] return session.run(output_names, {input_name: X})[0][:, 0] @staticmethod - def _assert_almost_equal(actual: np.array, desired: np.array, decimal: int=7, frac: float=1.0): + def _assert_almost_equal( + actual: np.array, desired: np.array, decimal: int = 7, frac: float = 1.0 + ): """ - Assert that almost all rows in actual and desired are almost equal to each other. - Similar to np.testing.assert_almost_equal but allows to define a fraction of rows to be almost + Assert that almost all rows in actual and + desired are almost equal to each other. + Similar to np.testing.assert_almost_equal but allows + to define a fraction of rows to be almost equal instead of expecting all rows to be almost equal. """ assert 0 <= frac <= 1, "frac must be in range(0, 1)." - success_abs = (abs(actual - desired) <= (10 ** -decimal)).sum() + success_abs = (abs(actual - desired) <= (10**-decimal)).sum() success_rel = success_abs / len(actual) - assert success_rel >= frac, f"Only {success_abs} out of {len(actual)} rows are almost equal to {decimal} decimals." - + assert success_rel >= frac, ( + f"Only {success_abs} out of {len(actual)} " + f"rows are almost equal to {decimal} decimals." + ) def test_missing_values(self): """ - Test that an ONNX model for a LGBM regressor that was trained without having seen missing values + Test that an ONNX model for a LGBM regressor that was + trained without having seen missing values correctly predicts rows that contain missing values. """ regressor = LGBMRegressor( @@ -56,9 +66,12 @@ def test_missing_values(self): min_data_in_leaf=1, n_estimators=1, learning_rate=1, - num_thread=1) + num_thread=1, + ) regressor.fit(_X_train, _y) - regressor_onnx: ModelProto = convert_lightgbm(regressor, initial_types=_INITIAL_TYPES, target_opset=TARGET_OPSET) + regressor_onnx: ModelProto = convert_lightgbm( + regressor, initial_types=_INITIAL_TYPES, target_opset=TARGET_OPSET + ) y_pred = regressor.predict(_X_test) y_pred_onnx = self._predict_with_onnx(regressor_onnx, _X_test) self._assert_almost_equal( @@ -69,6 +82,5 @@ def test_missing_values(self): ) - if __name__ == "__main__": unittest.main() diff --git a/tests/lightgbm/test_lightgbm_tree_structure.py b/tests/lightgbm/test_lightgbm_tree_structure.py index c7adcabc..afe17d7d 100644 --- a/tests/lightgbm/test_lightgbm_tree_structure.py +++ b/tests/lightgbm/test_lightgbm_tree_structure.py @@ -3,7 +3,8 @@ import unittest import copy from onnxmltools.convert.lightgbm.operator_converters.LightGbm import ( - modify_tree_for_rule_in_set) + modify_tree_for_rule_in_set, +) def count_nodes(tree, done=None): @@ -14,205 +15,246 @@ def count_nodes(tree, done=None): return 0 done[tid] = tree nb = 1 - if 'right_child' in tree: - nb += count_nodes(tree['right_child'], done) - if 'left_child' in tree: - nb += count_nodes(tree['left_child'], done) + if "right_child" in tree: + nb += count_nodes(tree["right_child"], done) + if "left_child" in tree: + nb += count_nodes(tree["left_child"], done) return nb tree2_t1 = { - 'decision_type': '==', - 'default_left': True, - 'internal_count': 1612, - 'internal_value': 0, - 'left_child': { - 'decision_type': '<=', - 'default_left': True, - 'internal_count': 1367, - 'internal_value': 1.02414, - 'left_child': { - 'decision_type': '<=', - 'default_left': True, - 'internal_count': 623, - 'internal_value': 1.02414, - 'left_child': { - 'leaf_count': 253, - 'leaf_index': 0, - 'leaf_value': 3.7749963852295396}, - 'missing_type': 'None', - 'right_child': { - 'leaf_count': 370, - 'leaf_index': 5, - 'leaf_value': 3.7749963852295396}, - 'split_feature': 3, - 'split_gain': 1.7763600157738027e-15, - 'split_index': 4, - 'threshold': 3.5000000000000004}, - 'missing_type': 'None', - 'right_child': { - 'decision_type': '<=', - 'default_left': True, - 'internal_count': 744, - 'internal_value': 1.02414, - 'left_child': { - 'leaf_count': 291, - 'leaf_index': 3, - 'leaf_value': 3.7749963852295396}, - 'missing_type': 'None', - 'right_child': { - 'leaf_count': 453, - 'leaf_index': 4, - 'leaf_value': 3.7749963852295396}, - 'split_feature': 3, - 'split_gain': 3.552710078910475e-15, - 'split_index': 3, - 'threshold': 3.5000000000000004}, - 'split_feature': 2, - 'split_gain': 7.105429898699844e-15, - 'split_index': 2, - 'threshold': 3.5000000000000004}, - 'missing_type': 'None', - 'right_child': { - 'decision_type': '<=', - 'default_left': True, - 'internal_count': 245, - 'internal_value': -5.7143, - 'left_child': { - 'leaf_count': 128, - 'leaf_index': 1, - 'leaf_value': 3.130106784685405}, - 'missing_type': 'None', - 'right_child': { - 'leaf_count': 117, - 'leaf_index': 2, - 'leaf_value': 3.7749963852295396}, - 'split_feature': 3, - 'split_gain': 234.05499267578125, - 'split_index': 1, - 'threshold': 6.500000000000001}, - 'split_feature': 2, - 'split_gain': 217.14300537109375, - 'split_index': 0, - 'threshold': '8||9||10'} + "decision_type": "==", + "default_left": True, + "internal_count": 1612, + "internal_value": 0, + "left_child": { + "decision_type": "<=", + "default_left": True, + "internal_count": 1367, + "internal_value": 1.02414, + "left_child": { + "decision_type": "<=", + "default_left": True, + "internal_count": 623, + "internal_value": 1.02414, + "left_child": { + "leaf_count": 253, + "leaf_index": 0, + "leaf_value": 3.7749963852295396, + }, + "missing_type": "None", + "right_child": { + "leaf_count": 370, + "leaf_index": 5, + "leaf_value": 3.7749963852295396, + }, + "split_feature": 3, + "split_gain": 1.7763600157738027e-15, + "split_index": 4, + "threshold": 3.5000000000000004, + }, + "missing_type": "None", + "right_child": { + "decision_type": "<=", + "default_left": True, + "internal_count": 744, + "internal_value": 1.02414, + "left_child": { + "leaf_count": 291, + "leaf_index": 3, + "leaf_value": 3.7749963852295396, + }, + "missing_type": "None", + "right_child": { + "leaf_count": 453, + "leaf_index": 4, + "leaf_value": 3.7749963852295396, + }, + "split_feature": 3, + "split_gain": 3.552710078910475e-15, + "split_index": 3, + "threshold": 3.5000000000000004, + }, + "split_feature": 2, + "split_gain": 7.105429898699844e-15, + "split_index": 2, + "threshold": 3.5000000000000004, + }, + "missing_type": "None", + "right_child": { + "decision_type": "<=", + "default_left": True, + "internal_count": 245, + "internal_value": -5.7143, + "left_child": { + "leaf_count": 128, + "leaf_index": 1, + "leaf_value": 3.130106784685405, + }, + "missing_type": "None", + "right_child": { + "leaf_count": 117, + "leaf_index": 2, + "leaf_value": 3.7749963852295396, + }, + "split_feature": 3, + "split_gain": 234.05499267578125, + "split_index": 1, + "threshold": 6.500000000000001, + }, + "split_feature": 2, + "split_gain": 217.14300537109375, + "split_index": 0, + "threshold": "8||9||10", +} tree2_t2 = { - 'decision_type': '<=', - 'default_left': True, - 'internal_count': 1612, - 'internal_value': 0, - 'left_child': { - 'leaf_count': 1367, - 'leaf_index': 0, - 'leaf_value': 0.05114685710677944}, - 'missing_type': 'None', - 'right_child': { - 'decision_type': '<=', - 'default_left': True, - 'internal_count': 245, - 'internal_value': -3.89759, - 'left_child': { - 'leaf_count': 128, - 'leaf_index': 1, - 'leaf_value': -0.3177225912983217}, - 'missing_type': 'None', - 'right_child': { - 'leaf_count': 117, - 'leaf_index': 2, - 'leaf_value': 0.05114685710677942}, - 'split_feature': 3, - 'split_gain': 93.09839630126953, - 'split_index': 1, - 'threshold': 6.500000000000001}, - 'split_feature': 2, - 'split_gain': 148.33299255371094, - 'split_index': 0, - 'threshold': 8.500000000000002} + "decision_type": "<=", + "default_left": True, + "internal_count": 1612, + "internal_value": 0, + "left_child": { + "leaf_count": 1367, + "leaf_index": 0, + "leaf_value": 0.05114685710677944, + }, + "missing_type": "None", + "right_child": { + "decision_type": "<=", + "default_left": True, + "internal_count": 245, + "internal_value": -3.89759, + "left_child": { + "leaf_count": 128, + "leaf_index": 1, + "leaf_value": -0.3177225912983217, + }, + "missing_type": "None", + "right_child": { + "leaf_count": 117, + "leaf_index": 2, + "leaf_value": 0.05114685710677942, + }, + "split_feature": 3, + "split_gain": 93.09839630126953, + "split_index": 1, + "threshold": 6.500000000000001, + }, + "split_feature": 2, + "split_gain": 148.33299255371094, + "split_index": 0, + "threshold": 8.500000000000002, +} -tree2 = {'average_output': False, - 'feature_names': ['c1', 'c2', 'c3', 'c4'], - 'label_index': 0, - 'max_feature_idx': 3, - 'name': 'tree', - 'num_class': 1, - 'num_tree_per_iteration': 1, - 'objective': 'binary sigmoid:1', - 'pandas_categorical': None, - 'tree_info': [{'num_cat': 0, - 'num_leaves': 6, - 'shrinkage': 1, - 'tree_index': 0, - 'tree_structure': tree2_t1}, - {'num_cat': 0, - 'num_leaves': 3, - 'shrinkage': 0.05, - 'tree_index': 1, - 'tree_structure': tree2_t2}], - 'version': 'v2'} +tree2 = { + "average_output": False, + "feature_names": ["c1", "c2", "c3", "c4"], + "label_index": 0, + "max_feature_idx": 3, + "name": "tree", + "num_class": 1, + "num_tree_per_iteration": 1, + "objective": "binary sigmoid:1", + "pandas_categorical": None, + "tree_info": [ + { + "num_cat": 0, + "num_leaves": 6, + "shrinkage": 1, + "tree_index": 0, + "tree_structure": tree2_t1, + }, + { + "num_cat": 0, + "num_leaves": 3, + "shrinkage": 0.05, + "tree_index": 1, + "tree_structure": tree2_t2, + }, + ], + "version": "v2", +} class TestLightGbmTreeStructur(unittest.TestCase): - def test_onnxrt_python_lightgbm_categorical(self): - val = {'decision_type': '==', - 'default_left': True, - 'internal_count': 6805, - 'internal_value': 0.117558, - 'left_child': {'leaf_count': 4293, - 'leaf_index': 18, - 'leaf_value': 0.003519117642745049}, - 'missing_type': 'None', - 'right_child': {'leaf_count': 2512, - 'leaf_index': 25, - 'leaf_value': 0.012305307958365394}, - 'split_feature': 24, - 'split_gain': 12.233599662780762, - 'split_index': 24, - 'threshold': '10||12||13'} + val = { + "decision_type": "==", + "default_left": True, + "internal_count": 6805, + "internal_value": 0.117558, + "left_child": { + "leaf_count": 4293, + "leaf_index": 18, + "leaf_value": 0.003519117642745049, + }, + "missing_type": "None", + "right_child": { + "leaf_count": 2512, + "leaf_index": 25, + "leaf_value": 0.012305307958365394, + }, + "split_feature": 24, + "split_gain": 12.233599662780762, + "split_index": 24, + "threshold": "10||12||13", + } - t2 = {'decision_type': '==', - 'default_left': True, - 'internal_count': 6805, - 'internal_value': 0.117558, - 'left_child': {'leaf_count': 4293, - 'leaf_index': 18, - 'leaf_value': 0.003519117642745049}, - 'missing_type': 'None', - 'right_child': {'decision_type': '==', - 'default_left': True, - 'internal_count': 6805, - 'internal_value': 0.117558, - 'left_child': { - 'leaf_count': 4293, - 'leaf_index': 18, - 'leaf_value': 0.003519117642745049}, - 'missing_type': 'None', - 'right_child': { - 'leaf_count': 2512, - 'leaf_index': 25, - 'leaf_value': 0.012305307958365394}, - 'split_feature': 24, - 'split_gain': 12.233599662780762, - 'split_index': 24, - 'threshold': 13}, - 'split_feature': 24, - 'split_gain': 12.233599662780762, - 'split_index': 24, - 'threshold': 12} + t2 = { + "decision_type": "==", + "default_left": True, + "internal_count": 6805, + "internal_value": 0.117558, + "left_child": { + "leaf_count": 4293, + "leaf_index": 18, + "leaf_value": 0.003519117642745049, + }, + "missing_type": "None", + "right_child": { + "decision_type": "==", + "default_left": True, + "internal_count": 6805, + "internal_value": 0.117558, + "left_child": { + "leaf_count": 4293, + "leaf_index": 18, + "leaf_value": 0.003519117642745049, + }, + "missing_type": "None", + "right_child": { + "leaf_count": 2512, + "leaf_index": 25, + "leaf_value": 0.012305307958365394, + }, + "split_feature": 24, + "split_gain": 12.233599662780762, + "split_index": 24, + "threshold": 13, + }, + "split_feature": 24, + "split_gain": 12.233599662780762, + "split_index": 24, + "threshold": 12, + } - exp = {'decision_type': '==', - 'default_left': True, - 'internal_count': 6805, - 'internal_value': 0.117558, - 'left_child': {'leaf_count': 4293, - 'leaf_index': 18, - 'leaf_value': 0.003519117642745049}, - 'missing_type': 'None', - 'right_child': t2, - 'split_feature': 24, - 'split_gain': 12.233599662780762, - 'split_index': 24, - 'threshold': 10} + exp = { + "decision_type": "==", + "default_left": True, + "internal_count": 6805, + "internal_value": 0.117558, + "left_child": { + "leaf_count": 4293, + "leaf_index": 18, + "leaf_value": 0.003519117642745049, + }, + "missing_type": "None", + "right_child": t2, + "split_feature": 24, + "split_gain": 12.233599662780762, + "split_index": 24, + "threshold": 10, + } nb1 = count_nodes(val) modify_tree_for_rule_in_set(val) @@ -220,15 +262,15 @@ def test_onnxrt_python_lightgbm_categorical(self): self.assertEqual(nb1, 3) self.assertEqual(nb2, 5) sval = str(val) - self.assertNotIn('||', sval) + self.assertNotIn("||", sval) self.maxDiff = None self.assertEqual(exp, val) def test_onnxrt_python_lightgbm_categorical2(self): val = copy.deepcopy(tree2) - nb1 = sum(count_nodes(t['tree_structure']) for t in val['tree_info']) + nb1 = sum(count_nodes(t["tree_structure"]) for t in val["tree_info"]) modify_tree_for_rule_in_set(val) - nb2 = sum(count_nodes(t['tree_structure']) for t in val['tree_info']) + nb2 = sum(count_nodes(t["tree_structure"]) for t in val["tree_info"]) self.assertEqual(nb1, 16) self.assertEqual(nb2, 18) diff --git a/tests/lightgbm/test_objective_functions.py b/tests/lightgbm/test_objective_functions.py index ffb2ae7c..220d7180 100644 --- a/tests/lightgbm/test_objective_functions.py +++ b/tests/lightgbm/test_objective_functions.py @@ -14,9 +14,9 @@ from lightgbm import LGBMRegressor -_N_ROWS=10_000 -_N_COLS=10 -_N_DECIMALS=5 +_N_ROWS = 10_000 +_N_COLS = 10 +_N_DECIMALS = 5 _FRAC = 0.9997 _X = pd.DataFrame(np.random.random(size=(_N_ROWS, _N_COLS))) @@ -31,19 +31,16 @@ class ObjectiveTest(unittest.TestCase): - - _objectives: Tuple[str] = ( - "regression", - "poisson", - "gamma", - "quantile" - ) + _objectives: Tuple[str] = ("regression", "poisson", "gamma", "quantile") @staticmethod def _calc_initial_types(X: DataFrame) -> List[Tuple[str, TensorType]]: dtypes = set(str(dtype) for dtype in X.dtypes) if len(dtypes) > 1: - raise RuntimeError(f"Test expects homogenous input matrix. Found multiple dtypes: {dtypes}.") + raise RuntimeError( + f"Test expects homogenous input matrix. " + f"Found multiple dtypes: {dtypes}." + ) dtype = dtypes.pop() tensor_type = _DTYPE_MAP[dtype] return [("input", tensor_type(X.shape))] @@ -54,38 +51,57 @@ def _predict_with_onnx(model: ModelProto, X: DataFrame) -> np.array: output_names = [s_output.name for s_output in session.get_outputs()] input_names = [s_input.name for s_input in session.get_inputs()] if len(input_names) > 1: - raise RuntimeError(f"Test expects one input. Found multiple inputs: {input_names}.") + raise RuntimeError( + f"Test expects one input. Found multiple inputs: {input_names}." + ) input_name = input_names[0] return session.run(output_names, {input_name: X.values})[0][:, 0] @staticmethod - def _assert_almost_equal(actual: np.array, desired: np.array, decimal: int=7, frac: float=1.0): + def _assert_almost_equal( + actual: np.array, desired: np.array, decimal: int = 7, frac: float = 1.0 + ): """ - Assert that almost all rows in actual and desired are almost equal to each other. + Assert that almost all rows in actual and desired + are almost equal to each other. - Similar to np.testing.assert_almost_equal but allows to define a fraction of rows to be almost + Similar to np.testing.assert_almost_equal but allows to define + a fraction of rows to be almost equal instead of expecting all rows to be almost equal. """ assert 0 <= frac <= 1, "frac must be in range(0, 1)." - success_abs = (abs(actual - desired) <= (10 ** -decimal)).sum() + success_abs = (abs(actual - desired) <= (10**-decimal)).sum() success_rel = success_abs / len(actual) - assert success_rel >= frac, f"Only {success_abs} out of {len(actual)} rows are almost equal to {decimal} decimals." - - @unittest.skipIf(tuple(int(ver) for ver in onnxruntime.__version__.split(".")[:2]) < (1, 3), "not supported in this library version") + assert success_rel >= frac, ( + f"Only {success_abs} out of {len(actual)} " + f"rows are almost equal to {decimal} decimals." + ) + + @unittest.skipIf( + tuple(int(ver) for ver in onnxruntime.__version__.split(".")[:2]) < (1, 3), + "not supported in this library version", + ) def test_objective(self): """ - Test if a LGBMRegressor a with certain objective (e.g. 'poisson') can be converted to ONNX - and whether the ONNX graph and the original model produce almost equal predictions. - - Note that this tests is a bit flaky because of precision differences with ONNX and LightGBM - and therefore sometimes fails randomly. In these cases, a retry should resolve the issue. + Test if a LGBMRegressor a with certain objective (e.g. 'poisson') + can be converted to ONNX + and whether the ONNX graph and the original model produce + almost equal predictions. + + Note that this tests is a bit flaky because of precision + differences with ONNX and LightGBM + and therefore sometimes fails randomly. In these cases, + a retry should resolve the issue. """ for objective in self._objectives: with self.subTest(X=_X, objective=objective): regressor = LGBMRegressor(objective=objective, num_thread=1) regressor.fit(_X, _Y) - regressor_onnx: ModelProto = convert_lightgbm(regressor, initial_types=self._calc_initial_types(_X), - target_opset=TARGET_OPSET) + regressor_onnx: ModelProto = convert_lightgbm( + regressor, + initial_types=self._calc_initial_types(_X), + target_opset=TARGET_OPSET, + ) y_pred = regressor.predict(_X) y_pred_onnx = self._predict_with_onnx(regressor_onnx, _X) self._assert_almost_equal( diff --git a/tests/sparkml/__init__.py b/tests/sparkml/__init__.py index d5acb7c3..64a1eec8 100644 --- a/tests/sparkml/__init__.py +++ b/tests/sparkml/__init__.py @@ -4,13 +4,17 @@ from tests.sparkml.sparkml_test_base import SparkMlTestCase except ImportError as e: import os + raise ImportError( "Unable to import local test submodule " "'tests.sparkml.sparkml_test_base'. " - "Current directory: %r, PYTHONPATH=%r, in folder=%r." % ( - os.getcwd(), os.environ.get('PYTHONPATH', '-'), - os.listdir("."))) from e + "Current directory: %r, PYTHONPATH=%r, in folder=%r." + % (os.getcwd(), os.environ.get("PYTHONPATH", "-"), os.listdir(".")) + ) from e from tests.sparkml.sparkml_test_utils import ( - start_spark, stop_spark, dump_data_and_sparkml_model, - dataframe_to_nparray) + start_spark, + stop_spark, + dump_data_and_sparkml_model, + dataframe_to_nparray, +) diff --git a/tests/sparkml/profile_pipeline.py b/tests/sparkml/profile_pipeline.py index 9febf72d..a21ca938 100644 --- a/tests/sparkml/profile_pipeline.py +++ b/tests/sparkml/profile_pipeline.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 import unittest -import sys import inspect import os import time @@ -14,27 +13,35 @@ from onnxmltools import convert_sparkml from onnxmltools.convert.sparkml import buildInitialTypesSimple, buildInputDictSimple from onnxmltools.utils.utils_backend import OnnxRuntimeAssertionError, compare_outputs -from onnxmltools.utils.utils_backend_onnxruntime import run_with_runtime, _compare_expected +from onnxmltools.utils.utils_backend_onnxruntime import ( + run_with_runtime, + _compare_expected, +) from tests.sparkml import SparkMlTestCase -class ProfileSparkmlPipeline(SparkMlTestCase): +class ProfileSparkmlPipeline1(SparkMlTestCase): def _get_spark_options(self): # add additional jar files before creating SparkSession - return {'spark.jars.packages': "ml.combust.mleap:mleap-spark_2.11:0.13.0"} + return {"spark.jars.packages": "ml.combust.mleap:mleap-spark_2.11:0.13.0"} -class ProfileSparkmlPipeline(SparkMlTestCase): - +class ProfileSparkmlPipeline2(SparkMlTestCase): def test_profile_sparkml_pipeline(self): - import mleap.pyspark - from mleap.pyspark.spark_support import SimpleSparkSerializer + pass # add additional jar files before creating SparkSession - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "AdultCensusIncomeOriginal.csv") - full_data = self.spark.read.format('csv') \ - .options(header='true', inferschema='true').load(input_path) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "AdultCensusIncomeOriginal.csv" + ) + full_data = ( + self.spark.read.format("csv") + .options(header="true", inferschema="true") + .load(input_path) + ) training_data, test_data = full_data.randomSplit([0.9, 0.1], seed=1) label = "income" @@ -50,11 +57,17 @@ def test_profile_sparkml_pipeline(self): feature_cols.append(feature_col) tmp_col = "-".join([key, "tmp"]) - si_xvars.append(StringIndexer(inputCol=key, outputCol=tmp_col, handleInvalid="skip")) - ohe_xvars.append(OneHotEncoder(inputCols=[tmp_col], outputCols=[feature_col], dropLast=False)) + si_xvars.append( + StringIndexer(inputCol=key, outputCol=tmp_col, handleInvalid="skip") + ) + ohe_xvars.append( + OneHotEncoder( + inputCols=[tmp_col], outputCols=[feature_col], dropLast=False + ) + ) else: feature_cols.append(key) - si_label = StringIndexer(inputCol=label, outputCol='label') + si_label = StringIndexer(inputCol=label, outputCol="label") assembler = VectorAssembler(inputCols=feature_cols, outputCol="features") lr = LogisticRegression(regParam=0.001) pipeline = Pipeline(stages=si_xvars + ohe_xvars + [si_label, assembler, lr]) @@ -64,9 +77,12 @@ def test_profile_sparkml_pipeline(self): test_data = test_data.limit(1) # create Spark and Onnx models model = pipeline.fit(training_data) - model_onnx = convert_sparkml(model, 'Sparkml Pipeline', buildInitialTypesSimple(test_data)) + model_onnx = convert_sparkml( + model, "Sparkml Pipeline", buildInitialTypesSimple(test_data) + ) # save Onnx model for runtime usage - if model_onnx is None: raise AssertionError("Failed to create the onnx model") + if model_onnx is None: + raise AssertionError("Failed to create the onnx model") model_path = os.path.join("tests", "profile_pipeline_model.onnx") with open(model_path, "wb") as f: f.write(model_onnx.SerializeToString()) @@ -112,10 +128,13 @@ def test_profile_sparkml_pipeline(self): expected = [ spark_prediction.toPandas().label.values.astype(numpy.float32), spark_prediction.toPandas().prediction.values.astype(numpy.float32), - spark_prediction.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype( - numpy.float32) + spark_prediction.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - _compare_expected(expected, output, session, model_path, decimal=5, onnx_shape=None) + _compare_expected( + expected, output, session, model_path, decimal=5, onnx_shape=None + ) gen_plot(spark_times, mleap_times, runtime_times) @@ -128,8 +147,12 @@ def _compare_mleap_pyspark(mleap_prediction, spark_prediction): msg = compare_outputs(spark_predicted_labels, mleap_predicted_labels, decimal=5) if msg: raise OnnxRuntimeAssertionError("Predictions in mleap and spark do not match") - spark_probability = spark_pandas.probability.apply(lambda x: pandas.Series(x.toArray())).values - mleap_probability = mleap_pandas.probability.apply(lambda x: pandas.Series(x.toArray())).values + spark_probability = spark_pandas.probability.apply( + lambda x: pandas.Series(x.toArray()) + ).values + mleap_probability = mleap_pandas.probability.apply( + lambda x: pandas.Series(x.toArray()) + ).values msg = compare_outputs(spark_probability, mleap_probability, decimal=5) if msg: raise OnnxRuntimeAssertionError("Probabilities in mleap and spark do not match") @@ -137,18 +160,18 @@ def _compare_mleap_pyspark(mleap_prediction, spark_prediction): def gen_plot(spark_times, mleap_times, runtime_times): import matplotlib.pyplot as pyplot - pyplot.hist(spark_times, label='pyspark') - pyplot.hist(mleap_times, label='MLeap') - pyplot.hist(runtime_times, label='onnxruntime') - pyplot.ylabel('Frequency') - pyplot.xlabel('Prediction Time(ms)') + + pyplot.hist(spark_times, label="pyspark") + pyplot.hist(mleap_times, label="MLeap") + pyplot.hist(runtime_times, label="onnxruntime") + pyplot.ylabel("Frequency") + pyplot.xlabel("Prediction Time(ms)") pyplot.legend() fig = pyplot.gcf() - #pyplot.show() + # pyplot.show() pyplot.draw() - fig.savefig('tests/spark-perf-histogram.png') + fig.savefig("tests/spark-perf-histogram.png") if __name__ == "__main__": unittest.main() - diff --git a/tests/sparkml/r_pipeline.py b/tests/sparkml/r_pipeline.py index 904b2f14..4ac0c090 100644 --- a/tests/sparkml/r_pipeline.py +++ b/tests/sparkml/r_pipeline.py @@ -2,36 +2,50 @@ import inspect import unittest -import sys import os import numpy import pandas from pyspark.ml import PipelineModel from onnxmltools import convert_sparkml from onnxmltools.convert.sparkml import buildInitialTypesSimple, buildInputDictSimple -from onnxmltools.utils.utils_backend_onnxruntime import run_with_runtime, _compare_expected +from onnxmltools.utils.utils_backend_onnxruntime import ( + run_with_runtime, + _compare_expected, +) from tests.sparkml import SparkMlTestCase class RPipeline(SparkMlTestCase): - def test_sparkml_r_pipeline(self): # add additional jar files before creating SparkSession - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "iris.csv") - data = self.spark.read.format('csv') \ - .options(header='true', inferschema='true').load(input_path) \ - .drop('_index_') + data = ( + self.spark.read.format("csv") + .options(header="true", inferschema="true") + .load(input_path) + .drop("_index_") + ) # read the model from disk pipeline_path = os.path.join(this_script_dir, "mlpmodel") model = PipelineModel.load(path=pipeline_path) # create Onnx model - model_onnx = convert_sparkml(model, 'Sparkml R Pipeline', buildInitialTypesSimple(data), spark_session=self.spark) + model_onnx = convert_sparkml( + model, + "Sparkml R Pipeline", + buildInitialTypesSimple(data), + spark_session=self.spark, + ) # save Onnx model for runtime usage - if model_onnx is None: raise AssertionError("Failed to create the onnx model") - model_path = os.path.join(this_script_dir, "tests_dump", "r_pipeline_model.onnx") + if model_onnx is None: + raise AssertionError("Failed to create the onnx model") + model_path = os.path.join( + this_script_dir, "tests_dump", "r_pipeline_model.onnx" + ) with open(model_path, "wb") as f: f.write(model_onnx.SerializeToString()) @@ -45,12 +59,14 @@ def test_sparkml_r_pipeline(self): expected = [ spark_prediction.toPandas().label.values.astype(numpy.float32), spark_prediction.toPandas().prediction.values.astype(numpy.float32), - spark_prediction.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype( - numpy.float32) + spark_prediction.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - _compare_expected(expected, output, session, model_path, decimal=5, onnx_shape=None) + _compare_expected( + expected, output, session, model_path, decimal=5, onnx_shape=None + ) if __name__ == "__main__": unittest.main() - diff --git a/tests/sparkml/sparkml_test_base.py b/tests/sparkml/sparkml_test_base.py index a7897a72..df99749d 100644 --- a/tests/sparkml/sparkml_test_base.py +++ b/tests/sparkml/sparkml_test_base.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -''' +""" Testcase Base class for SparkML tests -''' +""" import os import inspect import unittest @@ -14,12 +14,13 @@ def _get_spark_options(self): return None def setUp(self): - if os.name == 'nt' and os.environ.get('HADOOP_HOME') is None: - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - print('setting HADOOP_HOME to: ', this_script_dir) - os.environ['HADOOP_HOME'] = this_script_dir + if os.name == "nt" and os.environ.get("HADOOP_HOME") is None: + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + print("setting HADOOP_HOME to: ", this_script_dir) + os.environ["HADOOP_HOME"] = this_script_dir self.spark = start_spark(self._get_spark_options()) def tearDown(self): stop_spark(self.spark) - diff --git a/tests/sparkml/sparkml_test_utils.py b/tests/sparkml/sparkml_test_utils.py index 45db647f..3ea1ddf0 100644 --- a/tests/sparkml/sparkml_test_utils.py +++ b/tests/sparkml/sparkml_test_utils.py @@ -2,20 +2,20 @@ import pickle import os -import warnings import sys import numpy import onnxruntime from onnxruntime.capi.onnxruntime_pybind11_state import InvalidArgument, Fail import pyspark -from pyspark import SparkConf from pyspark.sql import SparkSession -from pyspark.ml.linalg import VectorUDT -from pyspark.sql.types import ArrayType, FloatType, DoubleType from onnxmltools.utils.utils_backend import ( - compare_backend, extract_options, is_backend_enabled, - OnnxRuntimeAssertionError, compare_outputs, ExpectedAssertionError) -from onnxmltools.utils.utils_backend_onnxruntime import _create_column + compare_backend, + extract_options, + is_backend_enabled, + OnnxRuntimeAssertionError, + compare_outputs, + ExpectedAssertionError, +) def start_spark(options): @@ -26,7 +26,7 @@ def start_spark(options): builder = SparkSession.builder.appName("pyspark-unittesting").master("local[1]") if options: - for k,v in options.items(): + for k, v in options.items(): builder.config(k, v) spark = builder.getOrCreate() # spark.sparkContext.setLogLevel("ALL") @@ -37,10 +37,19 @@ def stop_spark(spark): spark.sparkContext.stop() -def save_data_models(input, expected, model, onnx_model, basename="model", folder=None, - save_spark_model=False, pickle_spark_model=False, pickle_data=False): +def save_data_models( + input, + expected, + model, + onnx_model, + basename="model", + folder=None, + save_spark_model=False, + pickle_spark_model=False, + pickle_data=False, +): if folder is None: - folder = os.environ.get('ONNXTESTDUMP', 'tests_dump') + folder = os.environ.get("ONNXTESTDUMP", "tests_dump") if not os.path.exists(folder): os.makedirs(folder) @@ -71,7 +80,7 @@ def save_data_models(input, expected, model, onnx_model, basename="model", folde def run_onnx_model(output_names, input, onnx_model): - sess = onnxruntime.InferenceSession(onnx_model, providers=['CPUExecutionProvider']) + sess = onnxruntime.InferenceSession(onnx_model, providers=["CPUExecutionProvider"]) if isinstance(input, dict): inputs = input elif isinstance(input, list): @@ -83,11 +92,13 @@ def run_onnx_model(output_names, input, onnx_model): inputs = {inp[0].name: input} else: raise OnnxRuntimeAssertionError( - "Wrong number of inputs onnx {0} != original shape {1}, onnx='{2}'".format( - len(inp), input.shape, onnx_model)) + "Wrong number of inputs onnx {0} != original shape " + "{1}, onnx='{2}'".format(len(inp), input.shape, onnx_model) + ) else: raise OnnxRuntimeAssertionError( - "Dict or list is expected, not {0}".format(type(input))) + "Dict or list is expected, not {0}".format(type(input)) + ) for k in inputs: if isinstance(inputs[k], list): @@ -102,13 +113,14 @@ def run_onnx_model(output_names, input, onnx_model): rows.append("output: {} - {} - {}".format(inp.name, inp.type, inp.shape)) rows.append("REQUIRED: {}".format(output_names)) for k, v in sorted(inputs.items()): - if hasattr(v, 'shape'): + if hasattr(v, "shape"): rows.append("{}={}-{}-{}".format(k, v.shape, v.dtype, v)) else: rows.append("{}={}".format(k, v)) raise AssertionError( - "Unable to run onnxruntime\n{}".format("\n".join(rows))) from e - + "Unable to run onnxruntime\n{}".format("\n".join(rows)) + ) from e + output_shapes = [_.shape for _ in sess.get_outputs()] return output, output_shapes @@ -119,13 +131,17 @@ def compare_results(expected, output, decimal=5): if isinstance(output, list): if len(expected) != len(output): raise OnnxRuntimeAssertionError( - "Unexpected number of outputs: expected={0}, got={1}".format(len(expected), len(output))) + "Unexpected number of outputs: expected={0}, got={1}".format( + len(expected), len(output) + ) + ) for exp, out in zip(expected, output): compare_results(exp, out, decimal=decimal) tested += 1 else: raise OnnxRuntimeAssertionError( - "Type mismatch: output type is {0}".format(type(output))) + "Type mismatch: output type is {0}".format(type(output)) + ) elif isinstance(expected, dict): if not isinstance(output, dict): raise OnnxRuntimeAssertionError("Type mismatch fo") @@ -134,12 +150,15 @@ def compare_results(expected, output, decimal=5): continue msg = compare_outputs(expected[k], v, decimal=decimal) if msg: - raise OnnxRuntimeAssertionError("Unexpected output '{0}': \n{2}".format(k, msg)) + raise OnnxRuntimeAssertionError( + "Unexpected output '{0}': \n{1}".format(k, msg) + ) tested += 1 elif isinstance(expected, numpy.ndarray): if isinstance(output, list): if expected.shape[0] == len(output) and isinstance(output[0], dict): import pandas + output = pandas.DataFrame(output) output = output[list(sorted(output.columns))] output = output.values @@ -149,37 +168,49 @@ def compare_results(expected, output, decimal=5): if len(ex) > 70: ex = ex[:70] + "..." raise OnnxRuntimeAssertionError( - "More than one output when 1 is expected\n{0}".format(ex)) + "More than one output when 1 is expected\n{0}".format(ex) + ) output = output[-1] if not isinstance(output, numpy.ndarray): raise OnnxRuntimeAssertionError( - "output must be an array not {0}".format(type(output))) + "output must be an array not {0}".format(type(output)) + ) msg = compare_outputs(expected, output, decimal=decimal) if isinstance(msg, ExpectedAssertionError): raise msg if msg: - raise OnnxRuntimeAssertionError( - "Unexpected output\n{}".format(msg)) + raise OnnxRuntimeAssertionError("Unexpected output\n{}".format(msg)) tested += 1 else: from scipy.sparse.csr import csr_matrix + if isinstance(expected, csr_matrix): # DictVectorizer one_array = numpy.array(output) msg = compare_outputs(expected.todense(), one_array, decimal=decimal) if msg: - raise OnnxRuntimeAssertionError("Unexpected output\n{1}".format(msg)) + raise OnnxRuntimeAssertionError("Unexpected output\n{0}".format(msg)) tested += 1 else: raise OnnxRuntimeAssertionError( - "Unexpected type for expected output ({0})".format(type(expected))) + "Unexpected type for expected output ({0})".format(type(expected)) + ) if tested == 0: raise OnnxRuntimeAssertionError("No test for model") -def dump_data_and_sparkml_model(input, expected, model, onnx=None, basename="model", folder=None, - backend="onnxruntime", context=None, - allow_failure=None, verbose=False): +def dump_data_and_sparkml_model( + input, + expected, + model, + onnx=None, + basename="model", + folder=None, + backend="onnxruntime", + context=None, + allow_failure=None, + verbose=False, +): """ Saves data with pickle, saves the model with pickle and *onnx*, runs and saves the predictions for the given model. @@ -218,14 +249,17 @@ def dump_data_and_sparkml_model(input, expected, model, onnx=None, basename="mod * ``-CannotLoad``: the model can be converted but the runtime cannot load it * ``-Dec3``: compares expected and computed outputs up to 3 decimals (5 by default) * ``-Dec4``: compares expected and computed outputs up to 4 decimals (5 by default) - * ``-NoProb``: The original models computed probabilites for two classes *size=(N, 2)* - but the runtime produces a vector of size *N*, the test will compare the second column + * ``-NoProb``: The original models computed probabilites + for two classes *size=(N, 2)* + but the runtime produces a vector of size *N*, + the test will compare the second column to the column * ``-OneOff``: the ONNX runtime cannot computed the prediction for several inputs, it must be called for each of them and computed output. * ``-Out0``: only compares the first output on both sides - * ``-Reshape``: merges all outputs into one single vector and resizes it before comparing + * ``-Reshape``: merges all outputs into one single vector + and resizes it before comparing * ``-SkipDim1``: before comparing expected and computed output, arrays with a shape like *(2, 1, 2)* becomes *(2, 2)* @@ -236,11 +270,11 @@ def dump_data_and_sparkml_model(input, expected, model, onnx=None, basename="mod runtime_test = dict(model=model, data=input) if folder is None: - folder = os.environ.get('ONNXTESTDUMP', 'tests_dump') + folder = os.environ.get("ONNXTESTDUMP", "tests_dump") if not os.path.exists(folder): os.makedirs(folder) - runtime_test['expected'] = expected + runtime_test["expected"] = expected names = [] dest = os.path.join(folder, basename + ".expected.pkl") @@ -273,19 +307,13 @@ def dump_data_and_sparkml_model(input, expected, model, onnx=None, basename="mod continue if isinstance(allow_failure, str): raise NotImplementedError("allow_failure is deprecated.") - if allow is None: - output = compare_backend(b, runtime_test, options=extract_options(basename), - context=context, verbose=verbose) - else: - try: - output = compare_backend(b, runtime_test, options=extract_options(basename), - context=context, verbose=verbose) - except AssertionError as e: - if isinstance(allow, bool) and allow: - warnings.warn("Issue with '{0}' due to {1}".format(basename, e)) - continue - else: - raise e + output = compare_backend( + b, + runtime_test, + options=extract_options(basename), + context=context, + verbose=verbose, + ) if output is not None: dest = os.path.join(folder, basename + ".backend.{0}.pkl".format(b)) names.append(dest) @@ -297,12 +325,18 @@ def dump_data_and_sparkml_model(input, expected, model, onnx=None, basename="mod def dataframe_to_nparray(df): from pyspark.ml.linalg import VectorUDT + schema = df.schema npcols = [] for i in range(0, len(df.columns)): if isinstance(schema.fields[i].dataType, VectorUDT): - npcols.append(df.select(df.columns[i]).toPandas().apply( - lambda x : numpy.array(x[0].toArray())).as_matrix().reshape(-1, 1)) + npcols.append( + df.select(df.columns[i]) + .toPandas() + .apply(lambda x: numpy.array(x[0].toArray())) + .as_matrix() + .reshape(-1, 1) + ) else: npcols.append(df.select(df.columns[i]).collect()) return numpy.array(npcols) diff --git a/tests/sparkml/test_PCA.py b/tests/sparkml/test_PCA.py index b72fe4c9..3934e81c 100644 --- a/tests/sparkml/test_PCA.py +++ b/tests/sparkml/test_PCA.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,30 +22,48 @@ class TestSparkmlPCA(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_polynomial_expansion(self): - data = self.spark.createDataFrame([ - (Vectors.sparse(5, [(1, 1.0), (3, 7.0)]),), - (Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]),), - (Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0]),) - ], ["features"]) + data = self.spark.createDataFrame( + [ + (Vectors.sparse(5, [(1, 1.0), (3, 7.0)]),), + (Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]),), + (Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0]),), + ], + ["features"], + ) pca = PCA(k=2, inputCol="features", outputCol="pca_features") model = pca.fit(data) # the input name should match that of what StringIndexer.inputCol feature_count = data.first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml PCA', [('features', FloatTensorType([None, feature_count]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml PCA", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().pca_features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlPCA") + expected = ( + predicted.toPandas() + .pca_features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlPCA" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['pca_features'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["pca_features"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_aft_survival_regression.py b/tests/sparkml/test_aft_survival_regression.py index d0867065..19d5dfad 100644 --- a/tests/sparkml/test_aft_survival_regression.py +++ b/tests/sparkml/test_aft_survival_regression.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,31 +22,42 @@ class TestSparkmAFTSurvivalRegression(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_aft_regression_survival(self): - data = self.spark.createDataFrame([ - (1.0, Vectors.dense(1.0), 1.0), - (1e-40, Vectors.sparse(1, [], []), 0.0) - ], ["label", "features", "censor"]) + data = self.spark.createDataFrame( + [(1.0, Vectors.dense(1.0), 1.0), (1e-40, Vectors.sparse(1, [], []), 0.0)], + ["label", "features", "censor"], + ) gbt = AFTSurvivalRegression() model = gbt.fit(data) feature_count = data.first()[1].size - model_onnx = convert_sparkml(model, 'Sparkml AFTSurvivalRegression', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml AFTSurvivalRegression", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlAFTSurvivalRegression") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlAFTSurvivalRegression", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_binarizer.py b/tests/sparkml/test_binarizer.py index 2b154235..e4b638b7 100644 --- a/tests/sparkml/test_binarizer.py +++ b/tests/sparkml/test_binarizer.py @@ -8,7 +8,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -16,25 +20,31 @@ class TestSparkmlBinarizer(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_binarizer(self): - data = self.spark.createDataFrame([(0, 0.1), (1, 0.8), (2, 0.2) ], ["id", "feature"]) - model = Binarizer(inputCol='feature', outputCol='binarized') + data = self.spark.createDataFrame( + [(0, 0.1), (1, 0.8), (2, 0.2)], ["id", "feature"] + ) + model = Binarizer(inputCol="feature", outputCol="binarized") # the input name should match that of what StringIndexer.inputCol - model_onnx = convert_sparkml(model, 'Sparkml Binarizer', [('feature', FloatTensorType([None, 1]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Binarizer", + [("feature", FloatTensorType([None, 1]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) expected = predicted.select("binarized").toPandas().values.astype(numpy.float32) - data_np = data.select('feature').toPandas().values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlBinarizer") + data_np = data.select("feature").toPandas().values.astype(numpy.float32) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlBinarizer" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['binarized'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["binarized"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_bucketed_random_projection_lsh.py b/tests/sparkml/test_bucketed_random_projection_lsh.py index bef85c32..78139996 100644 --- a/tests/sparkml/test_bucketed_random_projection_lsh.py +++ b/tests/sparkml/test_bucketed_random_projection_lsh.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,39 +22,66 @@ class TestBucketedRandomProjectionLSH(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_bucketed_random_projection_lsh(self): - data = self.spark.createDataFrame([ - (0, Vectors.dense([-1.0, -1.0 ]),), - (1, Vectors.dense([-1.0, 1.0 ]),), - (2, Vectors.dense([1.0, -1.0 ]),), - (3, Vectors.dense([1.0, 1.0]),) - ], ["id", "features"]) - mh = BucketedRandomProjectionLSH(inputCol="features", outputCol="hashes", seed=12345, bucketLength=1.0) + data = self.spark.createDataFrame( + [ + ( + 0, + Vectors.dense([-1.0, -1.0]), + ), + ( + 1, + Vectors.dense([-1.0, 1.0]), + ), + ( + 2, + Vectors.dense([1.0, -1.0]), + ), + ( + 3, + Vectors.dense([1.0, 1.0]), + ), + ], + ["id", "features"], + ) + mh = BucketedRandomProjectionLSH( + inputCol="features", outputCol="hashes", seed=12345, bucketLength=1.0 + ) model = mh.fit(data) feature_count = data.first()[1].size - model_onnx = convert_sparkml(model, 'Sparkml BucketedRandomProjectionLSH', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml BucketedRandomProjectionLSH", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply( - lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ - predicted.toPandas().hashes.apply(lambda x: pandas.Series(x) - .map(lambda y: y.values[0])).values.astype(numpy.float32), + predicted.toPandas() + .hashes.apply(lambda x: pandas.Series(x).map(lambda y: y.values[0])) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlBucketedRandomProjectionLSH") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlBucketedRandomProjectionLSH", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['hashes'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["hashes"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_bucketizer.py b/tests/sparkml/test_bucketizer.py index 1868fcaa..e051e078 100644 --- a/tests/sparkml/test_bucketizer.py +++ b/tests/sparkml/test_bucketizer.py @@ -8,7 +8,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -16,26 +20,33 @@ class TestSparkmlBucketizer(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_bucketizer(self): values = [(0.1,), (0.4,), (1.2,), (1.5,)] data = self.spark.createDataFrame(values, ["features"]) - model = Bucketizer(splits=[-float("inf"), 0.5, 1.4, float("inf")], inputCol="features", outputCol="buckets") - - feature_count = len(data.select('features').first()) - model_onnx = convert_sparkml(model, 'Sparkml Bucketizer', [ - ('features', FloatTensorType([None, feature_count])) - ], target_opset=TARGET_OPSET) + model = Bucketizer( + splits=[-float("inf"), 0.5, 1.4, float("inf")], + inputCol="features", + outputCol="buckets", + ) + + feature_count = len(data.select("features").first()) + model_onnx = convert_sparkml( + model, + "Sparkml Bucketizer", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.setHandleInvalid("error").transform(data) expected = predicted.select("buckets").toPandas().values.astype(numpy.float32) data_np = [data.toPandas().values.astype(numpy.float32)] - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlBucketizer") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlBucketizer" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['buckets'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["buckets"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_chi_sql_selector.py b/tests/sparkml/test_chi_sql_selector.py index a221c510..acf7ee88 100644 --- a/tests/sparkml/test_chi_sql_selector.py +++ b/tests/sparkml/test_chi_sql_selector.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,31 +22,48 @@ class TestSparkmlChiSqSelector(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_chi_sq_selector(self): - data = self.spark.createDataFrame([ - (Vectors.dense([0.0, 0.0, 18.0, 1.0]), 1.0), - (Vectors.dense([0.0, 1.0, 12.0, 0.0]), 0.0), - (Vectors.dense([1.0, 0.0, 15.0, 0.1]), 0.0) - ], ["features", "label"]) + data = self.spark.createDataFrame( + [ + (Vectors.dense([0.0, 0.0, 18.0, 1.0]), 1.0), + (Vectors.dense([0.0, 1.0, 12.0, 0.0]), 0.0), + (Vectors.dense([1.0, 0.0, 15.0, 0.1]), 0.0), + ], + ["features", "label"], + ) selector = ChiSqSelector(numTopFeatures=1, outputCol="selectedFeatures") model = selector.fit(data) # the input name should match that of what StringIndexer.inputCol feature_count = data.first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml ChiSqSelector', [('features', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml ChiSqSelector", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().selectedFeatures.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlChiSqSelector") + expected = ( + predicted.toPandas() + .selectedFeatures.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlChiSqSelector" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['selectedFeatures'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["selectedFeatures"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_count_vectorizer.py b/tests/sparkml/test_count_vectorizer.py index b24bc1cf..19a59485 100644 --- a/tests/sparkml/test_count_vectorizer.py +++ b/tests/sparkml/test_count_vectorizer.py @@ -9,79 +9,126 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) -class TestSparkmlCountVectorizer(SparkMlTestCase): - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") +class TestSparkmlCountVectorizer(SparkMlTestCase): + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_count_vectorizer_default(self): - data = self.spark.createDataFrame([ - ("A B C".split(" "), ), - ("A B B C A".split(" "), ), - ], ["text"]) - count_vec = CountVectorizer(inputCol="text", outputCol="result", minTF=1.0, binary=False) + data = self.spark.createDataFrame( + [ + ("A B C".split(" "),), + ("A B B C A".split(" "),), + ], + ["text"], + ) + count_vec = CountVectorizer( + inputCol="text", outputCol="result", minTF=1.0, binary=False + ) model: CountVectorizerModel = count_vec.fit(data) result = model.transform(data) - model_onnx = convert_sparkml(model, 'Sparkml CountVectorizer', [('text', StringTensorType([None, None]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml CountVectorizer", + [("text", StringTensorType([None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) - + data_pd = data.toPandas() data_np = { "text": data_pd.text.apply(lambda x: pandas.Series(x)).values.astype(str), } expected = { - "prediction_result": numpy.asarray(result.toPandas().result.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32)), + "prediction_result": numpy.asarray( + result.toPandas() + .result.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ), } - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlCountVectorizerModel_Default") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlCountVectorizerModel_Default", + ) onnx_model_path = paths[-1] - - output_names = ['result'] + + output_names = ["result"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) actual_output = dict(zip(output_names, output)) - + assert output_shapes[0] == [None, 3] - compare_results(expected["prediction_result"], actual_output["result"], decimal=5) + compare_results( + expected["prediction_result"], actual_output["result"], decimal=5 + ) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_count_vectorizer_binary(self): - data = self.spark.createDataFrame([ - ("A B C".split(" "), ), - ("A B B C A".split(" "), ), - ("B B B D".split(" "), ), - ], ["text"]) - count_vec = CountVectorizer(inputCol="text", outputCol="result", minTF=2.0, binary=True) + data = self.spark.createDataFrame( + [ + ("A B C".split(" "),), + ("A B B C A".split(" "),), + ("B B B D".split(" "),), + ], + ["text"], + ) + count_vec = CountVectorizer( + inputCol="text", outputCol="result", minTF=2.0, binary=True + ) model: CountVectorizerModel = count_vec.fit(data) result = model.transform(data) - model_onnx = convert_sparkml(model, 'Sparkml CountVectorizer', [('text', StringTensorType([None, None]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml CountVectorizer", + [("text", StringTensorType([None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) - + data_pd = data.toPandas() data_np = { "text": data_pd.text.apply(lambda x: pandas.Series(x)).values.astype(str), } expected = { - "prediction_result": numpy.asarray(result.toPandas().result.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32)), + "prediction_result": numpy.asarray( + result.toPandas() + .result.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ), } - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlCountVectorizerModel_Binary") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlCountVectorizerModel_Binary", + ) onnx_model_path = paths[-1] - - output_names = ['result'] + + output_names = ["result"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) actual_output = dict(zip(output_names, output)) - + assert output_shapes[0] == [None, 4] - compare_results(expected["prediction_result"], actual_output["result"], decimal=5) + compare_results( + expected["prediction_result"], actual_output["result"], decimal=5 + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/sparkml/test_dct.py b/tests/sparkml/test_dct.py index 18e76647..f290645c 100644 --- a/tests/sparkml/test_dct.py +++ b/tests/sparkml/test_dct.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,28 +22,38 @@ class TestSparkmlDCT(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_dct(self): - data = self.spark.createDataFrame( - [(Vectors.dense([5.0, 8.0, 6.0]),)], - ["vec"]) + data = self.spark.createDataFrame([(Vectors.dense([5.0, 8.0, 6.0]),)], ["vec"]) model = DCT(inverse=False, inputCol="vec", outputCol="resultVec") # the input name should match that of what inputCol feature_count = data.first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml DCT', [('vec', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml DCT", + [("vec", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().resultVec.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().vec.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlDCT") + expected = ( + predicted.toPandas() + .resultVec.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .vec.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlDCT" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['resultVec'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["resultVec"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_decision_tree_classifier.py b/tests/sparkml/test_decision_tree_classifier.py index c435b99e..5929d15c 100644 --- a/tests/sparkml/test_decision_tree_classifier.py +++ b/tests/sparkml/test_decision_tree_classifier.py @@ -14,7 +14,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType, FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, compare_results, run_onnx_model +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + compare_results, + run_onnx_model, +) from tests.sparkml import SparkMlTestCase from pyspark.ml.feature import StringIndexer, VectorIndexer @@ -23,135 +27,216 @@ class TestSparkmDecisionTreeClassifier(SparkMlTestCase): - # @unittest.skipIf(True, reason="Mismatched input dimensions.") - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_tree_pipeline(self): import os - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) # # truncate the features # feature_count = 5 - self.spark.udf.register("truncateFeatures", - lambda x: SparseVector(feature_count, range(0, feature_count), x.toArray()[125:130]), - VectorUDT()) - data = original_data.selectExpr("cast(label as string) as label", "truncateFeatures(features) as features") - label_indexer = StringIndexer(inputCol="label", outputCol="indexedLabel", handleInvalid='error') - feature_indexer = VectorIndexer(inputCol="features", outputCol="indexedFeatures", - maxCategories=10, handleInvalid='error') + self.spark.udf.register( + "truncateFeatures", + lambda x: SparseVector( + feature_count, range(0, feature_count), x.toArray()[125:130] + ), + VectorUDT(), + ) + data = original_data.selectExpr( + "cast(label as string) as label", "truncateFeatures(features) as features" + ) + label_indexer = StringIndexer( + inputCol="label", outputCol="indexedLabel", handleInvalid="error" + ) + feature_indexer = VectorIndexer( + inputCol="features", + outputCol="indexedFeatures", + maxCategories=10, + handleInvalid="error", + ) - dt = DecisionTreeClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures") + dt = DecisionTreeClassifier( + labelCol="indexedLabel", featuresCol="indexedFeatures" + ) pipeline = Pipeline(stages=[label_indexer, feature_indexer, dt]) model = pipeline.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree Pipeline', [ - ('label', StringTensorType([None, 1])), - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree Pipeline", + [ + ("label", StringTensorType([None, 1])), + ("features", FloatTensorType([None, feature_count])), + ], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data.limit(1)) data_np = { - 'label': data.limit(1).toPandas().label.values.reshape((-1, 1)), - 'features': data.limit(1).toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + "label": data.limit(1).toPandas().label.values.reshape((-1, 1)), + "features": data.limit(1) + .toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), } expected = [ predicted.toPandas().indexedLabel.values.astype(numpy.int64), predicted.toPandas().prediction.values.astype(numpy.int64), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreePipeline") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlDecisionTreePipeline" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['indexedLabel', 'prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["indexedLabel", "prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) # @unittest.skipIf(True, reason="Mismatched input dimensions.") - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_tree_one_class_classification(self): - features = [[0., 1.], [1., 1.], [2., 0.]] + features = [[0.0, 1.0], [1.0, 1.0], [2.0, 0.0]] features = numpy.array(features, dtype=numpy.float32) labels = [1, 1, 0] dd = [(labels[i], Vectors.dense(features[i])) for i in range(len(labels))] - data = self.spark.createDataFrame(self.spark.sparkContext.parallelize(dd), schema=["label", "features"]) + data = self.spark.createDataFrame( + self.spark.sparkContext.parallelize(dd), schema=["label", "features"] + ) dt = DecisionTreeClassifier(labelCol="label", featuresCol="features") model = dt.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree One Class', [ - ('features', FloatTensorType([None, 2])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree One Class", + [("features", FloatTensorType([None, 2]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) predicted = model.transform(data) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreeBinaryClass") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlDecisionTreeBinaryClass", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_tree_binary_classification(self): features = [[0, 1], [1, 1], [2, 0]] features = numpy.array(features, dtype=numpy.float32) labels = [0, 1, 0] dd = [(labels[i], Vectors.dense(features[i])) for i in range(len(labels))] - data = self.spark.createDataFrame(self.spark.sparkContext.parallelize(dd), schema=["label", "features"]) + data = self.spark.createDataFrame( + self.spark.sparkContext.parallelize(dd), schema=["label", "features"] + ) dt = DecisionTreeClassifier(labelCol="label", featuresCol="features") model = dt.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree Binary Class', [ - ('features', FloatTensorType([None, 2])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree Binary Class", + [("features", FloatTensorType([None, 2]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) predicted = model.transform(data) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreeBinaryClass") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlDecisionTreeBinaryClass", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_tree_multiple_classification(self): features = [[0, 1], [1, 1], [2, 0], [0.5, 0.5], [1.1, 1.1], [2.1, 0.1]] features = numpy.array(features, dtype=numpy.float32) labels = [0, 1, 2, 1, 1, 2] dd = [(labels[i], Vectors.dense(features[i])) for i in range(len(labels))] - data = self.spark.createDataFrame(self.spark.sparkContext.parallelize(dd), schema=["label", "features"]) + data = self.spark.createDataFrame( + self.spark.sparkContext.parallelize(dd), schema=["label", "features"] + ) dt = DecisionTreeClassifier(labelCol="label", featuresCol="features") model = dt.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree Multi Class', [ - ('features', FloatTensorType([None, 2])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree Multi Class", + [("features", FloatTensorType([None, 2]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) predicted = model.transform(data) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreeMultiClass") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlDecisionTreeMultiClass", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_decision_tree_classifier_category.py b/tests/sparkml/test_decision_tree_classifier_category.py index 95acde89..4d581188 100644 --- a/tests/sparkml/test_decision_tree_classifier_category.py +++ b/tests/sparkml/test_decision_tree_classifier_category.py @@ -1,51 +1,46 @@ # SPDX-License-Identifier: Apache-2.0 import sys -import os -import inspect import unittest import packaging.version as pv import onnx import pandas import numpy -from sklearn.datasets import dump_svmlight_file + try: from sklearn.utils._testing import ignore_warnings except ImportError: from sklearn.utils.testing import ignore_warnings -from pyspark.ml import Pipeline from pyspark.ml.classification import DecisionTreeClassifier -from pyspark.ml.linalg import VectorUDT, SparseVector, Vectors from pyspark.ml.feature import VectorAssembler -from pyspark.sql.functions import col from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml -from onnxmltools.convert.common.data_types import StringTensorType, FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, compare_results, run_onnx_model +from onnxmltools.convert.common.data_types import FloatTensorType +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + compare_results, + run_onnx_model, +) from tests.sparkml import SparkMlTestCase -from pyspark.ml.feature import StringIndexer, VectorIndexer TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) class TestSparkmDecisionTreeClassifierBig(SparkMlTestCase): - # @unittest.skipIf(True, reason="Mismatched input dimensions.") @ignore_warnings(category=(ResourceWarning, DeprecationWarning)) - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_tree_pipeline_category(self): ok = {1, 6, 10, 56, 78, 34} ok2 = {1, 7, 8, 23, 35, 45} - ok3 = {6, 9, 3, 29, 31, 40} def f_label(x): - if x[1] in ok or x[0] in ok: return 1 if x[1] in ok or x[0] in ok2: @@ -59,38 +54,62 @@ def f_label(x): features_names = [f"c{i}" for i in range(df.shape[1])] df.columns = features_names cat = set(s for s in df["c0"]) - df["c0"] = pandas.Categorical(list(map(lambda s: int(s), df['c0'])), categories=cat, ordered=False) + df["c0"] = pandas.Categorical( + list(map(lambda s: int(s), df["c0"])), categories=cat, ordered=False + ) cat = set(s for s in df["c1"]) - df["c1"] = pandas.Categorical(list(map(lambda s: int(s), df['c1'])), categories=cat, ordered=False) + df["c1"] = pandas.Categorical( + list(map(lambda s: int(s), df["c1"])), categories=cat, ordered=False + ) df["label"] = labels - sparkDF = self.spark.createDataFrame(df) + sparkDF = self.spark.createDataFrame(df) - data = sparkDF # self.spark.read.csv(input_path, header=True, inferSchema=True) - va = VectorAssembler(inputCols=features_names, outputCol='features') + data = sparkDF # self.spark.read.csv(input_path, header=True, inferSchema=True) + va = VectorAssembler(inputCols=features_names, outputCol="features") va_df = va.transform(data) - va_df = va_df.select(['features', 'label']) + va_df = va_df.select(["features", "label"]) - dt = DecisionTreeClassifier(labelCol="label", featuresCol='features', maxDepth=3, maxBins=50) + dt = DecisionTreeClassifier( + labelCol="label", featuresCol="features", maxDepth=3, maxBins=50 + ) model = dt.fit(va_df) # print(model.toDebugString) - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree Binary Class', [ - ('features', FloatTensorType([None, n_features])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) - data_np = va_df.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree Binary Class", + [("features", FloatTensorType([None, n_features]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) + data_np = ( + va_df.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) predicted = model.transform(va_df) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreeBinaryClassCategory") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlDecisionTreeBinaryClassCategory", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) if __name__ == "__main__": import logging - logging.basicConfig(encoding='utf-8', level=logging.INFO) + + logging.basicConfig(encoding="utf-8", level=logging.INFO) unittest.main() diff --git a/tests/sparkml/test_decision_tree_regressor.py b/tests/sparkml/test_decision_tree_regressor.py index bfb2654d..c467ad3b 100644 --- a/tests/sparkml/test_decision_tree_regressor.py +++ b/tests/sparkml/test_decision_tree_regressor.py @@ -14,7 +14,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase from pyspark.ml.feature import VectorIndexer @@ -23,73 +27,107 @@ class TestSparkmDecisionTreeRegressor(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_decision_tree_regressor_pipeline(self): import os - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) feature_count = 5 - self.spark.udf.register("truncateFeatures", - lambda x: SparseVector(feature_count, range(0,feature_count), x.toArray()[125:130]), - VectorUDT()) - data = original_data.selectExpr("label", "truncateFeatures(features) as features") + self.spark.udf.register( + "truncateFeatures", + lambda x: SparseVector( + feature_count, range(0, feature_count), x.toArray()[125:130] + ), + VectorUDT(), + ) + data = original_data.selectExpr( + "label", "truncateFeatures(features) as features" + ) - featureIndexer = \ - VectorIndexer(inputCol="features", outputCol="indexedFeatures", maxCategories=4, handleInvalid='error') + featureIndexer = VectorIndexer( + inputCol="features", + outputCol="indexedFeatures", + maxCategories=4, + handleInvalid="error", + ) (trainingData, testData) = data.randomSplit([0.7, 0.3]) dt = DecisionTreeRegressor(featuresCol="indexedFeatures") pipeline = Pipeline(stages=[featureIndexer, dt]) model = pipeline.fit(trainingData) - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree Regressor Pipeline', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree Regressor Pipeline", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(testData) - data_np = testData.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - expected = [ - predicted.toPandas().prediction.values.astype(numpy.float32) - ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreeRegressorPipeline") + data_np = ( + testData.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + expected = [predicted.toPandas().prediction.values.astype(numpy.float32)] + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlDecisionTreeRegressorPipeline", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_decision_tree_regressor(self): features = [[0, 1], [1, 1], [2, 0]] features = numpy.array(features, dtype=numpy.float32) labels = [100, -10, 50] dd = [(labels[i], Vectors.dense(features[i])) for i in range(len(labels))] - data = self.spark.createDataFrame(self.spark.sparkContext.parallelize(dd), schema=["label", "features"]) + data = self.spark.createDataFrame( + self.spark.sparkContext.parallelize(dd), schema=["label", "features"] + ) dt = DecisionTreeRegressor(labelCol="label", featuresCol="features") model = dt.fit(data) - feature_count = data.select('features').first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml Decision Tree Regressor', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + feature_count = data.select("features").first()[0].size + model_onnx = convert_sparkml( + model, + "Sparkml Decision Tree Regressor", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) predicted = model.transform(data) - expected = [ - predicted.toPandas().prediction.values.astype(numpy.float32) - ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlDecisionTreeRegressor") + expected = [predicted.toPandas().prediction.values.astype(numpy.float32)] + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlDecisionTreeRegressor", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_decision_tree_rules.py b/tests/sparkml/test_decision_tree_rules.py index 89f7282b..08056802 100644 --- a/tests/sparkml/test_decision_tree_rules.py +++ b/tests/sparkml/test_decision_tree_rules.py @@ -5,131 +5,150 @@ class TestSparkmDecisionTreeClassifier(unittest.TestCase): - def test_rule_number(self): attrs = { - 'class_ids': [0, 0, 0, 0, 0, 0], - 'class_nodeids': [2, 4, 5, 8, 9, 10], - 'class_treeids': [0, 0, 0, 0, 0, 0], - 'class_weights': [0.8462194428652643, - 0.2781875658587987, - 0.5437174290677474, - 0.6656197654941374, - 0.4343004513217279, - 0.2975769813225644], - 'classlabels_int64s': [0, 1], - 'nodes_falsenodeids': [6, 3, 0, 5, 0, 0, 10, 9, 0, 0, 0], - 'nodes_featureids': [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], - 'nodes_hitrates': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - 'nodes_missing_value_tracks_true': [False, - False, - False, - False, - False, - False, - False, - False, - False, - False, - False], - 'nodes_modes': ['BRANCH_LEQ', - 'BRANCH_LEQ', - 'LEAF', - 'BRANCH_LEQ', - 'LEAF', - 'LEAF', - 'BRANCH_LEQ', - 'BRANCH_LEQ', - 'LEAF', - 'LEAF', - 'LEAF'], - 'nodes_nodeids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - 'nodes_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'nodes_truenodeids': [1, 2, 0, 4, 0, 0, 7, 8, 0, 0, 0], - 'nodes_values': [21.65000057220459, - 100.95000076293945, - 0.0, - -22.84999942779541, - 0.0, - 0.0, - 98.0999984741211, - 37.14999961853027, - 0.0, - 0.0, - 0.0], - 'post_transform': 'NONE'} + "class_ids": [0, 0, 0, 0, 0, 0], + "class_nodeids": [2, 4, 5, 8, 9, 10], + "class_treeids": [0, 0, 0, 0, 0, 0], + "class_weights": [ + 0.8462194428652643, + 0.2781875658587987, + 0.5437174290677474, + 0.6656197654941374, + 0.4343004513217279, + 0.2975769813225644, + ], + "classlabels_int64s": [0, 1], + "nodes_falsenodeids": [6, 3, 0, 5, 0, 0, 10, 9, 0, 0, 0], + "nodes_featureids": [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "nodes_hitrates": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + "nodes_missing_value_tracks_true": [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + ], + "nodes_modes": [ + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "LEAF", + ], + "nodes_nodeids": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + "nodes_treeids": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "nodes_truenodeids": [1, 2, 0, 4, 0, 0, 7, 8, 0, 0, 0], + "nodes_values": [ + 21.65000057220459, + 100.95000076293945, + 0.0, + -22.84999942779541, + 0.0, + 0.0, + 98.0999984741211, + 37.14999961853027, + 0.0, + 0.0, + 0.0, + ], + "post_transform": "NONE", + } root, _ = Node.create(attrs) root.unfold_rule_or() new_attrs = root.to_attrs( - post_transform=attrs['post_transform'], - classlabels_int64s=attrs["classlabels_int64s"]) - assert len(attrs['nodes_nodeids']) <= len(new_attrs['nodes_nodeids']) + post_transform=attrs["post_transform"], + classlabels_int64s=attrs["classlabels_int64s"], + ) + assert len(attrs["nodes_nodeids"]) <= len(new_attrs["nodes_nodeids"]) for i in range(len(new_attrs["nodes_truenodeids"])): - if new_attrs["nodes_modes"][i] == 'LEAF': + if new_attrs["nodes_modes"][i] == "LEAF": continue assert new_attrs["nodes_truenodeids"][i] > i for i in range(len(new_attrs["nodes_falsenodeids"])): - if new_attrs["nodes_modes"][i] == 'LEAF': + if new_attrs["nodes_modes"][i] == "LEAF": continue assert new_attrs["nodes_falsenodeids"][i] > i attrs = { - 'class_ids': [0, 0, 0, 0, 0, 0], - 'class_nodeids': [102, 104, 105, 108, 109, 1010], - 'class_treeids': [0, 0, 0, 0, 0, 0], - 'class_weights': [0.8462194428652643, - 0.2781875658587987, - 0.5437174290677474, - 0.6656197654941374, - 0.4343004513217279, - 0.2975769813225644], - 'classlabels_int64s': [0, 1], - 'nodes_falsenodeids': [106, 103, 0, 105, 0, 0, 1010, 109, 0, 0, 0], - 'nodes_featureids': [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], - 'nodes_hitrates': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - 'nodes_missing_value_tracks_true': [False, - False, - False, - False, - False, - False, - False, - False, - False, - False, - False], - 'nodes_modes': ['BRANCH_LEQ', - 'BRANCH_LEQ', - 'LEAF', - 'BRANCH_LEQ', - 'LEAF', - 'LEAF', - 'BRANCH_LEQ', - 'BRANCH_LEQ', - 'LEAF', - 'LEAF', - 'LEAF'], - 'nodes_nodeids': [100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 1010], - 'nodes_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'nodes_truenodeids': [101, 102, 0, 104, 0, 0, 107, 108, 0, 0, 0], - 'nodes_values': [21.65000057220459, - 100.95000076293945, - 0.0, - -22.84999942779541, - 0.0, - 0.0, - 98.0999984741211, - 37.14999961853027, - 0.0, - 0.0, - 0.0], - 'post_transform': 'NONE'} + "class_ids": [0, 0, 0, 0, 0, 0], + "class_nodeids": [102, 104, 105, 108, 109, 1010], + "class_treeids": [0, 0, 0, 0, 0, 0], + "class_weights": [ + 0.8462194428652643, + 0.2781875658587987, + 0.5437174290677474, + 0.6656197654941374, + 0.4343004513217279, + 0.2975769813225644, + ], + "classlabels_int64s": [0, 1], + "nodes_falsenodeids": [106, 103, 0, 105, 0, 0, 1010, 109, 0, 0, 0], + "nodes_featureids": [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "nodes_hitrates": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + "nodes_missing_value_tracks_true": [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + ], + "nodes_modes": [ + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "LEAF", + ], + "nodes_nodeids": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 1010], + "nodes_treeids": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "nodes_truenodeids": [101, 102, 0, 104, 0, 0, 107, 108, 0, 0, 0], + "nodes_values": [ + 21.65000057220459, + 100.95000076293945, + 0.0, + -22.84999942779541, + 0.0, + 0.0, + 98.0999984741211, + 37.14999961853027, + 0.0, + 0.0, + 0.0, + ], + "post_transform": "NONE", + } root, _ = Node.create(attrs) root.unfold_rule_or() new_attrs2 = root.to_attrs( - post_transform=attrs['post_transform'], - classlabels_int64s=attrs["classlabels_int64s"]) - assert len(attrs['nodes_nodeids']) <= len(new_attrs2['nodes_nodeids']) + post_transform=attrs["post_transform"], + classlabels_int64s=attrs["classlabels_int64s"], + ) + assert len(attrs["nodes_nodeids"]) <= len(new_attrs2["nodes_nodeids"]) for k in new_attrs: if new_attrs[k] != new_attrs2[k]: raise ValueError(f"Key {k!r}, {new_attrs[k]} != {new_attrs2[k]}") @@ -137,62 +156,72 @@ def test_rule_number(self): def test_rule_in_set(self): attrs = { - 'class_ids': [0, 0, 0, 0, 0, 0], - 'class_nodeids': [2, 4, 5, 8, 9, 10], - 'class_treeids': [0, 0, 0, 0, 0, 0], - 'class_weights': [0.8462194428652643, - 0.2781875658587987, - 0.5437174290677474, - 0.6656197654941374, - 0.4343004513217279, - 0.2975769813225644], - 'classlabels_int64s': [0, 1], - 'nodes_falsenodeids': [6, 3, 0, 5, 0, 0, 10, 9, 0, 0, 0], - 'nodes_featureids': [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], - 'nodes_hitrates': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - 'nodes_missing_value_tracks_true': [False, - False, - False, - False, - False, - False, - False, - False, - False, - False, - False], - 'nodes_modes': ['||', - 'BRANCH_LEQ', - 'LEAF', - 'BRANCH_LEQ', - 'LEAF', - 'LEAF', - 'BRANCH_LEQ', - 'BRANCH_LEQ', - 'LEAF', - 'LEAF', - 'LEAF'], - 'nodes_nodeids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], - 'nodes_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'nodes_truenodeids': [1, 2, 0, 4, 0, 0, 7, 8, 0, 0, 0], - 'nodes_values': [[21.56, 22.56], - 100.95000076293945, - 0.0, - -22.84999942779541, - 0.0, - 0.0, - 98.0999984741211, - 37.14999961853027, - 0.0, - 0.0, - 0.0], - 'post_transform': 'NONE'} + "class_ids": [0, 0, 0, 0, 0, 0], + "class_nodeids": [2, 4, 5, 8, 9, 10], + "class_treeids": [0, 0, 0, 0, 0, 0], + "class_weights": [ + 0.8462194428652643, + 0.2781875658587987, + 0.5437174290677474, + 0.6656197654941374, + 0.4343004513217279, + 0.2975769813225644, + ], + "classlabels_int64s": [0, 1], + "nodes_falsenodeids": [6, 3, 0, 5, 0, 0, 10, 9, 0, 0, 0], + "nodes_featureids": [0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0], + "nodes_hitrates": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], + "nodes_missing_value_tracks_true": [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + ], + "nodes_modes": [ + "||", + "BRANCH_LEQ", + "LEAF", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "LEAF", + ], + "nodes_nodeids": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10], + "nodes_treeids": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "nodes_truenodeids": [1, 2, 0, 4, 0, 0, 7, 8, 0, 0, 0], + "nodes_values": [ + [21.56, 22.56], + 100.95000076293945, + 0.0, + -22.84999942779541, + 0.0, + 0.0, + 98.0999984741211, + 37.14999961853027, + 0.0, + 0.0, + 0.0, + ], + "post_transform": "NONE", + } root, _ = Node.create(attrs) root.unfold_rule_or() new_attrs = root.to_attrs( - post_transform=attrs['post_transform'], - classlabels_int64s=attrs["classlabels_int64s"]) - assert len(attrs['nodes_nodeids']) < len(new_attrs['nodes_nodeids']) + post_transform=attrs["post_transform"], + classlabels_int64s=attrs["classlabels_int64s"], + ) + assert len(attrs["nodes_nodeids"]) < len(new_attrs["nodes_nodeids"]) assert "BRANCH_EQ" in new_attrs["nodes_modes"] assert "||" not in new_attrs["nodes_modes"] diff --git a/tests/sparkml/test_element_wise_product.py b/tests/sparkml/test_element_wise_product.py index 79cd3817..9e12ef7d 100644 --- a/tests/sparkml/test_element_wise_product.py +++ b/tests/sparkml/test_element_wise_product.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,28 +22,42 @@ class TestSparkmlElementwiseProduct(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_element_wise_product(self): - data = self.spark.createDataFrame([(Vectors.dense([2.0, 1.0, 3.0]),)], ["features"]) - model = ElementwiseProduct(scalingVec=Vectors.dense([1.0, 2.0, 3.0]), - inputCol="features", outputCol="eprod") + data = self.spark.createDataFrame( + [(Vectors.dense([2.0, 1.0, 3.0]),)], ["features"] + ) + model = ElementwiseProduct( + scalingVec=Vectors.dense([1.0, 2.0, 3.0]), + inputCol="features", + outputCol="eprod", + ) feature_count = data.first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml ElementwiseProduct', - [('features', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml ElementwiseProduct", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) expected = [ - predicted.toPandas().eprod.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - ] - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlElementwiseProduct") + predicted.toPandas() + .eprod.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ] + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlElementwiseProduct" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['eprod'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["eprod"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_gbt_classifier.py b/tests/sparkml/test_gbt_classifier.py index 5a619b9e..3388d5bb 100644 --- a/tests/sparkml/test_gbt_classifier.py +++ b/tests/sparkml/test_gbt_classifier.py @@ -12,7 +12,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase from pyspark.ml.feature import StringIndexer @@ -21,38 +25,50 @@ class TestSparkmTreeEnsembleClassifier(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_gbt_classifier(self): - raw_data = self.spark.createDataFrame([ - (1.0, Vectors.dense(1.0)), - (0.0, Vectors.sparse(1, [], [])) - ], ["label", "features"]) + raw_data = self.spark.createDataFrame( + [(1.0, Vectors.dense(1.0)), (0.0, Vectors.sparse(1, [], []))], + ["label", "features"], + ) string_indexer = StringIndexer(inputCol="label", outputCol="indexed") si_model = string_indexer.fit(raw_data) data = si_model.transform(raw_data) gbt = GBTClassifier(maxIter=5, maxDepth=2, labelCol="indexed", seed=42) model = gbt.fit(data) feature_count = data.first()[1].size - model_onnx = convert_sparkml(model, 'Sparkml GBT Classifier', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml GBT Classifier", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlGBTClassifier") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlGBTClassifier" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_gbt_regressor.py b/tests/sparkml/test_gbt_regressor.py index c9c9a0a5..5f6b3665 100644 --- a/tests/sparkml/test_gbt_regressor.py +++ b/tests/sparkml/test_gbt_regressor.py @@ -10,7 +10,11 @@ from onnxmltools.convert.common.data_types import FloatTensorType from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,33 +22,39 @@ class TestSparkmTreeEnsembleClassifier(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_gbt_regressor(self): - data = self.spark.createDataFrame([ - (1.0, Vectors.dense(1.0)), - (0.0, Vectors.sparse(1, [], [])) - ], ["label", "features"]) + data = self.spark.createDataFrame( + [(1.0, Vectors.dense(1.0)), (0.0, Vectors.sparse(1, [], []))], + ["label", "features"], + ) gbt = GBTRegressor(maxIter=5, maxDepth=2, seed=42) model = gbt.fit(data) feature_count = data.first()[1].size - model_onnx = convert_sparkml(model, 'Sparkml GBTRegressor', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml GBTRegressor", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlGBTRegressor") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlGBTRegressor" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_imputer.py b/tests/sparkml/test_imputer.py index 090b2d54..8de27d6f 100644 --- a/tests/sparkml/test_imputer.py +++ b/tests/sparkml/test_imputer.py @@ -8,88 +8,109 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) -## For some reason during the spark bring up and shutdown something happens causing Imputer +## For some reason during the spark bring up and shutdown +## something happens causing Imputer ## tests to fail. For that you need to run each test here individually ## for now these will be commented out so as not to break the build -## AttributeError: 'NoneType' object has no attribute 'setCallSite' on model.surrogateDF -## Therefore we leave these tests out for now until a newere version of pyspark is availabe that address this issue +## AttributeError: 'NoneType' object has no attribute +# 'setCallSite' on model.surrogateDF +## Therefore we leave these tests out for now until a newere +## version of pyspark is availabe that address this issue class TestSparkmlImputer(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_imputer_single(self): self._imputer_test_single() - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_imputer_multi(self): self._imputer_test_multi() def _imputer_test_multi(self): - data = self.spark.createDataFrame([ - (1.0, float("nan")), - (2.0, float("nan")), - (float("nan"), 3.0), - (4.0, 4.0), - (5.0, 5.0) - ], ["a", "b"]) + data = self.spark.createDataFrame( + [ + (1.0, float("nan")), + (2.0, float("nan")), + (float("nan"), 3.0), + (4.0, 4.0), + (5.0, 5.0), + ], + ["a", "b"], + ) imputer = Imputer(inputCols=["a", "b"], outputCols=["out_a", "out_b"]) model = imputer.fit(data) - + # the input name should match the inputCols above - model_onnx = convert_sparkml(model, 'Sparkml Imputer Multi Input', [ - ('a', FloatTensorType([None, 1])), - ('b', FloatTensorType([None, 1]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Imputer Multi Input", + [("a", FloatTensorType([None, 1])), ("b", FloatTensorType([None, 1]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) - + # run the model predicted = model.transform(data) - + expected = { "out_a": predicted.select("out_a").toPandas().values.astype(numpy.int64), "out_b": predicted.select("out_b").toPandas().values.astype(numpy.int64), } data_np = data.toPandas().values.astype(numpy.float32) - data_np = {'a': data_np[:, :1], 'b': data_np[:, 1:]} - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlImputerMulti") + data_np = {"a": data_np[:, :1], "b": data_np[:, 1:]} + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlImputerMulti" + ) onnx_model_path = paths[-1] - output_names = ['out_a', 'out_b'] + output_names = ["out_a", "out_b"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) actual_output = dict(zip(output_names, output)) compare_results(expected, actual_output, decimal=5) - + def _imputer_test_single(self): - data = self.spark.createDataFrame([ - (1.0, float("nan")), - (2.0, float("nan")), - (float("nan"), 3.0), - (4.0, 4.0), - (5.0, 5.0) - ], ["a", "b"]) + data = self.spark.createDataFrame( + [ + (1.0, float("nan")), + (2.0, float("nan")), + (float("nan"), 3.0), + (4.0, 4.0), + (5.0, 5.0), + ], + ["a", "b"], + ) imputer = Imputer(inputCols=["a"], outputCols=["out_a"]) model = imputer.fit(data) - + # the input name should match the inputCols above - model_onnx = convert_sparkml(model, 'Sparkml Imputer', [ - ('a', FloatTensorType([None, 1]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Imputer", + [("a", FloatTensorType([None, 1]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) - + # run the model predicted = model.transform(data) expected = predicted.select("out_a").toPandas().values.astype(numpy.float32) data_np = data.toPandas().a.values.astype(numpy.float32) data_np = data_np.reshape((-1, 1)) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlImputerSingle") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlImputerSingle" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['out_a'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["out_a"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_index_to_string.py b/tests/sparkml/test_index_to_string.py index 294e0cf6..e1c313b3 100644 --- a/tests/sparkml/test_index_to_string.py +++ b/tests/sparkml/test_index_to_string.py @@ -10,7 +10,11 @@ from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import Int64TensorType from onnxmltools.convert.sparkml.utils import SparkMlConversionError -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,49 +22,61 @@ class TestSparkmlIndexToString(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") @pytest.mark.xfail(raises=SparkMlConversionError) def test_index_to_string_throws(self): original_data = self.spark.createDataFrame( [(0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c")], - ["id", "category"]) + ["id", "category"], + ) string_indexer = StringIndexer(inputCol="category", outputCol="categoryIndex") string_indexer_model = string_indexer.fit(original_data) - data = string_indexer_model.transform(original_data) + string_indexer_model.transform(original_data) model = IndexToString(inputCol="categoryIndex", outputCol="originalCategory") # the input name should match that of what IndexToString.inputCol - model_onnx = None with pytest.raises(SparkMlConversionError): - model_onnx = convert_sparkml(model, 'Sparkml IndexToString', [('categoryIndex', Int64TensorType([None, 1]))], - target_opset=TARGET_OPSET) + convert_sparkml( + model, + "Sparkml IndexToString", + [("categoryIndex", Int64TensorType([None, 1]))], + target_opset=TARGET_OPSET, + ) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_index_to_string(self): original_data = self.spark.createDataFrame( [(0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c")], - ["id", "category"]) + ["id", "category"], + ) string_indexer = StringIndexer(inputCol="category", outputCol="categoryIndex") string_indexer_model = string_indexer.fit(original_data) data = string_indexer_model.transform(original_data) - model = IndexToString(inputCol="categoryIndex", outputCol="originalCategory", - labels=['A', 'B', 'C']) + model = IndexToString( + inputCol="categoryIndex", + outputCol="originalCategory", + labels=["A", "B", "C"], + ) # the input name should match that of what IndexToString.inputCol - model_onnx = convert_sparkml(model, 'Sparkml IndexToString', [('categoryIndex', Int64TensorType([None, 1]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml IndexToString", + [("categoryIndex", Int64TensorType([None, 1]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) expected = predicted.select("originalCategory").toPandas().values - data_np = data.select('categoryIndex').toPandas().values.astype(numpy.int64) - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlIndexToString") + data_np = data.select("categoryIndex").toPandas().values.astype(numpy.int64) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlIndexToString" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['originalCategory'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["originalCategory"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_k_means.py b/tests/sparkml/test_k_means.py index c03ce40d..2b8bf9a2 100644 --- a/tests/sparkml/test_k_means.py +++ b/tests/sparkml/test_k_means.py @@ -11,7 +11,11 @@ from pyspark.ml import Pipeline from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -19,61 +23,115 @@ class TestSparkmlKMeansModel(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_k_means_euclidean(self): """ - Testing ONNX conversion for Spark KMeansModel when distanceMeasure is set to "euclidean". + Testing ONNX conversion for Spark KMeansModel + when distanceMeasure is set to "euclidean". """ - kmeans_euclidean = KMeans(k=3, distanceMeasure="euclidean", featuresCol="features_euclidean", predictionCol="prediction_euclidean") - kmeans_cosine = KMeans(k=3, distanceMeasure="cosine", featuresCol="features_cosine", predictionCol="prediction_cosine") - - data = self.spark.createDataFrame([ - (0, Vectors.dense([1.0, 3.1, -1.0]),Vectors.dense([1.0, 1.0, 1.0]),), - (1, Vectors.dense([1.1, 3.0, -1.1]),Vectors.dense([2.0, 2.0, 2.0]),), - (2, Vectors.dense([-3.0, 5.1, 9.0]),Vectors.dense([-1.0, 3.0, -5.0]),), - (3, Vectors.dense([-2.9, 4.9, 8.9]),Vectors.dense([-2.0, 6.0, -10.0]),), - (4, Vectors.dense([5.0, -3.5, 2.0]),Vectors.dense([1.0, -2.0, 4.0]),), - (5, Vectors.dense([5.1, -3.3, 2.1]),Vectors.dense([2.0, -4.0, 8.0]),), - ], ["id", "features_euclidean", "features_cosine"]) + kmeans_euclidean = KMeans( + k=3, + distanceMeasure="euclidean", + featuresCol="features_euclidean", + predictionCol="prediction_euclidean", + ) + kmeans_cosine = KMeans( + k=3, + distanceMeasure="cosine", + featuresCol="features_cosine", + predictionCol="prediction_cosine", + ) + + data = self.spark.createDataFrame( + [ + ( + 0, + Vectors.dense([1.0, 3.1, -1.0]), + Vectors.dense([1.0, 1.0, 1.0]), + ), + ( + 1, + Vectors.dense([1.1, 3.0, -1.1]), + Vectors.dense([2.0, 2.0, 2.0]), + ), + ( + 2, + Vectors.dense([-3.0, 5.1, 9.0]), + Vectors.dense([-1.0, 3.0, -5.0]), + ), + ( + 3, + Vectors.dense([-2.9, 4.9, 8.9]), + Vectors.dense([-2.0, 6.0, -10.0]), + ), + ( + 4, + Vectors.dense([5.0, -3.5, 2.0]), + Vectors.dense([1.0, -2.0, 4.0]), + ), + ( + 5, + Vectors.dense([5.1, -3.3, 2.1]), + Vectors.dense([2.0, -4.0, 8.0]), + ), + ], + ["id", "features_euclidean", "features_cosine"], + ) model = Pipeline(stages=[kmeans_euclidean, kmeans_cosine]).fit(data) model_onnx = convert_sparkml( - model, - 'Sparkml KMeansModel', - [('features_euclidean', FloatTensorType([None, 3])), ('features_cosine', FloatTensorType([None, 3]))], - target_opset=TARGET_OPSET - ) + model, + "Sparkml KMeansModel", + [ + ("features_euclidean", FloatTensorType([None, 3])), + ("features_cosine", FloatTensorType([None, 3])), + ], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(data).toPandas() - + data_pd = data.toPandas() data_np = { - "features_euclidean": data_pd.features_euclidean.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32), - "features_cosine": data_pd.features_cosine.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32), + "features_euclidean": data_pd.features_euclidean.apply( + lambda x: pandas.Series(x.toArray()) + ).values.astype(numpy.float32), + "features_cosine": data_pd.features_cosine.apply( + lambda x: pandas.Series(x.toArray()) + ).values.astype(numpy.float32), } expected = { - "prediction_euclidean": numpy.asarray(predicted.prediction_euclidean.values), + "prediction_euclidean": numpy.asarray( + predicted.prediction_euclidean.values + ), "prediction_cosine": numpy.asarray(predicted.prediction_cosine.values), } - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlKMeansModel") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlKMeansModel" + ) onnx_model_path = paths[-1] - - output_names = ['prediction_euclidean', 'prediction_cosine'] + + output_names = ["prediction_euclidean", "prediction_cosine"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) actual_output = dict(zip(output_names, output)) - + assert output_shapes[0] == [None] assert output_shapes[1] == [None] - compare_results(expected["prediction_euclidean"], actual_output["prediction_euclidean"], decimal=5) - compare_results(expected["prediction_cosine"], actual_output["prediction_cosine"], decimal=5) + compare_results( + expected["prediction_euclidean"], + actual_output["prediction_euclidean"], + decimal=5, + ) + compare_results( + expected["prediction_cosine"], actual_output["prediction_cosine"], decimal=5 + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/sparkml/test_linear_classifier.py b/tests/sparkml/test_linear_classifier.py index 7d9dcf7e..4d380381 100644 --- a/tests/sparkml/test_linear_classifier.py +++ b/tests/sparkml/test_linear_classifier.py @@ -12,7 +12,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -20,67 +24,100 @@ class TestSparkmlLogisticRegression(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_logistic_regression_binary_class(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) # # truncate the features # - self.spark.udf.register("truncateFeatures", lambda x: SparseVector(5, range(0,5), x.toArray()[125:130]), - VectorUDT()) - data = original_data.selectExpr("label", "truncateFeatures(features) as features") + self.spark.udf.register( + "truncateFeatures", + lambda x: SparseVector(5, range(0, 5), x.toArray()[125:130]), + VectorUDT(), + ) + data = original_data.selectExpr( + "label", "truncateFeatures(features) as features" + ) lr = LogisticRegression(maxIter=100, tol=0.0001) model = lr.fit(data) # the name of the input for Logistic Regression is 'features' C = model.numFeatures - model_onnx = convert_sparkml(model, 'sparkml logistic regression', [('features', FloatTensorType([None, C]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "sparkml logistic regression", + [("features", FloatTensorType([None, C]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] # known error in onnxruntime 0.3.0 case - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlLogisticRegression") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlLogisticRegression" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_linear_svc(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) # # truncate the features # - self.spark.udf.register("truncateFeatures", lambda x: SparseVector(5, range(0,5), x.toArray()[125:130]), - VectorUDT()) - data = original_data.selectExpr("label", "truncateFeatures(features) as features") + self.spark.udf.register( + "truncateFeatures", + lambda x: SparseVector(5, range(0, 5), x.toArray()[125:130]), + VectorUDT(), + ) + data = original_data.selectExpr( + "label", "truncateFeatures(features) as features" + ) lsvc = LinearSVC(maxIter=10, regParam=0.01) model = lsvc.fit(data) # the name of the input for Logistic Regression is 'features' C = model.numFeatures - model_onnx = convert_sparkml(model, 'Spark ML Linear SVC', [('features', FloatTensorType([None, C]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Spark ML Linear SVC", + [("features", FloatTensorType([None, C]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - expected = [ predicted.toPandas().prediction.values.astype(numpy.float32) ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlLinearSVC") + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + expected = [predicted.toPandas().prediction.values.astype(numpy.float32)] + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlLinearSVC" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_linear_regressor.py b/tests/sparkml/test_linear_regressor.py index 862f60f1..f7aa485a 100644 --- a/tests/sparkml/test_linear_regressor.py +++ b/tests/sparkml/test_linear_regressor.py @@ -12,7 +12,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -20,77 +24,118 @@ class TestSparkmlLinearRegression(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_linear_regression_basic(self): - data = self.spark.createDataFrame([ - (1.0, 2.0, Vectors.dense(1.0)), - (0.0, 2.0, Vectors.sparse(1, [], [])) - ], ["label", "weight", "features"]) - lr = LinearRegression(maxIter=5, regParam=0.0, solver="normal", weightCol="weight") + data = self.spark.createDataFrame( + [(1.0, 2.0, Vectors.dense(1.0)), (0.0, 2.0, Vectors.sparse(1, [], []))], + ["label", "weight", "features"], + ) + lr = LinearRegression( + maxIter=5, regParam=0.0, solver="normal", weightCol="weight" + ) model = lr.fit(data) # the name of the input is 'features' C = model.numFeatures - model_onnx = convert_sparkml(model, 'sparkml LinearRegressorBasic', [('features', FloatTensorType([None, C]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "sparkml LinearRegressorBasic", + [("features", FloatTensorType([None, C]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - expected = [ predicted.toPandas().prediction.values.astype(numpy.float32) ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlLinearRegressor_Basic") + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + expected = [predicted.toPandas().prediction.values.astype(numpy.float32)] + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlLinearRegressor_Basic", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_linear_regression(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "sample_linear_regression_data.txt") + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "sample_linear_regression_data.txt" + ) data = self.spark.read.format("libsvm").load(input_path) lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8) model = lr.fit(data) # the name of the input is 'features' C = model.numFeatures - model_onnx = convert_sparkml(model, 'sparkml LinearRegressor', [('features', FloatTensorType([None, C]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "sparkml LinearRegressor", + [("features", FloatTensorType([None, C]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - expected = [ predicted.toPandas().prediction.values.astype(numpy.float32) ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlLinearRegressor") + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + expected = [predicted.toPandas().prediction.values.astype(numpy.float32)] + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlLinearRegressor" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_generalized_linear_regression(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "sample_linear_regression_data.txt") + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "sample_linear_regression_data.txt" + ) data = self.spark.read.format("libsvm").load(input_path) lr = LinearRegression(maxIter=10, regParam=0.3, elasticNetParam=0.8) model = lr.fit(data) # the name of the input is 'features' C = model.numFeatures - model_onnx = convert_sparkml(model, 'sparkml GeneralizedLinearRegression', [('features', FloatTensorType([None, C]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "sparkml GeneralizedLinearRegression", + [("features", FloatTensorType([None, C]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - expected = [ predicted.toPandas().prediction.values.astype(numpy.float32) ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlGeneralizedLinearRegression") + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + expected = [predicted.toPandas().prediction.values.astype(numpy.float32)] + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlGeneralizedLinearRegression", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_min_hash_lsh.py b/tests/sparkml/test_min_hash_lsh.py index 0cf4ce93..024e3252 100644 --- a/tests/sparkml/test_min_hash_lsh.py +++ b/tests/sparkml/test_min_hash_lsh.py @@ -10,7 +10,11 @@ from onnxmltools.convert.common.data_types import FloatTensorType from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,37 +22,57 @@ class TestSparkmMinHashLSH(SparkMlTestCase): - @unittest.skipIf(True, reason="Discrepencies (Float -> Double?).") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_min_hash_lsh(self): - data = self.spark.createDataFrame([ - (0, Vectors.sparse(6, [0, 1, 2], [1.0, 1.0, 1.0]),), - (1, Vectors.sparse(6, [2, 3, 4], [1.0, 1.0, 1.0]),), - (2, Vectors.sparse(6, [0, 2, 4], [1.0, 1.0, 1.0]),) - ], ["id", "features"]) + data = self.spark.createDataFrame( + [ + ( + 0, + Vectors.sparse(6, [0, 1, 2], [1.0, 1.0, 1.0]), + ), + ( + 1, + Vectors.sparse(6, [2, 3, 4], [1.0, 1.0, 1.0]), + ), + ( + 2, + Vectors.sparse(6, [0, 2, 4], [1.0, 1.0, 1.0]), + ), + ], + ["id", "features"], + ) mh = MinHashLSH(inputCol="features", outputCol="hashes", numHashTables=5) model = mh.fit(data) feature_count = data.first()[1].size - model_onnx = convert_sparkml(model, 'Sparkml MinHashLSH', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml MinHashLSH", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data.limit(2)) - data_np = data.limit(2).toPandas().features.apply( - lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.limit(2) + .toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ - predicted.toPandas().hashes.apply( - lambda x: pandas.Series(x).map( - lambda y: y.values[0])).values.astype(numpy.float32)] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlMinHashLSH") + predicted.toPandas() + .hashes.apply(lambda x: pandas.Series(x).map(lambda y: y.values[0])) + .values.astype(numpy.float32) + ] + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlMinHashLSH" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['hashes'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["hashes"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_mlp_classifier.py b/tests/sparkml/test_mlp_classifier.py index a3b6e36c..301d1f3c 100644 --- a/tests/sparkml/test_mlp_classifier.py +++ b/tests/sparkml/test_mlp_classifier.py @@ -6,13 +6,20 @@ import os import numpy import pandas -from pyspark.ml.classification import MultilayerPerceptronClassifier, MultilayerPerceptronClassificationModel +from pyspark.ml.classification import ( + MultilayerPerceptronClassifier, + MultilayerPerceptronClassificationModel, +) from pyspark.ml.linalg import VectorUDT, SparseVector from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -22,19 +29,30 @@ class TestSparkmlMLPClassifier(SparkMlTestCase): @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_mlp_classifier_binary_class(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) # # truncate the features # self.spark.udf.register( - "truncateFeatures", lambda x: SparseVector(100, range(0, 100), x.toArray()[30:130]), VectorUDT() + "truncateFeatures", + lambda x: SparseVector(100, range(0, 100), x.toArray()[30:130]), + VectorUDT(), ) - data = original_data.selectExpr("label", "truncateFeatures(features) as features") + data = original_data.selectExpr( + "label", "truncateFeatures(features) as features" + ) - mlp = MultilayerPerceptronClassifier(maxIter=100, tol=0.0001, seed=137, layers=[100, 20, 5, 2],) + mlp = MultilayerPerceptronClassifier( + maxIter=100, + tol=0.0001, + seed=137, + layers=[100, 20, 5, 2], + ) model: MultilayerPerceptronClassificationModel = mlp.fit(data) # the name of the input for Logistic Regression is 'features' @@ -50,17 +68,28 @@ def test_model_mlp_classifier_binary_class(self): # run the model predicted = model.transform(data) - # predicted.select("prediction", "probability", "label").show(100, truncate=False) + # predicted.select("prediction", "probability", "label").show( + # 100, truncate=False) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32), + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlMLPClassifier") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlMLPClassifier" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(["prediction", "probability"], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_naive_bayes.py b/tests/sparkml/test_naive_bayes.py index d4cc8dd7..f2911ff8 100644 --- a/tests/sparkml/test_naive_bayes.py +++ b/tests/sparkml/test_naive_bayes.py @@ -11,7 +11,11 @@ from onnxmltools.convert.common.data_types import FloatTensorType from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -19,59 +23,92 @@ class TestSparkmlNaiveBayes(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_naive_bayes_bernoulli(self): - data = self.spark.createDataFrame([ - Row(label=0.0, weight=0.1, features=Vectors.dense([0.0, 0.0])), - Row(label=0.0, weight=0.5, features=Vectors.dense([0.0, 1.0])), - Row(label=1.0, weight=1.0, features=Vectors.dense([1.0, 0.0]))]) + data = self.spark.createDataFrame( + [ + Row(label=0.0, weight=0.1, features=Vectors.dense([0.0, 0.0])), + Row(label=0.0, weight=0.5, features=Vectors.dense([0.0, 1.0])), + Row(label=1.0, weight=1.0, features=Vectors.dense([1.0, 0.0])), + ] + ) nb = NaiveBayes(smoothing=1.0, modelType="bernoulli", weightCol="weight") model = nb.fit(data) - feature_count = data.select('features').first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml NaiveBayes Bernoulli', - [('features', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + feature_count = data.select("features").first()[0].size + model_onnx = convert_sparkml( + model, + "Sparkml NaiveBayes Bernoulli", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - ] - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlNaiveBayesBernoulli") + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), + ] + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlNaiveBayesBernoulli" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_naive_bayes_multinomial(self): - data = self.spark.createDataFrame([ - Row(label=0.0, weight=0.1, features=Vectors.dense([0.0, 0.0])), - Row(label=0.0, weight=0.5, features=Vectors.dense([0.0, 1.0])), - Row(label=1.0, weight=1.0, features=Vectors.dense([1.0, 0.0]))]) + data = self.spark.createDataFrame( + [ + Row(label=0.0, weight=0.1, features=Vectors.dense([0.0, 0.0])), + Row(label=0.0, weight=0.5, features=Vectors.dense([0.0, 1.0])), + Row(label=1.0, weight=1.0, features=Vectors.dense([1.0, 0.0])), + ] + ) nb = NaiveBayes(smoothing=1.0, modelType="multinomial", weightCol="weight") model = nb.fit(data) - feature_count = data.select('features').first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml NaiveBayes Multinomial', - [('features', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + feature_count = data.select("features").first()[0].size + model_onnx = convert_sparkml( + model, + "Sparkml NaiveBayes Multinomial", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - ] - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlNaiveBayesMultinomial") + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), + ] + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlNaiveBayesMultinomial", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_normalizer.py b/tests/sparkml/test_normalizer.py index 765aca6b..78a13dab 100644 --- a/tests/sparkml/test_normalizer.py +++ b/tests/sparkml/test_normalizer.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,49 +22,84 @@ class TestSparkmlNormalizer(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_normalizer_1(self): - data = self.spark.createDataFrame([ - (0, Vectors.dense(1.0, 0.5, -1.0)), - (1, Vectors.dense(2.0, 1.0, 1.0)), - (2, Vectors.dense(4.0, 10.0, 2.0)) - ]).toDF("id", "features") - model = Normalizer(inputCol='features', outputCol='norm_feature', p=1.0) + data = self.spark.createDataFrame( + [ + (0, Vectors.dense(1.0, 0.5, -1.0)), + (1, Vectors.dense(2.0, 1.0, 1.0)), + (2, Vectors.dense(4.0, 10.0, 2.0)), + ] + ).toDF("id", "features") + model = Normalizer(inputCol="features", outputCol="norm_feature", p=1.0) - model_onnx = convert_sparkml(model, 'Sparkml Normalizer', [('features', FloatTensorType([None, 3]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Normalizer", + [("features", FloatTensorType([None, 3]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().norm_feature.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlNormalizer") + expected = ( + predicted.toPandas() + .norm_feature.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlNormalizer" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['norm_feature'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["norm_feature"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_normalizer_2(self): - data = self.spark.createDataFrame([ - (0, Vectors.dense(1.0, 0.5, -1.0)), - (1, Vectors.dense(2.0, 1.0, 1.0)), - (2, Vectors.dense(4.0, 10.0, 2.0)) - ]).toDF("id", "features") - model = Normalizer(inputCol='features', outputCol='norm_feature', p=2.0) + data = self.spark.createDataFrame( + [ + (0, Vectors.dense(1.0, 0.5, -1.0)), + (1, Vectors.dense(2.0, 1.0, 1.0)), + (2, Vectors.dense(4.0, 10.0, 2.0)), + ] + ).toDF("id", "features") + model = Normalizer(inputCol="features", outputCol="norm_feature", p=2.0) - model_onnx = convert_sparkml(model, 'Sparkml Normalizer', [('features', FloatTensorType([None, 3]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Normalizer", + [("features", FloatTensorType([None, 3]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().norm_feature.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlNormalizer") + expected = ( + predicted.toPandas() + .norm_feature.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlNormalizer" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['norm_feature'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["norm_feature"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_one_vs_rest.py b/tests/sparkml/test_one_vs_rest.py index 6887198b..0a00130c 100644 --- a/tests/sparkml/test_one_vs_rest.py +++ b/tests/sparkml/test_one_vs_rest.py @@ -13,7 +13,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -21,35 +25,47 @@ class TestSparkmOneVsRest(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), - 'Need Greater Opset 9') + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_one_vs_rest(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "sample_multiclass_classification_data.txt") + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "sample_multiclass_classification_data.txt" + ) data = self.spark.read.format("libsvm").load(input_path) lr = LogisticRegression(maxIter=100, tol=0.0001, regParam=0.01) ovr = OneVsRest(classifier=lr) model = ovr.fit(data) feature_count = data.first()[1].size - model_onnx = convert_sparkml(model, 'Sparkml OneVsRest', [ - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml OneVsRest", + [("features", FloatTensorType([None, feature_count]))], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) expected = [ predicted.toPandas().prediction.values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlOneVsRest") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlOneVsRest" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["prediction"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_onehot_encoder.py b/tests/sparkml/test_onehot_encoder.py index 07898a97..482a7b29 100644 --- a/tests/sparkml/test_onehot_encoder.py +++ b/tests/sparkml/test_onehot_encoder.py @@ -8,7 +8,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -16,79 +20,194 @@ class TestSparkmlOneHotEncoder(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_onehot_encoder_1(self): """ - Testing ONNX conversion for Spark OneHotEncoder when handleInvalid is set to "error" and dropLast set to False. + Testing ONNX conversion for Spark OneHotEncoder when handleInvalid + is set to "error" and dropLast set to False. """ - encoder = OneHotEncoder(inputCols=['index1', 'index2'], outputCols=['index1Vec', 'index2Vec'], handleInvalid="error", dropLast=False) - data = self.spark.createDataFrame([(0.0,5.0,), (1.0,4.0,), (2.0,3.0,), (2.0,2.0,), (0.0,1.0,), (2.0,0.0,)], ['index1','index2']) + encoder = OneHotEncoder( + inputCols=["index1", "index2"], + outputCols=["index1Vec", "index2Vec"], + handleInvalid="error", + dropLast=False, + ) + data = self.spark.createDataFrame( + [ + ( + 0.0, + 5.0, + ), + ( + 1.0, + 4.0, + ), + ( + 2.0, + 3.0, + ), + ( + 2.0, + 2.0, + ), + ( + 0.0, + 1.0, + ), + ( + 2.0, + 0.0, + ), + ], + ["index1", "index2"], + ) model = encoder.fit(data) model_onnx = convert_sparkml( - model, 'Sparkml OneHotEncoder', [('index1', FloatTensorType([None, 1])), ('index2', FloatTensorType([None, 1]))], - target_opset=TARGET_OPSET) + model, + "Sparkml OneHotEncoder", + [ + ("index1", FloatTensorType([None, 1])), + ("index2", FloatTensorType([None, 1])), + ], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(data) data_np = { - "index1": data.select("index1").toPandas().values.astype(numpy.float32), - "index2": data.select("index2").toPandas().values.astype(numpy.float32) - } - - predicted_np_1 = predicted.select("index1Vec").toPandas().index1Vec.apply(lambda x: x.toArray()).values - predicted_np_2 = predicted.select("index2Vec").toPandas().index2Vec.apply(lambda x: x.toArray()).values - expected = {"index1Vec": numpy.asarray(predicted_np_1.tolist()), "index2Vec": numpy.asarray(predicted_np_2.tolist())} - - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlOneHotEncoder") + "index1": data.select("index1").toPandas().values.astype(numpy.float32), + "index2": data.select("index2").toPandas().values.astype(numpy.float32), + } + + predicted_np_1 = ( + predicted.select("index1Vec") + .toPandas() + .index1Vec.apply(lambda x: x.toArray()) + .values + ) + predicted_np_2 = ( + predicted.select("index2Vec") + .toPandas() + .index2Vec.apply(lambda x: x.toArray()) + .values + ) + expected = { + "index1Vec": numpy.asarray(predicted_np_1.tolist()), + "index2Vec": numpy.asarray(predicted_np_2.tolist()), + } + + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlOneHotEncoder" + ) onnx_model_path = paths[-1] - - output_names = ['index1Vec', 'index2Vec'] + + output_names = ["index1Vec", "index2Vec"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) actual_output = dict(zip(output_names, output)) compare_results(expected["index1Vec"], actual_output["index1Vec"], decimal=5) compare_results(expected["index2Vec"], actual_output["index2Vec"], decimal=5) - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_onehot_encoder_2(self): """ - Testing ONNX conversion for Spark OneHotEncoder when handleInvalid is set to "keep" and dropLast set to True. + Testing ONNX conversion for Spark OneHotEncoder when handleInvalid + is set to "keep" and dropLast set to True. """ - encoder = OneHotEncoder(inputCols=['index1', 'index2'], outputCols=['index1Vec', 'index2Vec'], handleInvalid="keep", dropLast=True) - data = self.spark.createDataFrame([(0.0,5.0,), (1.0,4.0,), (2.0,3.0,), (2.0,2.0,), (0.0,1.0,), (2.0,0.0,)], ['index1','index2']) - test = self.spark.createDataFrame([(3.0,7.0,)], ['index1','index2']) # invalid data + encoder = OneHotEncoder( + inputCols=["index1", "index2"], + outputCols=["index1Vec", "index2Vec"], + handleInvalid="keep", + dropLast=True, + ) + data = self.spark.createDataFrame( + [ + ( + 0.0, + 5.0, + ), + ( + 1.0, + 4.0, + ), + ( + 2.0, + 3.0, + ), + ( + 2.0, + 2.0, + ), + ( + 0.0, + 1.0, + ), + ( + 2.0, + 0.0, + ), + ], + ["index1", "index2"], + ) + test = self.spark.createDataFrame( + [ + ( + 3.0, + 7.0, + ) + ], + ["index1", "index2"], + ) # invalid data model = encoder.fit(data) model_onnx = convert_sparkml( - model, 'Sparkml OneHotEncoder', [('index1', FloatTensorType([None, 1])), ('index2', FloatTensorType([None, 1]))], - target_opset=TARGET_OPSET) + model, + "Sparkml OneHotEncoder", + [ + ("index1", FloatTensorType([None, 1])), + ("index2", FloatTensorType([None, 1])), + ], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(test) data_np = { - "index1": test.select("index1").toPandas().values.astype(numpy.float32), - "index2": test.select("index2").toPandas().values.astype(numpy.float32) - } - - predicted_np_1 = predicted.select("index1Vec").toPandas().index1Vec.apply(lambda x: x.toArray()).values - predicted_np_2 = predicted.select("index2Vec").toPandas().index2Vec.apply(lambda x: x.toArray()).values - expected = {"index1Vec": numpy.asarray(predicted_np_1.tolist()), "index2Vec": numpy.asarray(predicted_np_2.tolist())} - - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlOneHotEncoder") + "index1": test.select("index1").toPandas().values.astype(numpy.float32), + "index2": test.select("index2").toPandas().values.astype(numpy.float32), + } + + predicted_np_1 = ( + predicted.select("index1Vec") + .toPandas() + .index1Vec.apply(lambda x: x.toArray()) + .values + ) + predicted_np_2 = ( + predicted.select("index2Vec") + .toPandas() + .index2Vec.apply(lambda x: x.toArray()) + .values + ) + expected = { + "index1Vec": numpy.asarray(predicted_np_1.tolist()), + "index2Vec": numpy.asarray(predicted_np_2.tolist()), + } + + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlOneHotEncoder" + ) onnx_model_path = paths[-1] - - output_names = ['index1Vec', 'index2Vec'] + + output_names = ["index1Vec", "index2Vec"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) actual_output = dict(zip(output_names, output)) compare_results(expected["index1Vec"], actual_output["index1Vec"], decimal=5) compare_results(expected["index2Vec"], actual_output["index2Vec"], decimal=5) + if __name__ == "__main__": unittest.main() diff --git a/tests/sparkml/test_pipeline.py b/tests/sparkml/test_pipeline.py index cc49feb2..aa7ac69e 100644 --- a/tests/sparkml/test_pipeline.py +++ b/tests/sparkml/test_pipeline.py @@ -13,7 +13,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -21,141 +25,237 @@ class TestSparkmlPipeline(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_pipeline_4_stage(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "AdultCensusIncomeOriginal.csv") - full_data = self.spark.read.format('csv')\ - .options(header='true', inferschema='true').load(input_path) - cols = ['workclass', 'education', 'marital_status'] - training_data, test_data = full_data.select('income', *cols).limit(1000).randomSplit([0.9, 0.1],seed=1) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "AdultCensusIncomeOriginal.csv" + ) + full_data = ( + self.spark.read.format("csv") + .options(header="true", inferschema="true") + .load(input_path) + ) + cols = ["workclass", "education", "marital_status"] + training_data, test_data = ( + full_data.select("income", *cols) + .limit(1000) + .randomSplit([0.9, 0.1], seed=1) + ) stages = [] for col in cols: - stages.append(StringIndexer(inputCol=col, outputCol=col+'_index', handleInvalid='skip')) - stages.append(OneHotEncoder(inputCols=[col+'_index'], outputCols=[col+'_vec'], dropLast=False)) + stages.append( + StringIndexer( + inputCol=col, outputCol=col + "_index", handleInvalid="skip" + ) + ) + stages.append( + OneHotEncoder( + inputCols=[col + "_index"], + outputCols=[col + "_vec"], + dropLast=False, + ) + ) - stages.append(VectorAssembler(inputCols=[c+'_vec' for c in cols], outputCol='features')) - stages.append(StringIndexer(inputCol='income', outputCol='label', handleInvalid='skip')) + stages.append( + VectorAssembler(inputCols=[c + "_vec" for c in cols], outputCol="features") + ) + stages.append( + StringIndexer(inputCol="income", outputCol="label", handleInvalid="skip") + ) stages.append(LogisticRegression(maxIter=100, tol=0.0001)) pipeline = Pipeline(stages=stages) model = pipeline.fit(training_data) - model_onnx = convert_sparkml(model, 'Sparkml Pipeline', [ - ('income', StringTensorType([None, 1])), - ('workclass', StringTensorType([None, 1])), - ('education', StringTensorType([None, 1])), - ('marital_status', StringTensorType([None, 1])) - ], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Pipeline", + [ + ("income", StringTensorType([None, 1])), + ("workclass", StringTensorType([None, 1])), + ("education", StringTensorType([None, 1])), + ("marital_status", StringTensorType([None, 1])), + ], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(test_data) data_np = { - 'income': test_data.select('income').toPandas().values, - 'workclass': test_data.select('workclass').toPandas().values, - 'education': test_data.select('education').toPandas().values, - 'marital_status': test_data.select('marital_status').toPandas().values + "income": test_data.select("income").toPandas().values, + "workclass": test_data.select("workclass").toPandas().values, + "education": test_data.select("education").toPandas().values, + "marital_status": test_data.select("marital_status").toPandas().values, } expected = [ predicted.toPandas().label.values.astype(numpy.float32), predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlPipeline_4Stage") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlPipeline_4Stage" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['label', 'prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["label", "prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_pipeline_3_stage(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "AdultCensusIncomeOriginal.csv") - full_data = self.spark.read.format('csv')\ - .options(header='true', inferschema='true').load(input_path) - cols = ['workclass', 'education', 'marital_status'] - training_data, test_data = full_data.select(*cols).limit(1000).randomSplit([0.9, 0.1], seed=1) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "AdultCensusIncomeOriginal.csv" + ) + full_data = ( + self.spark.read.format("csv") + .options(header="true", inferschema="true") + .load(input_path) + ) + cols = ["workclass", "education", "marital_status"] + training_data, test_data = ( + full_data.select(*cols).limit(1000).randomSplit([0.9, 0.1], seed=1) + ) stages = [] for col in cols: - stages.append(StringIndexer(inputCol=col, outputCol=col+'_index', handleInvalid='skip')) + stages.append( + StringIndexer( + inputCol=col, outputCol=col + "_index", handleInvalid="skip" + ) + ) # we need the dropLast option otherwise when assembled together (below) # we won't be able to expand the features without difficulties - stages.append(OneHotEncoder(inputCols=[col+'_index'], outputCols=[col+'_vec'], dropLast=False)) + stages.append( + OneHotEncoder( + inputCols=[col + "_index"], + outputCols=[col + "_vec"], + dropLast=False, + ) + ) - stages.append(VectorAssembler(inputCols=[c+'_vec' for c in cols], outputCol='features')) + stages.append( + VectorAssembler(inputCols=[c + "_vec" for c in cols], outputCol="features") + ) pipeline = Pipeline(stages=stages) model = pipeline.fit(training_data) - model_onnx = convert_sparkml(model, 'Sparkml Pipeline', [ - ('workclass', StringTensorType([None, 1])), - ('education', StringTensorType([None, 1])), - ('marital_status', StringTensorType([None, 1])) - ], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Pipeline", + [ + ("workclass", StringTensorType([None, 1])), + ("education", StringTensorType([None, 1])), + ("marital_status", StringTensorType([None, 1])), + ], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(test_data) data_np = { - 'workclass': test_data.select('workclass').toPandas().values, - 'education': test_data.select('education').toPandas().values, - 'marital_status': test_data.select('marital_status').toPandas().values + "workclass": test_data.select("workclass").toPandas().values, + "education": test_data.select("education").toPandas().values, + "marital_status": test_data.select("marital_status").toPandas().values, } - expected = predicted.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlPipeline_3Stage") + expected = ( + predicted.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlPipeline_3Stage" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['features'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["features"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_pipeline_2_stage(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) - input_path = os.path.join(this_script_dir, "data", "AdultCensusIncomeOriginal.csv") - full_data = self.spark.read.format('csv')\ - .options(header='true', inferschema='true').load(input_path) - cols = ['workclass', 'education', 'marital_status'] - training_data, test_data = full_data.select(*cols).limit(1000).randomSplit([0.9, 0.1], seed=1) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) + input_path = os.path.join( + this_script_dir, "data", "AdultCensusIncomeOriginal.csv" + ) + full_data = ( + self.spark.read.format("csv") + .options(header="true", inferschema="true") + .load(input_path) + ) + cols = ["workclass", "education", "marital_status"] + training_data, test_data = ( + full_data.select(*cols).limit(1000).randomSplit([0.9, 0.1], seed=1) + ) stages = [] for col in cols: - stages.append(StringIndexer(inputCol=col, outputCol=col+'_index', handleInvalid='skip')) - stages.append(OneHotEncoder(inputCols=[col+'_index'], outputCols=[col+'_vec'], dropLast=False)) + stages.append( + StringIndexer( + inputCol=col, outputCol=col + "_index", handleInvalid="skip" + ) + ) + stages.append( + OneHotEncoder( + inputCols=[col + "_index"], + outputCols=[col + "_vec"], + dropLast=False, + ) + ) pipeline = Pipeline(stages=stages) model = pipeline.fit(training_data) - model_onnx = convert_sparkml(model, 'Sparkml Pipeline', [ - ('workclass', StringTensorType([None, 1])), - ('education', StringTensorType([None, 1])), - ('marital_status', StringTensorType([None, 1])) - ], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Pipeline", + [ + ("workclass", StringTensorType([None, 1])), + ("education", StringTensorType([None, 1])), + ("marital_status", StringTensorType([None, 1])), + ], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(test_data) data_np = { - 'workclass': test_data.select('workclass').toPandas().values, - 'education': test_data.select('education').toPandas().values, - 'marital_status': test_data.select('marital_status').toPandas().values + "workclass": test_data.select("workclass").toPandas().values, + "education": test_data.select("education").toPandas().values, + "marital_status": test_data.select("marital_status").toPandas().values, } predicted_np = [ - predicted.toPandas().workclass_vec.apply(lambda x: pandas.Series(x.toArray())).values, - predicted.toPandas().education_vec.apply(lambda x: pandas.Series(x.toArray())).values, - predicted.toPandas().marital_status_vec.apply(lambda x: pandas.Series(x.toArray())).values - ] - + predicted.toPandas() + .workclass_vec.apply(lambda x: pandas.Series(x.toArray())) + .values, + predicted.toPandas() + .education_vec.apply(lambda x: pandas.Series(x.toArray())) + .values, + predicted.toPandas() + .marital_status_vec.apply(lambda x: pandas.Series(x.toArray())) + .values, + ] + expected = [numpy.asarray(row) for row in predicted_np] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlPipeline_2Stage") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlPipeline_2Stage" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['workclass_vec', 'education_vec', 'marital_status_vec'], - data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["workclass_vec", "education_vec", "marital_status_vec"], + data_np, + onnx_model_path, + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_polynomial_expansion.py b/tests/sparkml/test_polynomial_expansion.py index cbc4ca69..4e62c4c4 100644 --- a/tests/sparkml/test_polynomial_expansion.py +++ b/tests/sparkml/test_polynomial_expansion.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,30 +22,45 @@ class TestSparkmlPolynomialExpansion(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_polynomial_expansion(self): - data = self.spark.createDataFrame([ - (Vectors.dense([1.2, 3.2, 1.3, -5.6]),), - (Vectors.dense([4.3, -3.2, 5.7, 1.0]),), - (Vectors.dense([0, 3.2, 4.7, -8.9]),) - ], ["dense"]) + data = self.spark.createDataFrame( + [ + (Vectors.dense([1.2, 3.2, 1.3, -5.6]),), + (Vectors.dense([4.3, -3.2, 5.7, 1.0]),), + (Vectors.dense([0, 3.2, 4.7, -8.9]),), + ], + ["dense"], + ) model = PolynomialExpansion(degree=2, inputCol="dense", outputCol="expanded") # the input name should match that of what StringIndexer.inputCol feature_count = data.first()[0].size - model_onnx = convert_sparkml(model, 'Sparkml PolynomialExpansion', [('dense', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml PolynomialExpansion", + [("dense", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().expanded.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().dense.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlPolynomialExpansion") + expected = ( + predicted.toPandas() + .expanded.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .dense.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlPolynomialExpansion" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['expanded'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["expanded"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_random_forest_classifier.py b/tests/sparkml/test_random_forest_classifier.py index 6a5c2f9d..c51d60ac 100644 --- a/tests/sparkml/test_random_forest_classifier.py +++ b/tests/sparkml/test_random_forest_classifier.py @@ -15,7 +15,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType, FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase from pyspark.ml.feature import StringIndexer, VectorIndexer @@ -24,52 +28,81 @@ class TestSparkmRandomForestClassifier(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), - 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_random_forest_classification(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) # # truncate the features # feature_count = 5 - self.spark.udf.register("truncateFeatures", - lambda x: SparseVector(feature_count, range(0,feature_count), x.toArray()[125:130]), - VectorUDT()) - data = original_data.selectExpr("cast(label as string) as label", "truncateFeatures(features) as features") + self.spark.udf.register( + "truncateFeatures", + lambda x: SparseVector( + feature_count, range(0, feature_count), x.toArray()[125:130] + ), + VectorUDT(), + ) + data = original_data.selectExpr( + "cast(label as string) as label", "truncateFeatures(features) as features" + ) label_indexer = StringIndexer(inputCol="label", outputCol="indexedLabel") - feature_indexer = VectorIndexer(inputCol="features", outputCol="indexedFeatures", - maxCategories=10, handleInvalid='keep') + feature_indexer = VectorIndexer( + inputCol="features", + outputCol="indexedFeatures", + maxCategories=10, + handleInvalid="keep", + ) - rf = RandomForestClassifier(labelCol="indexedLabel", featuresCol="indexedFeatures", numTrees=10) + rf = RandomForestClassifier( + labelCol="indexedLabel", featuresCol="indexedFeatures", numTrees=10 + ) pipeline = Pipeline(stages=[label_indexer, feature_indexer, rf]) model = pipeline.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml RandomForest Classifier', [ - ('label', StringTensorType([None, 1])), - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml RandomForest Classifier", + [ + ("label", StringTensorType([None, 1])), + ("features", FloatTensorType([None, feature_count])), + ], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) data_np = { - 'label': data.toPandas().label.values.reshape((-1, 1)), - 'features': data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + "label": data.toPandas().label.values.reshape((-1, 1)), + "features": data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), } expected = [ predicted.toPandas().indexedLabel.values.astype(numpy.int64), predicted.toPandas().prediction.values.astype(numpy.float32), - predicted.toPandas().probability.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + predicted.toPandas() + .probability.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlRandomForestClassifier") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlRandomForestClassifier", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['indexedLabel', 'prediction', 'probability'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["indexedLabel", "prediction", "probability"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_random_forest_classifier_tree.py b/tests/sparkml/test_random_forest_classifier_tree.py index c36f8ad2..283253dd 100644 --- a/tests/sparkml/test_random_forest_classifier_tree.py +++ b/tests/sparkml/test_random_forest_classifier_tree.py @@ -6,19 +6,15 @@ import os import packaging.version as pv import onnx -import pandas import numpy from numpy.random import randint from onnxruntime import InferenceSession -from pyspark.ml import Pipeline from pyspark.ml.classification import RandomForestClassifier -from pyspark.ml.linalg import VectorUDT, SparseVector -from pyspark.ml.feature import StringIndexer, VectorIndexer, VectorAssembler +from pyspark.ml.feature import VectorAssembler from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml -from onnxmltools.convert.common.data_types import StringTensorType, FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from onnxmltools.convert.common.data_types import FloatTensorType from tests.sparkml import SparkMlTestCase @@ -26,42 +22,51 @@ class TestSparkmRandomForestClassifierTree(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), - 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_random_forest_classification_tree(self): FEATURE_LEN = 32 - def infer_from_onnx(model_onnx, input_list): + def infer_from_onnx(model_onnx, input_list): sess = InferenceSession(model_onnx.SerializeToString()) input_name = sess.get_inputs()[0].name - pred_onx = sess.run(None, {input_name: numpy.array(input_list, numpy.float32)}) + pred_onx = sess.run( + None, {input_name: numpy.array(input_list, numpy.float32)} + ) return pred_onx def export_as_onnx(model): model_onnx = convert_sparkml( - model, "Phish Classifier", + model, + "Phish Classifier", [("features", FloatTensorType([None, FEATURE_LEN]))], - spark_session=self.spark, target_opset=TARGET_OPSET) + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) return model_onnx def create_model(input_path): df = self.spark.read.csv(input_path, header=True, inferSchema=True) vec_assembler = VectorAssembler( - inputCols=["c" + str(i) for i in range(FEATURE_LEN)], outputCol="features") + inputCols=["c" + str(i) for i in range(FEATURE_LEN)], + outputCol="features", + ) data = vec_assembler.transform(df) - rf = RandomForestClassifier(labelCol="label", featuresCol="features", numTrees=5) + rf = RandomForestClassifier( + labelCol="label", featuresCol="features", numTrees=5 + ) model = rf.fit(dataset=data) # RandomForestClassificationModel # model.save("./dummy_spark_model/model/") return model - - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "features_32.csv") model = create_model(input_path) model_onnx = export_as_onnx(model) diff --git a/tests/sparkml/test_random_forest_regressor.py b/tests/sparkml/test_random_forest_regressor.py index 7e7ad6b1..3215e6f8 100644 --- a/tests/sparkml/test_random_forest_regressor.py +++ b/tests/sparkml/test_random_forest_regressor.py @@ -15,7 +15,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType, StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase from pyspark.ml.feature import VectorIndexer, StringIndexer @@ -24,51 +28,79 @@ class TestSparkmRandomForestRegressor(SparkMlTestCase): - - @unittest.skipIf(sys.platform == 'win32', - reason="UnsatisfiedLinkError") - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), - 'Need Greater Opset 9') + @unittest.skipIf(sys.platform == "win32", reason="UnsatisfiedLinkError") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_random_forest_regression(self): - this_script_dir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))) + this_script_dir = os.path.dirname( + os.path.abspath(inspect.getfile(inspect.currentframe())) + ) input_path = os.path.join(this_script_dir, "data", "sample_libsvm_data.txt") original_data = self.spark.read.format("libsvm").load(input_path) # # truncate the features # feature_count = 5 - self.spark.udf.register("truncateFeatures", - lambda x: SparseVector(feature_count, range(0,feature_count), x.toArray()[125:130]), - VectorUDT()) - data = original_data.selectExpr("cast(label as string) as label", "truncateFeatures(features) as features") + self.spark.udf.register( + "truncateFeatures", + lambda x: SparseVector( + feature_count, range(0, feature_count), x.toArray()[125:130] + ), + VectorUDT(), + ) + data = original_data.selectExpr( + "cast(label as string) as label", "truncateFeatures(features) as features" + ) label_indexer = StringIndexer(inputCol="label", outputCol="indexedLabel") - feature_indexer = VectorIndexer(inputCol="features", outputCol="indexedFeatures", - maxCategories=10, handleInvalid='error') + feature_indexer = VectorIndexer( + inputCol="features", + outputCol="indexedFeatures", + maxCategories=10, + handleInvalid="error", + ) - rf = RandomForestRegressor(labelCol="indexedLabel", featuresCol="indexedFeatures", numTrees=3) + rf = RandomForestRegressor( + labelCol="indexedLabel", featuresCol="indexedFeatures", numTrees=3 + ) pipeline = Pipeline(stages=[label_indexer, feature_indexer, rf]) model = pipeline.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml RandomForest Regressor', [ - ('label', StringTensorType([None, 1])), - ('features', FloatTensorType([None, feature_count])) - ], spark_session=self.spark, target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml RandomForest Regressor", + [ + ("label", StringTensorType([None, 1])), + ("features", FloatTensorType([None, feature_count])), + ], + spark_session=self.spark, + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data.limit(1)) data_np = { - 'label': data.limit(1).toPandas().label.values.reshape((-1, 1)), - 'features': data.limit(1).toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + "label": data.limit(1).toPandas().label.values.reshape((-1, 1)), + "features": data.limit(1) + .toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), } expected = [ predicted.toPandas().indexedLabel.values.astype(numpy.int64), - predicted.toPandas().prediction.values.astype(numpy.float32) + predicted.toPandas().prediction.values.astype(numpy.float32), ] - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlRandomForestRegressor") + paths = save_data_models( + data_np, + expected, + model, + model_onnx, + basename="SparkmlRandomForestRegressor", + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['indexedLabel', 'prediction'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["indexedLabel", "prediction"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_scaler.py b/tests/sparkml/test_scaler.py index 2827fdc3..d350c27d 100644 --- a/tests/sparkml/test_scaler.py +++ b/tests/sparkml/test_scaler.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,77 +22,165 @@ class TestSparkmlScaler(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_maxabs_scaler(self): - data = self.spark.createDataFrame([ - (0, Vectors.dense([1.0, 0.1, -1.0]),), - (1, Vectors.dense([2.0, 1.1, 1.0]),), - (2, Vectors.dense([3.0, 10.1, 3.0]),) - ], ["id", "features"]) - scaler = MaxAbsScaler(inputCol='features', outputCol='scaled_features') + data = self.spark.createDataFrame( + [ + ( + 0, + Vectors.dense([1.0, 0.1, -1.0]), + ), + ( + 1, + Vectors.dense([2.0, 1.1, 1.0]), + ), + ( + 2, + Vectors.dense([3.0, 10.1, 3.0]), + ), + ], + ["id", "features"], + ) + scaler = MaxAbsScaler(inputCol="features", outputCol="scaled_features") model = scaler.fit(data) # the input names must match the inputCol(s) above - model_onnx = convert_sparkml(model, 'Sparkml MaxAbsScaler', [('features', FloatTensorType([None, 3]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml MaxAbsScaler", + [("features", FloatTensorType([None, 3]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().scaled_features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlMaxAbsScaler") + expected = ( + predicted.toPandas() + .scaled_features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlMaxAbsScaler" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['scaled_features'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["scaled_features"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_minmax_scaler(self): - data = self.spark.createDataFrame([ - (0, Vectors.dense([1.0, 0.1, -1.0]),), - (1, Vectors.dense([2.0, 1.1, 1.0]),), - (2, Vectors.dense([3.0, 10.1, 3.0]),) - ], ["id", "features"]) - scaler = MinMaxScaler(inputCol='features', outputCol='scaled_features') + data = self.spark.createDataFrame( + [ + ( + 0, + Vectors.dense([1.0, 0.1, -1.0]), + ), + ( + 1, + Vectors.dense([2.0, 1.1, 1.0]), + ), + ( + 2, + Vectors.dense([3.0, 10.1, 3.0]), + ), + ], + ["id", "features"], + ) + scaler = MinMaxScaler(inputCol="features", outputCol="scaled_features") model = scaler.fit(data) # the input names must match the inputCol(s) above - model_onnx = convert_sparkml(model, 'Sparkml MinMaxScaler', [('features', FloatTensorType([None, 3]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml MinMaxScaler", + [("features", FloatTensorType([None, 3]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().scaled_features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlMinMaxScaler") + expected = ( + predicted.toPandas() + .scaled_features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlMinMaxScaler" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['scaled_features'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["scaled_features"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_standard_scaler(self): - data = self.spark.createDataFrame([ - (0, Vectors.dense([1.0, 0.1, -1.0]),), - (1, Vectors.dense([2.0, 1.1, 1.0]),), - (2, Vectors.dense([3.0, 10.1, 3.0]),) - ], ["id", "features"]) - scaler = StandardScaler(inputCol='features', outputCol='scaled_features', withStd=True, withMean=True) + data = self.spark.createDataFrame( + [ + ( + 0, + Vectors.dense([1.0, 0.1, -1.0]), + ), + ( + 1, + Vectors.dense([2.0, 1.1, 1.0]), + ), + ( + 2, + Vectors.dense([3.0, 10.1, 3.0]), + ), + ], + ["id", "features"], + ) + scaler = StandardScaler( + inputCol="features", + outputCol="scaled_features", + withStd=True, + withMean=True, + ) model = scaler.fit(data) # the input names must match the inputCol(s) above - model_onnx = convert_sparkml(model, 'Sparkml StandardScaler', [('features', FloatTensorType([None, 3]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml StandardScaler", + [("features", FloatTensorType([None, 3]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().scaled_features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlStandardScaler") + expected = ( + predicted.toPandas() + .scaled_features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlStandardScaler" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['scaled_features'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model( + ["scaled_features"], data_np, onnx_model_path + ) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_stop_words_remover.py b/tests/sparkml/test_stop_words_remover.py index e4c60b88..dea1ffaf 100644 --- a/tests/sparkml/test_stop_words_remover.py +++ b/tests/sparkml/test_stop_words_remover.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,26 +22,31 @@ class TestSparkmlStopWordsRemover(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.5'), - 'Need Greater Opset 10') + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.5"), "Need Greater Opset 10" + ) def test_stop_words_remover2(self): data = self.spark.createDataFrame([(["a", "b", "c"],)], ["text"]) model = StopWordsRemover(inputCol="text", outputCol="words", stopWords=["b"]) - model_onnx = convert_sparkml(model, 'Sparkml StopWordsRemover', - [('text', StringTensorType([None]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml StopWordsRemover", + [("text", StringTensorType([None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) expected = numpy.array(predicted.toPandas().words.values[0]) data_np = numpy.array(data.toPandas().text.values[0]) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlStopWordsRemover") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlStopWordsRemover" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['words'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["words"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_string_indexer.py b/tests/sparkml/test_string_indexer.py index 8687aa1a..7a188c58 100644 --- a/tests/sparkml/test_string_indexer.py +++ b/tests/sparkml/test_string_indexer.py @@ -7,7 +7,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase import numpy @@ -17,12 +21,19 @@ class TestSparkmlStringIndexer(SparkMlTestCase): @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_string_indexer(self): - indexer = StringIndexer(inputCol="cat1", outputCol="cat1_index", handleInvalid="skip") - data = self.spark.createDataFrame([("a",), ("b",), ("c",), ("a",), ("a",), ("c",)], ["cat1"]) + indexer = StringIndexer( + inputCol="cat1", outputCol="cat1_index", handleInvalid="skip" + ) + data = self.spark.createDataFrame( + [("a",), ("b",), ("c",), ("a",), ("a",), ("c",)], ["cat1"] + ) model = indexer.fit(data) # the input name should match that of what StringIndexer.inputCol model_onnx = convert_sparkml( - model, "Sparkml StringIndexer", [("cat1", StringTensorType([None, 1]))], target_opset=TARGET_OPSET + model, + "Sparkml StringIndexer", + [("cat1", StringTensorType([None, 1]))], + target_opset=TARGET_OPSET, ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) @@ -30,7 +41,9 @@ def test_model_string_indexer(self): predicted = model.transform(data) expected = predicted.select("cat1_index").toPandas().values data_np = data.select("cat1").toPandas().values - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlStringIndexer") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlStringIndexer" + ) onnx_model_path = paths[-1] output, output_shapes = run_onnx_model(["cat1_index"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) @@ -38,17 +51,48 @@ def test_model_string_indexer(self): @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_string_indexer_multiple_cols(self): indexer = StringIndexer( - inputCols=["cat1", "cat2"], outputCols=["cat1_index", "cat2_index"], handleInvalid="skip" + inputCols=["cat1", "cat2"], + outputCols=["cat1_index", "cat2_index"], + handleInvalid="skip", ) data = self.spark.createDataFrame( - [("a", "x",), ("b", "x",), ("c", "y",), ("a", "y",), ("a", "z",), ("c", "z",)], ["cat1", "cat2"] + [ + ( + "a", + "x", + ), + ( + "b", + "x", + ), + ( + "c", + "y", + ), + ( + "a", + "y", + ), + ( + "a", + "z", + ), + ( + "c", + "z", + ), + ], + ["cat1", "cat2"], ) model = indexer.fit(data) # the input name should match that of what StringIndexer.inputCol model_onnx = convert_sparkml( model, "Sparkml StringIndexer", - [("cat1", StringTensorType([None, 1])), ("cat2", StringTensorType([None, 1]))], + [ + ("cat1", StringTensorType([None, 1])), + ("cat2", StringTensorType([None, 1])), + ], target_opset=TARGET_OPSET, ) self.assertTrue(model_onnx is not None) @@ -57,14 +101,20 @@ def test_model_string_indexer_multiple_cols(self): predicted = model.transform(data) expected = { - "cat1_index": predicted.select("cat1_index").toPandas().values.astype(numpy.int64), - "cat2_index": predicted.select("cat2_index").toPandas().values.astype(numpy.int64), + "cat1_index": predicted.select("cat1_index") + .toPandas() + .values.astype(numpy.int64), + "cat2_index": predicted.select("cat2_index") + .toPandas() + .values.astype(numpy.int64), } data_np = { "cat1": data.select("cat1").toPandas().values.astype(str), "cat2": data.select("cat2").toPandas().values.astype(str), } - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlStringIndexer") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlStringIndexer" + ) onnx_model_path = paths[-1] output_names = ["cat1_index", "cat2_index"] output, output_shapes = run_onnx_model(output_names, data_np, onnx_model_path) diff --git a/tests/sparkml/test_tokenizer.py b/tests/sparkml/test_tokenizer.py index 21c80015..86294ce1 100644 --- a/tests/sparkml/test_tokenizer.py +++ b/tests/sparkml/test_tokenizer.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,25 +22,30 @@ class TestSparkmlTokenizer(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.5'), - 'Need Greater Opset 10') + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.5"), "Need Greater Opset 10" + ) def test_tokenizer(self): data = self.spark.createDataFrame([("a b c",)], ["text"]) - model = Tokenizer(inputCol='text', outputCol='words') + model = Tokenizer(inputCol="text", outputCol="words") predicted = model.transform(data) - model_onnx = convert_sparkml(model, 'Sparkml Tokenizer', [ - ('text', StringTensorType([None]))], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Tokenizer", + [("text", StringTensorType([None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model expected = predicted.toPandas().words.apply(pandas.Series).values data_np = data.toPandas().text.values.reshape([-1]) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlTokenizer") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlTokenizer" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['words'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["words"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_tree_helper.py b/tests/sparkml/test_tree_helper.py index e5d33ef0..5bd4826e 100644 --- a/tests/sparkml/test_tree_helper.py +++ b/tests/sparkml/test_tree_helper.py @@ -1,43 +1,188 @@ # SPDX-License-Identifier: Apache-2.0 import unittest -from distutils.version import StrictVersion import numpy as np from onnx import TensorProto from onnx.defs import onnx_opset_version -from onnx.helper import make_node, make_graph, make_model, make_tensor_value_info, make_opsetid +from onnx.helper import ( + make_node, + make_graph, + make_model, + make_tensor_value_info, + make_opsetid, +) from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER -from onnxruntime import InferenceSession, __version__ as ort_version +from onnxruntime import InferenceSession from onnxmltools.convert.sparkml.operator_converters.tree_helper import Node TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) class TestSparkmDecisionTreeClassifierBig(unittest.TestCase): - @unittest.skipIf(TARGET_OPSET < 17, reason="Opset 17 is needed") def test_split(self): - attrs = { - 'class_ids': [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], - 'class_nodeids': [3, 3, 3, 4, 4, 4, 6, 6, 6, 7, 7, 7, 10, 10, 10, 11, 11, 11, 12, 12, 12], - 'class_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'class_weights': [0.4, 0.6, 0.0, 0.95, 0.04, 0.0, 0.0, 0.185, 0.814, 0.372, 0.62, 0.0, 0.74, 0.21, - 0.03, 0.0, 1.0, 0.0, 0.87, 0.05, 0.071], - 'classlabels_int64s': [0, 1, 2], - 'name': 'TreeEnsembleClassifier', - 'nodes_falsenodeids': [8, 5, 4, 0, 0, 7, 0, 0, 12, 11, 0, 0, 0], - 'nodes_featureids': [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0], - 'nodes_hitrates': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - 'nodes_missing_value_tracks_true': [False, False, False, False, False, False, False, False, - False, False, False, False, False], - 'nodes_modes': ['BRANCH_LEQ', '||', 'BRANCH_LEQ', 'LEAF', 'LEAF', 'BRANCH_LEQ', 'LEAF', - 'LEAF', 'BRANCH_LEQ', 'BRANCH_LEQ', 'LEAF', 'LEAF', 'LEAF'], - 'nodes_nodeids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], - 'nodes_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'nodes_truenodeids': [1, 2, 3, 0, 0, 6, 0, 0, 9, 10, 0, 0, 0], - 'nodes_values': [10.5, [55, 59, 65], 1.5, 0.0, 0.0, 8.5, 0.0, 0.0, 10.5, 9.5, 0.0, 0.0, 0.0], - 'post_transform': None} + "class_ids": [ + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + ], + "class_nodeids": [ + 3, + 3, + 3, + 4, + 4, + 4, + 6, + 6, + 6, + 7, + 7, + 7, + 10, + 10, + 10, + 11, + 11, + 11, + 12, + 12, + 12, + ], + "class_treeids": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ], + "class_weights": [ + 0.4, + 0.6, + 0.0, + 0.95, + 0.04, + 0.0, + 0.0, + 0.185, + 0.814, + 0.372, + 0.62, + 0.0, + 0.74, + 0.21, + 0.03, + 0.0, + 1.0, + 0.0, + 0.87, + 0.05, + 0.071, + ], + "classlabels_int64s": [0, 1, 2], + "name": "TreeEnsembleClassifier", + "nodes_falsenodeids": [8, 5, 4, 0, 0, 7, 0, 0, 12, 11, 0, 0, 0], + "nodes_featureids": [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0], + "nodes_hitrates": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ], + "nodes_missing_value_tracks_true": [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + ], + "nodes_modes": [ + "BRANCH_LEQ", + "||", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "LEAF", + ], + "nodes_nodeids": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], + "nodes_treeids": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "nodes_truenodeids": [1, 2, 3, 0, 0, 6, 0, 0, 9, 10, 0, 0, 0], + "nodes_values": [ + 10.5, + [55, 59, 65], + 1.5, + 0.0, + 0.0, + 8.5, + 0.0, + 0.0, + 10.5, + 9.5, + 0.0, + 0.0, + 0.0, + ], + "post_transform": None, + } root, nodes = Node.create(attrs) for node in root: key = node.nodes_treeids, node.nodes_nodeids @@ -58,60 +203,215 @@ def test_split(self): ns[key] = node # back to onnx - new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s= [0, 1, 2], domain="ai.onnx.ml") + new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) for k in attrs: if k in {"post_transform"}: continue if len(attrs[k]) > len(new_attrs[k]): raise AssertionError( - f"Issue with attribute {k!r}\n{attrs['nodes_modes']}\nbefore {attrs[k]!r}" + f"Issue with attribute {k!r}\n" + f"{attrs['nodes_modes']}\nbefore {attrs[k]!r}" f"\nafter {new_attrs[k]!r}\n{new_attrs['nodes_modes']}" ) - node = make_node(op_type="TreeEnsembleClassifier", inputs=["X"], outputs=["L", "Y"], **new_attrs) - X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None]) - Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None, None]) - L = make_tensor_value_info('L', TensorProto.INT64, [None]) + node = make_node( + op_type="TreeEnsembleClassifier", + inputs=["X"], + outputs=["L", "Y"], + **new_attrs, + ) + X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None]) + Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None]) + L = make_tensor_value_info("L", TensorProto.INT64, [None]) graph = make_graph([node], "n", [X], [L, Y]) opset_imports = [make_opsetid("", 17), make_opsetid("ai.onnx.ml", 3)] model = make_model(graph, opset_imports=opset_imports) sess = InferenceSession(model.SerializeToString()) x = np.arange(20).reshape((-1, 2)).astype(np.float32) - got = sess.run(None, {'X': x}) + got = sess.run(None, {"X": x}) self.assertEqual(len(got), 2) self.assertEqual(got[0].tolist(), [2, 2, 2, 2, 2, 1, 0, 0, 0, 0]) # again root, nodes = Node.create(new_attrs) root.unfold_rule_or() - new_new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s= [0, 1, 2], domain="ai.onnx.ml") + new_new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) self.assertEqual(new_attrs, new_new_attrs) @unittest.skipIf(TARGET_OPSET < 17, reason="Opset 17 is needed") def test_split_non_contiguous_ids(self): - attrs = { - 'class_ids': [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], - 'class_nodeids': [3, 3, 3, 4, 4, 4, 6, 6, 6, 7, 7, 7, 10, 10, 10, 11, 11, 11, 13, 13, 13], - 'class_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'class_weights': [0.4, 0.6, 0.0, 0.95, 0.04, 0.0, 0.0, 0.185, 0.814, 0.372, 0.62, 0.0, 0.74, 0.21, - 0.03, 0.0, 1.0, 0.0, 0.87, 0.05, 0.071], - 'classlabels_int64s': [0, 1, 2], - 'name': 'TreeEnsembleClassifier', - 'nodes_falsenodeids': [8, 5, 4, 0, 0, 7, 0, 0, 13, 11, 0, 0, 0], - 'nodes_featureids': [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0], - 'nodes_hitrates': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - 'nodes_missing_value_tracks_true': [False, False, False, False, False, False, False, False, - False, False, False, False, False], - 'nodes_modes': ['BRANCH_LEQ', '||', 'BRANCH_LEQ', 'LEAF', 'LEAF', 'BRANCH_LEQ', 'LEAF', - 'LEAF', 'BRANCH_LEQ', 'BRANCH_LEQ', 'LEAF', 'LEAF', 'LEAF'], - 'nodes_nodeids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13], - 'nodes_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'nodes_truenodeids': [1, 2, 3, 0, 0, 6, 0, 0, 9, 10, 0, 0, 0], - 'nodes_values': [10.5, [55, 59, 65], 1.5, 0.0, 0.0, 8.5, 0.0, 0.0, 10.5, 9.5, 0.0, 0.0, 0.0], - 'post_transform': None} + "class_ids": [ + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + ], + "class_nodeids": [ + 3, + 3, + 3, + 4, + 4, + 4, + 6, + 6, + 6, + 7, + 7, + 7, + 10, + 10, + 10, + 11, + 11, + 11, + 13, + 13, + 13, + ], + "class_treeids": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ], + "class_weights": [ + 0.4, + 0.6, + 0.0, + 0.95, + 0.04, + 0.0, + 0.0, + 0.185, + 0.814, + 0.372, + 0.62, + 0.0, + 0.74, + 0.21, + 0.03, + 0.0, + 1.0, + 0.0, + 0.87, + 0.05, + 0.071, + ], + "classlabels_int64s": [0, 1, 2], + "name": "TreeEnsembleClassifier", + "nodes_falsenodeids": [8, 5, 4, 0, 0, 7, 0, 0, 13, 11, 0, 0, 0], + "nodes_featureids": [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0], + "nodes_hitrates": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ], + "nodes_missing_value_tracks_true": [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + ], + "nodes_modes": [ + "BRANCH_LEQ", + "||", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "LEAF", + ], + "nodes_nodeids": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13], + "nodes_treeids": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "nodes_truenodeids": [1, 2, 3, 0, 0, 6, 0, 0, 9, 10, 0, 0, 0], + "nodes_values": [ + 10.5, + [55, 59, 65], + 1.5, + 0.0, + 0.0, + 8.5, + 0.0, + 0.0, + 10.5, + 9.5, + 0.0, + 0.0, + 0.0, + ], + "post_transform": None, + } root, nodes = Node.create(attrs) for node in root: key = node.nodes_treeids, node.nodes_nodeids @@ -132,61 +432,215 @@ def test_split_non_contiguous_ids(self): ns[key] = node # back to onnx - new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s= [0, 1, 2], domain="ai.onnx.ml") + new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) for k in attrs: if k in {"post_transform"}: continue if len(attrs[k]) > len(new_attrs[k]): raise AssertionError( - f"Issue with attribute {k!r}\n{attrs['nodes_modes']}\nbefore {attrs[k]!r}" + f"Issue with attribute {k!r}\n" + f"{attrs['nodes_modes']}\nbefore {attrs[k]!r}" f"\nafter {new_attrs[k]!r}\n{new_attrs['nodes_modes']}" ) - node = make_node(op_type="TreeEnsembleClassifier", inputs=["X"], outputs=["L", "Y"], **new_attrs) - X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None]) - Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None, None]) - L = make_tensor_value_info('L', TensorProto.INT64, [None]) + node = make_node( + op_type="TreeEnsembleClassifier", + inputs=["X"], + outputs=["L", "Y"], + **new_attrs, + ) + X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None]) + Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None]) + L = make_tensor_value_info("L", TensorProto.INT64, [None]) graph = make_graph([node], "n", [X], [L, Y]) opset_imports = [make_opsetid("", 17), make_opsetid("ai.onnx.ml", 3)] model = make_model(graph, opset_imports=opset_imports) sess = InferenceSession(model.SerializeToString()) x = np.arange(20).reshape((-1, 2)).astype(np.float32) - got = sess.run(None, {'X': x}) + got = sess.run(None, {"X": x}) self.assertEqual(len(got), 2) self.assertEqual(got[0].tolist(), [2, 2, 2, 2, 2, 1, 0, 0, 0, 0]) # again root, nodes = Node.create(new_attrs) root.unfold_rule_or() - new_new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s= [0, 1, 2], domain="ai.onnx.ml") + new_new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) self.assertEqual(new_attrs, new_new_attrs) @unittest.skipIf(TARGET_OPSET < 17, reason="Opset 17 is needed") def test_split_more_complex(self): - attrs = { - 'class_ids': [0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2], - 'class_nodeids': [3, 3, 3, 4, 4, 4, 6, 6, 6, 7, 7, 7, 10, 10, 10, 11, 11, 11, 12, 12, 12], - 'class_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'class_weights': [0.4, 0.6, 0.0, 0.95, 0.04, 0.0, 0.0, 0.185, 0.814, 0.372, 0.62, 0.0, 0.74, 0.21, - 0.03, 0.0, 1.0, 0.0, 0.87, 0.05, 0.071], - 'classlabels_int64s': [0, 1, 2], - 'name': 'TreeEnsembleClassifier', - 'nodes_falsenodeids': [8, 5, 4, 0, 0, 7, 0, 0, 12, 11, 0, 0, 0], - 'nodes_featureids': [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0], - 'nodes_hitrates': [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0], - 'nodes_missing_value_tracks_true': [False, False, False, False, False, False, False, False, - False, False, False, False, False], - 'nodes_modes': ['BRANCH_LEQ', 'BRANCH_LEQ', 'BRANCH_LEQ', 'LEAF', 'LEAF', 'BRANCH_LEQ', 'LEAF', - 'LEAF', 'BRANCH_LEQ', 'BRANCH_LEQ', 'LEAF', 'LEAF', 'LEAF'], - 'nodes_nodeids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], - 'nodes_treeids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - 'nodes_truenodeids': [1, 2, 3, 0, 0, 6, 0, 0, 9, 10, 0, 0, 0], - 'nodes_values': [[100, 101, 102], [55, 59, 65], [10, 11, 12], 0.0, 0.0, [1000, 1001, 1002], - 0.0, 0.0, [10000, 10001], [100000, 100001], 0.0, 0.0, 0.0], - 'post_transform': None} + "class_ids": [ + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + 0, + 1, + 2, + ], + "class_nodeids": [ + 3, + 3, + 3, + 4, + 4, + 4, + 6, + 6, + 6, + 7, + 7, + 7, + 10, + 10, + 10, + 11, + 11, + 11, + 12, + 12, + 12, + ], + "class_treeids": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + ], + "class_weights": [ + 0.4, + 0.6, + 0.0, + 0.95, + 0.04, + 0.0, + 0.0, + 0.185, + 0.814, + 0.372, + 0.62, + 0.0, + 0.74, + 0.21, + 0.03, + 0.0, + 1.0, + 0.0, + 0.87, + 0.05, + 0.071, + ], + "classlabels_int64s": [0, 1, 2], + "name": "TreeEnsembleClassifier", + "nodes_falsenodeids": [8, 5, 4, 0, 0, 7, 0, 0, 12, 11, 0, 0, 0], + "nodes_featureids": [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0], + "nodes_hitrates": [ + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + ], + "nodes_missing_value_tracks_true": [ + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + False, + ], + "nodes_modes": [ + "BRANCH_LEQ", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "BRANCH_LEQ", + "BRANCH_LEQ", + "LEAF", + "LEAF", + "LEAF", + ], + "nodes_nodeids": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], + "nodes_treeids": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + "nodes_truenodeids": [1, 2, 3, 0, 0, 6, 0, 0, 9, 10, 0, 0, 0], + "nodes_values": [ + [100, 101, 102], + [55, 59, 65], + [10, 11, 12], + 0.0, + 0.0, + [1000, 1001, 1002], + 0.0, + 0.0, + [10000, 10001], + [100000, 100001], + 0.0, + 0.0, + 0.0, + ], + "post_transform": None, + } root, nodes = Node.create(attrs) for node in root: key = node.nodes_treeids, node.nodes_nodeids @@ -207,53 +661,72 @@ def test_split_more_complex(self): ns[key] = node # back to onnx - new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s=[0, 1, 2], domain="ai.onnx.ml") + new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) for k in attrs: if k in {"post_transform"}: continue if len(attrs[k]) > len(new_attrs[k]): raise AssertionError( - f"Issue with attribute {k!r}\n{attrs['nodes_modes']}\nbefore {attrs[k]!r}" + f"Issue with attribute {k!r}\n{attrs['nodes_modes']}" + f"\nbefore {attrs[k]!r}" f"\nafter {new_attrs[k]!r}\n{new_attrs['nodes_modes']}" ) self.assertEqual(len(new_attrs["nodes_modes"]), len(attrs["nodes_modes"]) + 10) - node = make_node(op_type="TreeEnsembleClassifier", inputs=["X"], outputs=["L", "Y"], **new_attrs) - X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None]) - Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None, None]) - L = make_tensor_value_info('L', TensorProto.INT64, [None]) + node = make_node( + op_type="TreeEnsembleClassifier", + inputs=["X"], + outputs=["L", "Y"], + **new_attrs, + ) + X = make_tensor_value_info("X", TensorProto.FLOAT, [None, None]) + Y = make_tensor_value_info("Y", TensorProto.FLOAT, [None, None]) + L = make_tensor_value_info("L", TensorProto.INT64, [None]) graph = make_graph([node], "n", [X], [L, Y]) opset_imports = [make_opsetid("", 17), make_opsetid("ai.onnx.ml", 3)] model = make_model(graph, opset_imports=opset_imports) sess = InferenceSession(model.SerializeToString()) x = np.arange(20).reshape((-1, 2)).astype(np.float32) - got = sess.run(None, {'X': x}) + got = sess.run(None, {"X": x}) self.assertEqual(len(got), 2) self.assertEqual(got[0].tolist(), [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]) # again root, nodes = Node.create(new_attrs) root.unfold_rule_or() - new_new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s= [0, 1, 2], domain="ai.onnx.ml") + new_new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) self.assertEqual(new_attrs, new_new_attrs) def test_debug(self): try: from onnxmltools.debug1 import attrs - except ImportError as e: + except ImportError: return root, nodes = Node.create(attrs) root.unfold_rule_or() - new_attrs = root.to_attrs(post_transform=None, name="TreeEnsembleClassifier", - classlabels_int64s=[0, 1, 2], domain="ai.onnx.ml") + new_attrs = root.to_attrs( + post_transform=None, + name="TreeEnsembleClassifier", + classlabels_int64s=[0, 1, 2], + domain="ai.onnx.ml", + ) for k in attrs: if k in {"post_transform"}: continue if len(attrs[k]) > len(new_attrs[k]): raise AssertionError( - f"Issue with attribute {k!r}\n{len(new_attrs[k])}\nbefore {len(attrs[k])}." + f"Issue with attribute {k!r}\n{len(new_attrs[k])}" + f"\nbefore {len(attrs[k])}." ) diff --git a/tests/sparkml/test_vector_assembler.py b/tests/sparkml/test_vector_assembler.py index 7fdf229f..492f2c13 100644 --- a/tests/sparkml/test_vector_assembler.py +++ b/tests/sparkml/test_vector_assembler.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,30 +22,41 @@ class TestSparkmlVectorAssembler(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_model_vector_assembler(self): col_names = ["a", "b"] - model = VectorAssembler(inputCols=col_names, outputCol='features') - data = self.spark.createDataFrame([(1., Vectors.dense([1.0, 2.0, 3.0]))], col_names) - model_onnx = convert_sparkml(model, 'Sparkml VectorAssembler', [ - ('a', FloatTensorType([None, 1])), - ('b', FloatTensorType([None, 3])) - ], target_opset=TARGET_OPSET) + model = VectorAssembler(inputCols=col_names, outputCol="features") + data = self.spark.createDataFrame( + [(1.0, Vectors.dense([1.0, 2.0, 3.0]))], col_names + ) + model_onnx = convert_sparkml( + model, + "Sparkml VectorAssembler", + [("a", FloatTensorType([None, 1])), ("b", FloatTensorType([None, 3]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) self.assertTrue(model_onnx.graph.node is not None) # run the model predicted = model.transform(data) - expected = predicted.select("features").toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values + expected = ( + predicted.select("features") + .toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values + ) data_np = { - 'a': data.select('a').toPandas().values.astype(numpy.float32), - 'b': data.select('b').toPandas()['b'].apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32), + "a": data.select("a").toPandas().values.astype(numpy.float32), + "b": data.select("b") + .toPandas()["b"] + .apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32), } - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlVectorAssembler") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlVectorAssembler" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['features'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["features"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) assert output_shapes == [[None, 4]] diff --git a/tests/sparkml/test_vector_indexer.py b/tests/sparkml/test_vector_indexer.py index 9fd96ee8..4a95982f 100644 --- a/tests/sparkml/test_vector_indexer.py +++ b/tests/sparkml/test_vector_indexer.py @@ -9,11 +9,14 @@ from onnx.defs import onnx_opset_version from pyspark.ml.feature import VectorIndexer from pyspark.ml.linalg import Vectors -from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -21,63 +24,92 @@ class TestSparkmlVectorIndexer(SparkMlTestCase): - @unittest.skipIf( - True, reason=( - "discrepency, unfound values are replaced by -1 by ONNX and 0 " - "by spark.")) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), - 'Need Greater Opset 9') + True, + reason=( + "discrepency, unfound values are replaced by -1 by ONNX and 0 " "by spark." + ), + ) + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_model_vector_indexer_multi(self): vi = VectorIndexer(maxCategories=2, inputCol="a", outputCol="indexed") - data = self.spark.createDataFrame([ - (Vectors.dense([-1.0, 1.0, 3.1]),), - (Vectors.dense([0.0, 5.0, 6.2]),), - (Vectors.dense([0.0, 9.0, 3.1]),)], - ["a"] + data = self.spark.createDataFrame( + [ + (Vectors.dense([-1.0, 1.0, 3.1]),), + (Vectors.dense([0.0, 5.0, 6.2]),), + (Vectors.dense([0.0, 9.0, 3.1]),), + ], + ["a"], ) model = vi.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml VectorIndexer Multi', [ - ('a', FloatTensorType([None, model.numFeatures])) - ], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml VectorIndexer Multi", + [("a", FloatTensorType([None, model.numFeatures]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().indexed.apply(lambda x: pandas.Series(x.toArray())).values - data_np = data.toPandas().a.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlVectorIndexerMulti") + expected = ( + predicted.toPandas() + .indexed.apply(lambda x: pandas.Series(x.toArray())) + .values + ) + data_np = ( + data.toPandas() + .a.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlVectorIndexerMulti" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['indexed'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["indexed"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") - @unittest.skipIf(pv.Version(onnx.__version__) <= pv.Version('1.3'), - 'Need Greater Opset 9') + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") + @unittest.skipIf( + pv.Version(onnx.__version__) <= pv.Version("1.3"), "Need Greater Opset 9" + ) def test_model_vector_indexer_single(self): vi = VectorIndexer(maxCategories=3, inputCol="a", outputCol="indexed") - data = self.spark.createDataFrame([ - (Vectors.dense([-1.0]),), - (Vectors.dense([0.0]),), - (Vectors.dense([0.0]),)], - ["a"] + data = self.spark.createDataFrame( + [ + (Vectors.dense([-1.0]),), + (Vectors.dense([0.0]),), + (Vectors.dense([0.0]),), + ], + ["a"], ) model = vi.fit(data) - model_onnx = convert_sparkml(model, 'Sparkml VectorIndexer Single', [ - ('a', FloatTensorType([None, model.numFeatures])) - ], target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml VectorIndexer Single", + [("a", FloatTensorType([None, model.numFeatures]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().indexed.apply(lambda x: pandas.Series(x.toArray())).values - data_np = data.toPandas().a.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlVectorIndexerSingle") + expected = ( + predicted.toPandas() + .indexed.apply(lambda x: pandas.Series(x.toArray())) + .values + ) + data_np = ( + data.toPandas() + .a.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlVectorIndexerSingle" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['indexed'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["indexed"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_vector_slicer.py b/tests/sparkml/test_vector_slicer.py index 42d4d5e8..e8db4023 100644 --- a/tests/sparkml/test_vector_slicer.py +++ b/tests/sparkml/test_vector_slicer.py @@ -10,7 +10,11 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import FloatTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase @@ -18,29 +22,44 @@ class TestSparkmlVectorSlicer(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_vector_slicer(self): - data = self.spark.createDataFrame([ - (Vectors.dense([-2.0, 2.3, 0.0, 0.0, 1.0]), ), - (Vectors.dense([0.0, 0.0, 0.0, 0.0, 0.0]), ), - (Vectors.dense([0.6, -1.1, -3.0, 4.5, 3.3]), )], ["features"]) + data = self.spark.createDataFrame( + [ + (Vectors.dense([-2.0, 2.3, 0.0, 0.0, 1.0]),), + (Vectors.dense([0.0, 0.0, 0.0, 0.0, 0.0]),), + (Vectors.dense([0.6, -1.1, -3.0, 4.5, 3.3]),), + ], + ["features"], + ) model = VectorSlicer(inputCol="features", outputCol="sliced", indices=[1, 4]) feature_count = data.first()[0].array.size - model_onnx = convert_sparkml(model, 'Sparkml VectorSlicer', - [('features', FloatTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml VectorSlicer", + [("features", FloatTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data) - expected = predicted.toPandas().sliced.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - data_np = data.toPandas().features.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) - paths = save_data_models(data_np, expected, model, model_onnx, basename="SparkmlVectorSlicer") + expected = ( + predicted.toPandas() + .sliced.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + data_np = ( + data.toPandas() + .features.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlVectorSlicer" + ) onnx_model_path = paths[-1] - output, output_shapes = run_onnx_model(['sliced'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["sliced"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/sparkml/test_word2vec.py b/tests/sparkml/test_word2vec.py index 4adaba24..b47019ce 100644 --- a/tests/sparkml/test_word2vec.py +++ b/tests/sparkml/test_word2vec.py @@ -9,50 +9,69 @@ from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools import convert_sparkml from onnxmltools.convert.common.data_types import StringTensorType -from tests.sparkml.sparkml_test_utils import save_data_models, run_onnx_model, compare_results +from tests.sparkml.sparkml_test_utils import ( + save_data_models, + run_onnx_model, + compare_results, +) from tests.sparkml import SparkMlTestCase TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) -## For some reason during the spark bring up and shutdown something happens causing these +## For some reason during the spark bring up and shutdown +## something happens causing these ## tests to fail. For that you need to run each test here individually ## For now these will be commented out so as not to break the build -## AttributeError: 'NoneType' object has no attribute 'setCallSite' on model.surrogateDF -## Therefore we leave these tests out for now until a newere version of pyspark is availabe that address this issue +## AttributeError: 'NoneType' object has no attribute 'setCallSite' +## on model.surrogateDF +## Therefore we leave these tests out for now until a newere version +## of pyspark is availabe that address this issue class TestSparkmlWord2Vec(SparkMlTestCase): - - @unittest.skipIf(sys.version_info < (3, 8), - reason="pickle fails on python 3.7") + @unittest.skipIf(sys.version_info < (3, 8), reason="pickle fails on python 3.7") def test_word2vec(self): - data = self.spark.createDataFrame([ - ("Hi I heard about Spark".split(" "), ), - ("I wish Java could use case classes".split(" "), ), - ("Logistic regression models are neat".split(" "), ) - ], ["text"]) - word2Vec = Word2Vec(vectorSize=3, minCount=0, inputCol="text", outputCol="result") + data = self.spark.createDataFrame( + [ + ("Hi I heard about Spark".split(" "),), + ("I wish Java could use case classes".split(" "),), + ("Logistic regression models are neat".split(" "),), + ], + ["text"], + ) + word2Vec = Word2Vec( + vectorSize=3, minCount=0, inputCol="text", outputCol="result" + ) model = word2Vec.fit(data) vectors = model.getVectors() vectors.show(100, False) - + result = model.transform(data) result.show(100, False) - + # the input name should match that of inputCol feature_count = len(data.first()[0]) - model_onnx = convert_sparkml(model, 'Sparkml Word2Vec', [('text', StringTensorType([None, feature_count]))], - target_opset=TARGET_OPSET) + model_onnx = convert_sparkml( + model, + "Sparkml Word2Vec", + [("text", StringTensorType([None, feature_count]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(model_onnx is not None) # run the model predicted = model.transform(data.limit(1)) - expected = predicted.toPandas().result.apply(lambda x: pandas.Series(x.toArray())).values.astype(numpy.float32) + expected = ( + predicted.toPandas() + .result.apply(lambda x: pandas.Series(x.toArray())) + .values.astype(numpy.float32) + ) data_np = data.limit(1).toPandas().text.values - paths = save_data_models(data_np, expected, model, model_onnx, - basename="SparkmlWord2Vec") + paths = save_data_models( + data_np, expected, model, model_onnx, basename="SparkmlWord2Vec" + ) onnx_model_path = paths[-1] data_np = numpy.array(data_np[0]).reshape((1, -1)) - output, output_shapes = run_onnx_model(['result'], data_np, onnx_model_path) + output, output_shapes = run_onnx_model(["result"], data_np, onnx_model_path) compare_results(expected, output, decimal=5) diff --git a/tests/svmlib/test_SVMConverters.py b/tests/svmlib/test_SVMConverters.py index 49339529..32822cc2 100644 --- a/tests/svmlib/test_SVMConverters.py +++ b/tests/svmlib/test_SVMConverters.py @@ -5,25 +5,31 @@ """ import tempfile import numpy -import scipy import os + try: - from libsvm.svm import C_SVC as SVC, EPSILON_SVR as SVR, NU_SVC as NuSVC, NU_SVR as NuSVR - import libsvm.svm as svm + from libsvm.svm import ( + C_SVC as SVC, + EPSILON_SVR as SVR, + NU_SVC as NuSVC, + NU_SVR as NuSVR, + ) import libsvm.svmutil as svmutil + DISABLED = False except ImportError: try: - import svm - from svm import C_SVC as SVC, EPSILON_SVR as SVR, NU_SVC as NuSVC, NU_SVR as NuSVR + from svm import ( + C_SVC as SVC, + EPSILON_SVR as SVR, + NU_SVC as NuSVC, + NU_SVR as NuSVR, + ) import svmutil except ImportError: DISABLED = True -import onnxruntime -import numpy as np import unittest -import packaging.version as pv from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from sklearn.datasets import load_iris @@ -32,6 +38,7 @@ try: from libsvm.svm import PRINT_STRING_FUN, print_null + noprint = PRINT_STRING_FUN(print_null) except ImportError: # This was recently added. @@ -40,12 +47,14 @@ if not DISABLED: try: from scipy import ctypeslib + + DISABLED = ctypeslib is None except ImportError: DISABLED = True if not DISABLED: from onnxmltools.convert.libsvm import convert - + TARGET_OPSET = min(DEFAULT_OPSET_NUMBER, onnx_opset_version()) @@ -55,7 +64,7 @@ def __init__(self, model): self.model = model def predict(self, X, options=""): - if hasattr(X, 'shape'): + if hasattr(X, "shape"): X = X.tolist() res = svmutil.svm_predict([0 for i in X], list(X), self.model, options=options) return res @@ -64,13 +73,13 @@ def __getstate__(self): f = tempfile.NamedTemporaryFile(delete=False) svmutil.svm_save_model(f.name, self.model) with open(f.name, "rb") as h: - return {'data': h.read()} + return {"data": h.read()} os.remove(f) def __setstate__(self, data): f = tempfile.NamedTemporaryFile(delete=False) with open(f.name, "wb") as h: - h.write(data['by']) + h.write(data["by"]) self.model = svmutil.svm_load_model(f.name) os.remove(f) @@ -83,7 +92,6 @@ def predict(self, X): class SkAPIClProba2(SkAPI): - def predict(self, X): res = SkAPI.predict(self, X) ret = numpy.array(res[0]).ravel() @@ -97,7 +105,6 @@ def predict_proba(self, X): class SkAPICl(SkAPI): - def predict(self, X): res = SkAPI.predict(self, X) ret = numpy.array(res[0]).ravel() @@ -110,7 +117,6 @@ def decision_function(self, X): if pro.shape[1] == 1: pro = numpy.hstack([-pro, pro]) elif pro.shape[1] > 2: - # see from sklearn.utils.multiclass import _ovr_decision_function if False: conf = pro.copy() @@ -126,9 +132,7 @@ def decision_function(self, X): class TestSvmLibSVM(unittest.TestCase): - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmc_linear(self): iris = load_iris() @@ -148,13 +152,21 @@ def test_convert_svmc_linear(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmSvmcLinear", [('input', FloatTensorType())], target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmSvmcLinear", + [("input", FloatTensorType())], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPIClProba2(libsvm_model), node, - basename="LibSvmSvmcLinear-Dec2") - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPIClProba2(libsvm_model), + node, + basename="LibSvmSvmcLinear-Dec2", + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmc(self): iris = load_iris() @@ -174,13 +186,21 @@ def test_convert_svmc(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmSvmc", [('input', FloatTensorType())], target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmSvmc", + [("input", FloatTensorType())], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPIClProba2(libsvm_model), node, - basename="LibSvmSvmc-Dec2") - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPIClProba2(libsvm_model), + node, + basename="LibSvmSvmc-Dec2", + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmr_linear(self): iris = load_iris() @@ -197,13 +217,21 @@ def test_convert_svmr_linear(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmSvmrLinear", [('input', FloatTensorType())], target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmSvmrLinear", + [("input", FloatTensorType())], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPIReg(libsvm_model), node, - basename="LibSvmSvmrLinear-Dec3") - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPIReg(libsvm_model), + node, + basename="LibSvmSvmrLinear-Dec3", + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmr(self): iris = load_iris() @@ -221,13 +249,21 @@ def test_convert_svmr(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmSvmr", [('input', FloatTensorType())], target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmSvmr", + [("input", FloatTensorType())], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPIReg(libsvm_model), node, - basename="LibSvmSvmr") - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPIReg(libsvm_model), + node, + basename="LibSvmSvmr", + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_nusvmr(self): iris = load_iris() @@ -245,13 +281,21 @@ def test_convert_nusvmr(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmNuSvmr", [('input', FloatTensorType())], target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmNuSvmr", + [("input", FloatTensorType())], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPIReg(libsvm_model), node, - basename="LibSvmNuSvmr") - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPIReg(libsvm_model), + node, + basename="LibSvmNuSvmr", + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_nusvmc(self): iris = load_iris() @@ -271,14 +315,21 @@ def test_convert_nusvmc(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmNuSvmc", [('input', FloatTensorType(shape=['None', 'None']))], - target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmNuSvmc", + [("input", FloatTensorType(shape=["None", "None"]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPIClProba2(libsvm_model), node, - basename="LibSvmNuSvmc-Dec2") - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPIClProba2(libsvm_model), + node, + basename="LibSvmNuSvmc-Dec2", + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmc_linear_raw(self): iris = load_iris() @@ -298,15 +349,23 @@ def test_convert_svmc_linear_raw(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmSvmcLinearRaw", [('input', FloatTensorType(shape=['None', 'None']))], - target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmSvmcLinearRaw", + [("input", FloatTensorType(shape=["None", "None"]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) # known svm runtime dimension error in ONNX Runtime - dump_data_and_model(X[:5].astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmSvmcLinearRaw-Dec3", verbose=False) - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPICl(libsvm_model), + node, + basename="LibSvmSvmcLinearRaw-Dec3", + verbose=False, + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmc_raw(self): iris = load_iris() @@ -327,15 +386,22 @@ def test_convert_svmc_raw(self): libsvm_model = svmutil.svm_train(prob, param) # known svm runtime dimension error in ONNX Runtime - node = convert(libsvm_model, "LibSvmSvmcRaw", [('input', FloatTensorType(shape=['None', 'None']))], - target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmSvmcRaw", + [("input", FloatTensorType(shape=["None", "None"]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) - dump_data_and_model(X[:5].astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmSvmcRaw") + dump_data_and_model( + X[:5].astype(numpy.float32), + SkAPICl(libsvm_model), + node, + basename="LibSvmSvmcRaw", + ) @unittest.skip(reason="libsvm crashes.") - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_nusvmc_linear_raw(self): iris = load_iris() @@ -355,15 +421,23 @@ def test_convert_nusvmc_linear_raw(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmNuSvmcRaw", [('input', FloatTensorType(shape=['None', 'None']))], - target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmNuSvmcRaw", + [("input", FloatTensorType(shape=["None", "None"]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) X2 = numpy.vstack([X[:5], X[60:65]]) # 5x0, 5x1 - dump_data_and_model(X2.astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmNuSvmcRaw", verbose=False) - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X2.astype(numpy.float32), + SkAPICl(libsvm_model), + node, + basename="LibSvmNuSvmcRaw", + verbose=False, + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmc_rbf_raw_multi(self): iris = load_iris() @@ -383,15 +457,23 @@ def test_convert_svmc_rbf_raw_multi(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmNuSvmcMultiRbfRaw", [('input', FloatTensorType(shape=['None', 'None']))], - target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmNuSvmcMultiRbfRaw", + [("input", FloatTensorType(shape=["None", "None"]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) X2 = numpy.vstack([X[:2], X[60:62], X[110:112], X[147:149]]) # 5x0, 5x1 - dump_data_and_model(X2.astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmNuSvmcRaw", verbose=False) - - @unittest.skipIf(DISABLED, - reason="svmlib not really maintained") + dump_data_and_model( + X2.astype(numpy.float32), + SkAPICl(libsvm_model), + node, + basename="LibSvmNuSvmcRaw", + verbose=False, + ) + + @unittest.skipIf(DISABLED, reason="svmlib not really maintained") def test_convert_svmc_linear_raw_multi(self): iris = load_iris() @@ -411,12 +493,21 @@ def test_convert_svmc_linear_raw_multi(self): libsvm_model = svmutil.svm_train(prob, param) - node = convert(libsvm_model, "LibSvmNuSvmcMultiRaw", [('input', FloatTensorType(shape=['None', 2]))], - target_opset=TARGET_OPSET) + node = convert( + libsvm_model, + "LibSvmNuSvmcMultiRaw", + [("input", FloatTensorType(shape=["None", 2]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(node is not None) X2 = numpy.vstack([X[:2], X[60:62], X[110:112], X[147:149]]) # 5x0, 5x1 - dump_data_and_model(X2.astype(numpy.float32), SkAPICl(libsvm_model), node, - basename="LibSvmSvmcRaw-Dec3", verbose=False) + dump_data_and_model( + X2.astype(numpy.float32), + SkAPICl(libsvm_model), + node, + basename="LibSvmSvmcRaw-Dec3", + verbose=False, + ) if __name__ == "__main__": diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index d5c3fee8..63833c57 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -5,7 +5,6 @@ """ import os import unittest -import warnings from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER import onnxmltools @@ -17,23 +16,28 @@ class TestUtils(unittest.TestCase): - @staticmethod def _parseEOL(file): - with open(file, 'r') as f: + with open(file, "r") as f: content = f.read() return content.replace("\r\n", "\n") def test_load_model(self): this = os.path.dirname(__file__) - onnx_file = os.path.join(this, "models", "coreml_OneHotEncoder_BikeSharing.onnx") + onnx_file = os.path.join( + this, "models", "coreml_OneHotEncoder_BikeSharing.onnx" + ) onnx_model = load_model(onnx_file) self.assertTrue(onnx_model is not None) def test_save_model(self): this = os.path.dirname(__file__) - onnx_file = os.path.join(this, "models", "coreml_OneHotEncoder_BikeSharing.onnx") - new_onnx_file = os.path.join(this, "models", "coreml_OneHotEncoder_BikeSharing2.onnx") + onnx_file = os.path.join( + this, "models", "coreml_OneHotEncoder_BikeSharing.onnx" + ) + new_onnx_file = os.path.join( + this, "models", "coreml_OneHotEncoder_BikeSharing2.onnx" + ) onnx_model = load_model(onnx_file) save_model(onnx_model, new_onnx_file) @@ -42,7 +46,9 @@ def test_save_model(self): def test_model_setters(self): this = os.path.dirname(__file__) - onnx_file = os.path.join(this, "models", "coreml_OneHotEncoder_BikeSharing.onnx") + onnx_file = os.path.join( + this, "models", "coreml_OneHotEncoder_BikeSharing.onnx" + ) onnx_model = load_model(onnx_file) set_model_version(onnx_model, 2) set_model_domain(onnx_model, "com.sample") @@ -54,22 +60,30 @@ def test_model_setters(self): def test_set_docstring_blank(self): this = os.path.dirname(__file__) - onnx_file = os.path.join(this, "models", "coreml_OneHotEncoder_BikeSharing.onnx") + onnx_file = os.path.join( + this, "models", "coreml_OneHotEncoder_BikeSharing.onnx" + ) onnx_model = load_model(onnx_file) set_model_doc_string(onnx_model, "sample") - self.assertRaises(ValueError, set_model_doc_string, onnx_model.doc_string, "sample") + self.assertRaises( + ValueError, set_model_doc_string, onnx_model.doc_string, "sample" + ) set_model_doc_string(onnx_model, "", True) self.assertEqual(onnx_model.doc_string, "") class TestWrapper(unittest.TestCase): - - @unittest.skipIf(True, reason="Needs this PR: https://github.com/onnx/tensorflow-onnx/pull/1563") + @unittest.skipIf( + True, reason="Needs this PR: https://github.com/onnx/tensorflow-onnx/pull/1563" + ) def test_keras_with_tf2onnx(self): import tensorflow.keras as keras + model = keras.Sequential() - model.add(keras.layers.Dense(units=4, input_shape=(10,), activation='relu')) - model.compile(loss='binary_crossentropy', optimizer='Adam', metrics=['binary_accuracy']) + model.add(keras.layers.Dense(units=4, input_shape=(10,), activation="relu")) + model.compile( + loss="binary_crossentropy", optimizer="Adam", metrics=["binary_accuracy"] + ) onnx_model = onnxmltools.convert_tensorflow(model, target_opset=TARGET_OPSET) self.assertTrue(len(onnx_model.graph.node) > 0) diff --git a/tests/xgboost/test_xgboost_13.py b/tests/xgboost/test_xgboost_13.py index ac877a57..18d31bab 100644 --- a/tests/xgboost/test_xgboost_13.py +++ b/tests/xgboost/test_xgboost_13.py @@ -9,7 +9,7 @@ from numpy.testing import assert_almost_equal import pandas from sklearn.model_selection import train_test_split -from xgboost import XGBRegressor, XGBClassifier, train, DMatrix +from xgboost import XGBClassifier from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER from onnxmltools.convert import convert_xgboost @@ -21,29 +21,41 @@ class TestXGBoost13(unittest.TestCase): - def test_xgb_regressor(self): this = os.path.dirname(__file__) df = pandas.read_csv(os.path.join(this, "data_fail_empty.csv")) - X, y = df.drop('y', axis=1), df['y'] + X, y = df.drop("y", axis=1), df["y"] X_train, X_test, y_train, y_test = train_test_split(X, y) clr = XGBClassifier( - max_delta_step= 0, tree_method='hist', n_estimators=100, - booster='gbtree', objective='binary:logistic', eval_metric='logloss', - learning_rate= 0.1, gamma=10, max_depth=7, min_child_weight=50, - subsample=0.75, colsample_bytree=0.75, random_state=42, - verbosity=0) + max_delta_step=0, + tree_method="hist", + n_estimators=100, + booster="gbtree", + objective="binary:logistic", + eval_metric="logloss", + learning_rate=0.1, + gamma=10, + max_depth=7, + min_child_weight=50, + subsample=0.75, + colsample_bytree=0.75, + random_state=42, + verbosity=0, + ) - clr.fit(X_train, y_train, eval_set=[(X_test, y_test)], - early_stopping_rounds=40) + clr.fit(X_train, y_train, eval_set=[(X_test, y_test)], early_stopping_rounds=40) - initial_type = [('float_input', FloatTensorType([None, 797]))] - onx = convert_xgboost(clr, initial_types=initial_type, target_opset=TARGET_OPSET) + initial_type = [("float_input", FloatTensorType([None, 797]))] + onx = convert_xgboost( + clr, initial_types=initial_type, target_opset=TARGET_OPSET + ) expected = clr.predict(X_test), clr.predict_proba(X_test) - sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"]) + sess = InferenceSession( + onx.SerializeToString(), providers=["CPUExecutionProvider"] + ) X_test = X_test.values.astype(np.float32) - got = sess.run(None, {'float_input': X_test}) + got = sess.run(None, {"float_input": X_test}) assert_almost_equal(expected[1], got[1]) assert_almost_equal(expected[0], got[0]) diff --git a/tests/xgboost/test_xgboost_converters.py b/tests/xgboost/test_xgboost_converters.py index 9845c1e0..0137ffa6 100644 --- a/tests/xgboost/test_xgboost_converters.py +++ b/tests/xgboost/test_xgboost_converters.py @@ -9,9 +9,21 @@ from numpy.testing import assert_almost_equal import pandas from sklearn.datasets import ( - load_diabetes, load_iris, make_classification, load_digits, make_regression) + load_diabetes, + load_iris, + make_classification, + load_digits, + make_regression, +) from sklearn.model_selection import train_test_split -from xgboost import XGBRegressor, XGBClassifier, train, DMatrix, Booster, train as train_xgb +from xgboost import ( + XGBRegressor, + XGBClassifier, + train, + DMatrix, + Booster, + train as train_xgb, +) from sklearn.preprocessing import StandardScaler from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER @@ -39,12 +51,15 @@ def fct_id(y): def _fit_classification_model(model, n_classes, is_str=False, dtype=None): - x, y = make_classification(n_classes=n_classes, n_features=100, - n_samples=1000, - random_state=42, n_informative=7) + x, y = make_classification( + n_classes=n_classes, + n_features=100, + n_samples=1000, + random_state=42, + n_informative=7, + ) y = y.astype(np.str_) if is_str else y.astype(np.int64) - x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, - random_state=42) + x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, random_state=42) if dtype is not None: y_train = y_train.astype(dtype) model.fit(x_train, y_train) @@ -52,162 +67,222 @@ def _fit_classification_model(model, n_classes, is_str=False, dtype=None): class TestXGBoostModels(unittest.TestCase): - def test_xgb_regressor(self): iris = load_diabetes() x = iris.data y = iris.target - x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, - random_state=42) + x_train, x_test, y_train, _ = train_test_split( + x, y, test_size=0.5, random_state=42 + ) xgb = XGBRegressor() xgb.fit(x_train, y_train) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) dump_data_and_model( x_test.astype("float32"), - xgb, conv_model, - basename="SklearnXGBRegressor-Dec3") + xgb, + conv_model, + basename="SklearnXGBRegressor-Dec3", + ) def test_xgb_classifier(self): xgb, x_test = _fit_classification_model(XGBClassifier(), 2) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) - dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifier") + dump_data_and_model(x_test, xgb, conv_model, basename="SklearnXGBClassifier") def test_xgb_classifier_uint8(self): - xgb, x_test = _fit_classification_model( - XGBClassifier(), 2, dtype=np.uint8) + xgb, x_test = _fit_classification_model(XGBClassifier(), 2, dtype=np.uint8) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=['None', 'None']))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=["None", "None"]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) - dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifier") + dump_data_and_model(x_test, xgb, conv_model, basename="SklearnXGBClassifier") def test_xgb_classifier_multi(self): xgb, x_test = _fit_classification_model(XGBClassifier(), 3) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifierMulti") + x_test, xgb, conv_model, basename="SklearnXGBClassifierMulti" + ) def test_xgb_classifier_multi_reglog(self): xgb, x_test = _fit_classification_model( - XGBClassifier(objective='reg:logistic'), 4) + XGBClassifier(objective="reg:logistic"), 4 + ) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifierMultiRegLog") + x_test, xgb, conv_model, basename="SklearnXGBClassifierMultiRegLog" + ) def test_xgb_classifier_reglog(self): xgb, x_test = _fit_classification_model( - XGBClassifier(objective='reg:logistic'), 2) + XGBClassifier(objective="reg:logistic"), 2 + ) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifierRegLog") + x_test, xgb, conv_model, basename="SklearnXGBClassifierRegLog" + ) def test_xgb_classifier_multi_discrete_int_labels(self): iris = load_iris() x = iris.data[:, :2] y = iris.target - x_train, x_test, y_train, _ = train_test_split(x, - y, - test_size=0.5, - random_state=42) + x_train, x_test, y_train, _ = train_test_split( + x, y, test_size=0.5, random_state=42 + ) xgb = XGBClassifier(n_estimators=3) xgb.fit(x_train, y_train) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) dump_data_and_model( x_test.astype("float32"), - xgb, conv_model, - basename="SklearnXGBClassifierMultiDiscreteIntLabels") + xgb, + conv_model, + basename="SklearnXGBClassifierMultiDiscreteIntLabels", + ) def test_xgboost_booster_classifier_bin(self): - x, y = make_classification(n_classes=2, n_features=5, - n_samples=100, - random_state=42, n_informative=3) - x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, - random_state=42) + x, y = make_classification( + n_classes=2, n_features=5, n_samples=100, random_state=42, n_informative=3 + ) + x_train, x_test, y_train, _ = train_test_split( + x, y, test_size=0.5, random_state=42 + ) data = DMatrix(x_train, label=y_train) - model = train({'objective': 'binary:logistic', - 'n_estimators': 3, 'min_child_samples': 1}, data) - model_onnx = convert_xgboost(model, 'tree-based classifier', - [('input', FloatTensorType([None, x.shape[1]]))], - target_opset=TARGET_OPSET) - dump_data_and_model(x_test.astype(np.float32), - model, model_onnx, basename="XGBBoosterMCl") + model = train( + {"objective": "binary:logistic", "n_estimators": 3, "min_child_samples": 1}, + data, + ) + model_onnx = convert_xgboost( + model, + "tree-based classifier", + [("input", FloatTensorType([None, x.shape[1]]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + x_test.astype(np.float32), model, model_onnx, basename="XGBBoosterMCl" + ) def test_xgboost_booster_classifier_multiclass_softprob(self): - x, y = make_classification(n_classes=3, n_features=5, - n_samples=100, - random_state=42, n_informative=3) - x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, - random_state=42) + x, y = make_classification( + n_classes=3, n_features=5, n_samples=100, random_state=42, n_informative=3 + ) + x_train, x_test, y_train, _ = train_test_split( + x, y, test_size=0.5, random_state=42 + ) data = DMatrix(x_train, label=y_train) - model = train({'objective': 'multi:softprob', - 'n_estimators': 3, 'min_child_samples': 1, - 'num_class': 3}, data) - model_onnx = convert_xgboost(model, 'tree-based classifier', - [('input', FloatTensorType([None, x.shape[1]]))], - target_opset=TARGET_OPSET) - dump_data_and_model(x_test.astype(np.float32), - model, model_onnx, basename="XGBBoosterMClSoftProb") + model = train( + { + "objective": "multi:softprob", + "n_estimators": 3, + "min_child_samples": 1, + "num_class": 3, + }, + data, + ) + model_onnx = convert_xgboost( + model, + "tree-based classifier", + [("input", FloatTensorType([None, x.shape[1]]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + x_test.astype(np.float32), + model, + model_onnx, + basename="XGBBoosterMClSoftProb", + ) def test_xgboost_booster_classifier_multiclass_softmax(self): - x, y = make_classification(n_classes=3, n_features=5, - n_samples=100, - random_state=42, n_informative=3) - x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, - random_state=42) + x, y = make_classification( + n_classes=3, n_features=5, n_samples=100, random_state=42, n_informative=3 + ) + x_train, x_test, y_train, _ = train_test_split( + x, y, test_size=0.5, random_state=42 + ) data = DMatrix(x_train, label=y_train) - model = train({'objective': 'multi:softmax', - 'n_estimators': 3, 'min_child_samples': 1, - 'num_class': 3}, data) - model_onnx = convert_xgboost(model, 'tree-based classifier', - [('input', FloatTensorType([None, x.shape[1]]))], - target_opset=TARGET_OPSET) - dump_data_and_model(x_test.astype(np.float32), - model, model_onnx, basename="XGBBoosterMClSoftMax") + model = train( + { + "objective": "multi:softmax", + "n_estimators": 3, + "min_child_samples": 1, + "num_class": 3, + }, + data, + ) + model_onnx = convert_xgboost( + model, + "tree-based classifier", + [("input", FloatTensorType([None, x.shape[1]]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + x_test.astype(np.float32), + model, + model_onnx, + basename="XGBBoosterMClSoftMax", + ) def test_xgboost_booster_classifier_reg(self): - x, y = make_classification(n_classes=2, n_features=5, - n_samples=100, - random_state=42, n_informative=3) + x, y = make_classification( + n_classes=2, n_features=5, n_samples=100, random_state=42, n_informative=3 + ) y = y.astype(np.float32) + 0.567 - x_train, x_test, y_train, _ = train_test_split(x, y, test_size=0.5, - random_state=42) + x_train, x_test, y_train, _ = train_test_split( + x, y, test_size=0.5, random_state=42 + ) data = DMatrix(x_train, label=y_train) - model = train({'objective': 'reg:squarederror', - 'n_estimators': 3, 'min_child_samples': 1}, data) - model_onnx = convert_xgboost(model, 'tree-based classifier', - [('input', FloatTensorType([None, x.shape[1]]))], - target_opset=TARGET_OPSET) - dump_data_and_model(x_test.astype(np.float32), - model, model_onnx, basename="XGBBoosterReg") + model = train( + { + "objective": "reg:squarederror", + "n_estimators": 3, + "min_child_samples": 1, + }, + data, + ) + model_onnx = convert_xgboost( + model, + "tree-based classifier", + [("input", FloatTensorType([None, x.shape[1]]))], + target_opset=TARGET_OPSET, + ) + dump_data_and_model( + x_test.astype(np.float32), model, model_onnx, basename="XGBBoosterReg" + ) def test_xgboost_10(self): this = os.path.abspath(os.path.dirname(__file__)) @@ -217,29 +292,42 @@ def test_xgboost_10(self): param_distributions = { "colsample_bytree": 0.5, "gamma": 0.2, - 'learning_rate': 0.3, - 'max_depth': 2, - 'min_child_weight': 1., - 'n_estimators': 1, - 'missing': np.nan, + "learning_rate": 0.3, + "max_depth": 2, + "min_child_weight": 1.0, + "n_estimators": 1, + "missing": np.nan, } train_df = pandas.read_csv(train) - X_train, y_train = train_df.drop('label', axis=1).values, train_df['label'].fillna(0).values + X_train, y_train = ( + train_df.drop("label", axis=1).values, + train_df["label"].fillna(0).values, + ) test_df = pandas.read_csv(test) - X_test, y_test = test_df.drop('label', axis=1).values, test_df['label'].fillna(0).values - - regressor = XGBRegressor(verbose=0, objective='reg:squarederror', **param_distributions) + X_test, _ = ( + test_df.drop("label", axis=1).values, + test_df["label"].fillna(0).values, + ) + + regressor = XGBRegressor( + verbose=0, objective="reg:squarederror", **param_distributions + ) regressor.fit(X_train, y_train) model_onnx = convert_xgboost( - regressor, 'bug', - [('input', FloatTensorType([None, X_train.shape[1]]))], - target_opset=TARGET_OPSET) + regressor, + "bug", + [("input", FloatTensorType([None, X_train.shape[1]]))], + target_opset=TARGET_OPSET, + ) dump_data_and_model( - X_test.astype(np.float32), regressor, model_onnx, - basename="XGBBoosterRegBug") + X_test.astype(np.float32), + regressor, + model_onnx, + basename="XGBBoosterRegBug", + ) def test_xgboost_classifier_i5450_softmax(self): iris = load_iris() @@ -247,19 +335,25 @@ def test_xgboost_classifier_i5450_softmax(self): X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=10) clr = XGBClassifier(objective="multi:softmax", max_depth=1, n_estimators=2) clr.fit(X_train, y_train, eval_set=[(X_test, y_test)], early_stopping_rounds=40) - initial_type = [('float_input', FloatTensorType([None, 4]))] - onx = convert_xgboost(clr, initial_types=initial_type, target_opset=TARGET_OPSET) - sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"]) - input_name = sess.get_inputs()[0].name - label_name = sess.get_outputs()[1].name - predict_list = [1., 20., 466., 0.] - predict_array = np.array(predict_list).reshape((1,-1)).astype(np.float32) - pred_onx = sess.run([label_name], {input_name: predict_array})[0] + initial_type = [("float_input", FloatTensorType([None, 4]))] + onx = convert_xgboost( + clr, initial_types=initial_type, target_opset=TARGET_OPSET + ) + sess = InferenceSession( + onx.SerializeToString(), providers=["CPUExecutionProvider"] + ) + sess.get_inputs()[0].name + sess.get_outputs()[1].name + predict_list = [1.0, 20.0, 466.0, 0.0] + np.array(predict_list).reshape((1, -1)).astype(np.float32) bst = clr.get_booster() - bst.dump_model('dump.raw.txt') + bst.dump_model("dump.raw.txt") dump_data_and_model( - X_test.astype(np.float32) + 1e-5, clr, onx, - basename="XGBClassifierIris-Out0") + X_test.astype(np.float32) + 1e-5, + clr, + onx, + basename="XGBClassifierIris-Out0", + ) def test_xgboost_classifier_i5450(self): iris = load_iris() @@ -267,20 +361,22 @@ def test_xgboost_classifier_i5450(self): X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=10) clr = XGBClassifier(objective="multi:softprob", max_depth=1, n_estimators=2) clr.fit(X_train, y_train, eval_set=[(X_test, y_test)], early_stopping_rounds=40) - initial_type = [('float_input', FloatTensorType([None, 4]))] - onx = convert_xgboost(clr, initial_types=initial_type, target_opset=TARGET_OPSET) - sess = InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"]) - input_name = sess.get_inputs()[0].name - label_name = sess.get_outputs()[1].name - predict_list = [1., 20., 466., 0.] - predict_array = np.array(predict_list).reshape((1,-1)).astype(np.float32) - pred_onx = sess.run([label_name], {input_name: predict_array})[0] - pred_xgboost = sessresults=clr.predict_proba(predict_array) + initial_type = [("float_input", FloatTensorType([None, 4]))] + onx = convert_xgboost( + clr, initial_types=initial_type, target_opset=TARGET_OPSET + ) + sess = InferenceSession( + onx.SerializeToString(), providers=["CPUExecutionProvider"] + ) + sess.get_inputs()[0].name + sess.get_outputs()[1].name + predict_list = [1.0, 20.0, 466.0, 0.0] + np.array(predict_list).reshape((1, -1)).astype(np.float32) bst = clr.get_booster() - bst.dump_model('dump.raw.txt') + bst.dump_model("dump.raw.txt") dump_data_and_model( - X_test.astype(np.float32) + 1e-5, clr, onx, - basename="XGBClassifierIris") + X_test.astype(np.float32) + 1e-5, clr, onx, basename="XGBClassifierIris" + ) def test_xgboost_example_mnist(self): """ @@ -299,12 +395,14 @@ def test_xgboost_example_mnist(self): sh = [None, X_train.shape[1]] onnx_model = convert_xgboost( - clf, initial_types=[('input', FloatTensorType(sh))], - target_opset=TARGET_OPSET) + clf, + initial_types=[("input", FloatTensorType(sh))], + target_opset=TARGET_OPSET, + ) dump_data_and_model( - X_test.astype(np.float32), clf, onnx_model, - basename="XGBoostExample") + X_test.astype(np.float32), clf, onnx_model, basename="XGBoostExample" + ) def test_xgb_empty_tree(self): xgb = XGBClassifier(n_estimators=2, max_depth=2) @@ -315,91 +413,129 @@ def test_xgb_empty_tree(self): y = [0, 1, 0] xgb.fit(X, y) conv_model = convert_xgboost( - xgb, initial_types=[ - ('input', FloatTensorType(shape=[None, X.shape[1]]))], - target_opset=TARGET_OPSET) - sess = InferenceSession(conv_model.SerializeToString(), providers=["CPUExecutionProvider"]) - res = sess.run(None, {'input': X.astype(np.float32)}) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, X.shape[1]]))], + target_opset=TARGET_OPSET, + ) + sess = InferenceSession( + conv_model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + res = sess.run(None, {"input": X.astype(np.float32)}) assert_almost_equal(xgb.predict_proba(X), res[1]) assert_almost_equal(xgb.predict(X), res[0]) def test_xgb_best_tree_limit(self): - # Train iris = load_iris() X, y = iris.data, iris.target X_train, X_test, y_train, y_test = train_test_split(X, y) dtrain = DMatrix(X_train, label=y_train) dtest = DMatrix(X_test) - param = {'objective': 'multi:softmax', 'num_class': 3} + param = {"objective": "multi:softmax", "num_class": 3} bst_original = train_xgb(param, dtrain, 10) - initial_type = [('float_input', FloatTensorType([None, 4]))] - bst_original.save_model('model.json') + initial_type = [("float_input", FloatTensorType([None, 4]))] + bst_original.save_model("model.json") onx_loaded = convert_xgboost( - bst_original, initial_types=initial_type, - target_opset=TARGET_OPSET) - sess = InferenceSession(onx_loaded.SerializeToString(), providers=["CPUExecutionProvider"]) - res = sess.run(None, {'float_input': X_test.astype(np.float32)}) - assert_almost_equal(bst_original.predict(dtest, output_margin=True), res[1], decimal=5) + bst_original, initial_types=initial_type, target_opset=TARGET_OPSET + ) + sess = InferenceSession( + onx_loaded.SerializeToString(), providers=["CPUExecutionProvider"] + ) + res = sess.run(None, {"float_input": X_test.astype(np.float32)}) + assert_almost_equal( + bst_original.predict(dtest, output_margin=True), res[1], decimal=5 + ) assert_almost_equal(bst_original.predict(dtest), res[0]) # After being restored, the loaded booster is not exactly the same # in memory. `best_ntree_limit` is not saved during `save_model`. bst_loaded = Booster() - bst_loaded.load_model('model.json') - bst_loaded.save_model('model2.json') - assert_almost_equal(bst_loaded.predict(dtest, output_margin=True), - bst_original.predict(dtest, output_margin=True), decimal=5) + bst_loaded.load_model("model.json") + bst_loaded.save_model("model2.json") + assert_almost_equal( + bst_loaded.predict(dtest, output_margin=True), + bst_original.predict(dtest, output_margin=True), + decimal=5, + ) assert_almost_equal(bst_loaded.predict(dtest), bst_original.predict(dtest)) onx_loaded = convert_xgboost( - bst_loaded, initial_types=initial_type, - target_opset=TARGET_OPSET) - sess = InferenceSession(onx_loaded.SerializeToString(), providers=["CPUExecutionProvider"]) - res = sess.run(None, {'float_input': X_test.astype(np.float32)}) - assert_almost_equal(bst_loaded.predict(dtest, output_margin=True), res[1], decimal=5) + bst_loaded, initial_types=initial_type, target_opset=TARGET_OPSET + ) + sess = InferenceSession( + onx_loaded.SerializeToString(), providers=["CPUExecutionProvider"] + ) + res = sess.run(None, {"float_input": X_test.astype(np.float32)}) + assert_almost_equal( + bst_loaded.predict(dtest, output_margin=True), res[1], decimal=5 + ) assert_almost_equal(bst_loaded.predict(dtest), res[0]) def test_onnxrt_python_xgbclassifier(self): x = np.random.randn(100, 10).astype(np.float32) - y = ((x.sum(axis=1) + np.random.randn(x.shape[0]) / 50 + 0.5) >= 0).astype(np.int64) + y = ((x.sum(axis=1) + np.random.randn(x.shape[0]) / 50 + 0.5) >= 0).astype( + np.int64 + ) x_train, x_test, y_train, y_test = train_test_split(x, y) bmy = np.mean(y_train) - + for bm, n_est in [(None, 1), (None, 3), (bmy, 1), (bmy, 3)]: - model_skl = XGBClassifier(n_estimators=n_est, - learning_rate=0.01, - subsample=0.5, objective="binary:logistic", - base_score=bm, max_depth=2) + model_skl = XGBClassifier( + n_estimators=n_est, + learning_rate=0.01, + subsample=0.5, + objective="binary:logistic", + base_score=bm, + max_depth=2, + ) model_skl.fit(x_train, y_train, eval_set=[(x_test, y_test)], verbose=0) model_onnx_skl = convert_xgboost( - model_skl, initial_types=[('X', FloatTensorType([None, x.shape[1]]))], - target_opset=TARGET_OPSET) + model_skl, + initial_types=[("X", FloatTensorType([None, x.shape[1]]))], + target_opset=TARGET_OPSET, + ) with self.subTest(base_score=bm, n_estimators=n_est): - oinf = InferenceSession(model_onnx_skl.SerializeToString(), providers=["CPUExecutionProvider"]) - res2 = oinf.run(None, {'X': x_test}) + oinf = InferenceSession( + model_onnx_skl.SerializeToString(), + providers=["CPUExecutionProvider"], + ) + res2 = oinf.run(None, {"X": x_test}) assert_almost_equal(model_skl.predict_proba(x_test), res2[1]) def test_xgb_cost(self): obj_classes = { - 'reg:logistic': (XGBClassifier, fct_cl2, - make_classification(n_features=4, n_classes=2, - n_clusters_per_class=1)), - 'binary:logistic': (XGBClassifier, fct_cl2, - make_classification(n_features=4, n_classes=2, - n_clusters_per_class=1)), - 'multi:softmax': (XGBClassifier, fct_id, - make_classification(n_features=4, n_classes=3, - n_clusters_per_class=1)), - 'multi:softprob': (XGBClassifier, fct_id, - make_classification(n_features=4, n_classes=3, - n_clusters_per_class=1)), - 'reg:squarederror': (XGBRegressor, fct_id, - make_regression(n_features=4, n_targets=1)), - 'reg:squarederror2': (XGBRegressor, fct_id, - make_regression(n_features=4, n_targets=2)), + "reg:logistic": ( + XGBClassifier, + fct_cl2, + make_classification(n_features=4, n_classes=2, n_clusters_per_class=1), + ), + "binary:logistic": ( + XGBClassifier, + fct_cl2, + make_classification(n_features=4, n_classes=2, n_clusters_per_class=1), + ), + "multi:softmax": ( + XGBClassifier, + fct_id, + make_classification(n_features=4, n_classes=3, n_clusters_per_class=1), + ), + "multi:softprob": ( + XGBClassifier, + fct_id, + make_classification(n_features=4, n_classes=3, n_clusters_per_class=1), + ), + "reg:squarederror": ( + XGBRegressor, + fct_id, + make_regression(n_features=4, n_targets=1), + ), + "reg:squarederror2": ( + XGBRegressor, + fct_id, + make_regression(n_features=4, n_targets=2), + ), } nb_tests = 0 for objective in obj_classes: # pylint: disable=C0206 @@ -412,18 +548,18 @@ def test_xgb_cost(self): X, y = iris.data, iris.target y = fct(y) X_train, X_test, y_train, _ = train_test_split( - X, y, random_state=11) + X, y, random_state=11 + ) probs.append((X_train, X_test, y_train)) X_train, X_test, y_train, _ = train_test_split( - *prob, random_state=11) + *prob, random_state=11 + ) probs.append((X_train, X_test, y_train)) for X_train, X_test, y_train in probs: - obj = objective.replace( - 'reg:squarederror2', 'reg:squarederror') - obj = obj.replace( - 'multi:softmax2', 'multi:softmax') + obj = objective.replace("reg:squarederror2", "reg:squarederror") + obj = obj.replace("multi:softmax2", "multi:softmax") clr = cl(objective=obj, n_estimators=n_estimators) if len(y_train.shape) == 2: y_train = y_train[:, 1] @@ -431,21 +567,26 @@ def test_xgb_cost(self): clr.fit(X_train, y_train) except ValueError as e: raise AssertionError( - "Unable to train with objective %r and data %r." % ( - objective, y_train)) from e + "Unable to train with objective %r and data %r." + % (objective, y_train) + ) from e model_def = convert_xgboost( - clr, initial_types=[('X', FloatTensorType([None, X.shape[1]]))], - target_opset=TARGET_OPSET) - - oinf = InferenceSession(model_def.SerializeToString(), - providers=["CPUExecutionProvider"]) - y = oinf.run(None, {'X': X_test.astype(np.float32)}) + clr, + initial_types=[("X", FloatTensorType([None, X.shape[1]]))], + target_opset=TARGET_OPSET, + ) + + oinf = InferenceSession( + model_def.SerializeToString(), + providers=["CPUExecutionProvider"], + ) + y = oinf.run(None, {"X": X_test.astype(np.float32)}) if cl == XGBRegressor: exp = clr.predict(X_test) assert_almost_equal(exp, y[0].ravel(), decimal=5) else: - if 'softmax' not in obj: + if "softmax" not in obj: exp = clr.predict_proba(X_test) got = pandas.DataFrame(y[1]).values assert_almost_equal(exp, got, decimal=5) @@ -459,36 +600,53 @@ def test_xgb_cost(self): def test_xgb_classifier_601(self): model = XGBClassifier( - base_score=0.5, booster='gbtree', colsample_bylevel=1, - colsample_bynode=1, colsample_bytree=1, gamma=0, - importance_type='gain', interaction_constraints='', - learning_rate=0.3, max_delta_step=0, max_depth=6, - min_child_weight=1, missing=np.nan, - n_estimators=3, n_jobs=0, num_parallel_tree=1, - objective='multi:softprob', random_state=0, reg_alpha=0, - reg_lambda=1, scale_pos_weight=None, subsample=1, - tree_method='exact', validate_parameters=1) + base_score=0.5, + booster="gbtree", + colsample_bylevel=1, + colsample_bynode=1, + colsample_bytree=1, + gamma=0, + importance_type="gain", + interaction_constraints="", + learning_rate=0.3, + max_delta_step=0, + max_depth=6, + min_child_weight=1, + missing=np.nan, + n_estimators=3, + n_jobs=0, + num_parallel_tree=1, + objective="multi:softprob", + random_state=0, + reg_alpha=0, + reg_lambda=1, + scale_pos_weight=None, + subsample=1, + tree_method="exact", + validate_parameters=1, + ) xgb, x_test = _fit_classification_model(model, 3) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) self.assertTrue(conv_model is not None) - dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifier601") + dump_data_and_model(x_test, xgb, conv_model, basename="SklearnXGBClassifier601") def test_xgb_classifier_hinge(self): model = XGBClassifier( - n_estimators=3, objective='binary:hinge', random_state=0, - max_depth=2) + n_estimators=3, objective="binary:hinge", random_state=0, max_depth=2 + ) xgb, x_test = _fit_classification_model(model, 2) conv_model = convert_xgboost( - xgb, initial_types=[('input', FloatTensorType(shape=[None, None]))], - target_opset=TARGET_OPSET) + xgb, + initial_types=[("input", FloatTensorType(shape=[None, None]))], + target_opset=TARGET_OPSET, + ) dump_data_and_model( - x_test, xgb, conv_model, - basename="SklearnXGBClassifierHinge") - + x_test, xgb, conv_model, basename="SklearnXGBClassifierHinge" + ) if __name__ == "__main__": diff --git a/tests/xgboost/test_xgboost_pipeline.py b/tests/xgboost/test_xgboost_pipeline.py index 4a3d62d7..ffc7bfc2 100644 --- a/tests/xgboost/test_xgboost_pipeline.py +++ b/tests/xgboost/test_xgboost_pipeline.py @@ -10,22 +10,25 @@ from numpy.testing import assert_almost_equal import pandas import onnxruntime as rt -from xgboost import XGBRegressor, XGBClassifier, train, DMatrix +from xgboost import XGBRegressor from sklearn.model_selection import train_test_split from sklearn.compose import ColumnTransformer from sklearn.pipeline import Pipeline from sklearn.preprocessing import MinMaxScaler, OneHotEncoder from onnx.defs import onnx_opset_version from onnxconverter_common.onnx_ex import DEFAULT_OPSET_NUMBER -from onnxmltools.convert import convert_xgboost, convert_sklearn +from onnxconverter_common import data_types as onnxtypes +from onnxmltools.convert import convert_sklearn from onnxmltools.convert.common.data_types import FloatTensorType -from onnxmltools.utils import dump_data_and_model -from onnxmltools.convert.xgboost.operator_converters.XGBoost import convert_xgboost as convert_xgb -from onnxmltools.convert.common.onnx_ex import get_maximum_opset_supported +from onnxmltools.convert.xgboost.operator_converters.XGBoost import ( + convert_xgboost as convert_xgb, +) try: from skl2onnx import update_registered_converter - from skl2onnx.common.shape_calculator import calculate_linear_regressor_output_shapes + from skl2onnx.common.shape_calculator import ( + calculate_linear_regressor_output_shapes, + ) can_test = True except ImportError: @@ -38,35 +41,32 @@ @unittest.skipIf(sys.version_info[:2] <= (3, 5), reason="not available") -@unittest.skipIf(not can_test, - reason="sklearn-onnx not recent enough") +@unittest.skipIf(not can_test, reason="sklearn-onnx not recent enough") class TestXGBoostModelsPipeline(unittest.TestCase): - def _column_tranformer_fitted_from_df(self, data): def transformer_for_column(column): - if column.dtype in ['float64', 'float32']: + if column.dtype in ["float64", "float32"]: return MinMaxScaler() - if column.dtype in ['bool']: - return 'passthrough' - if column.dtype in ['O']: + if column.dtype in ["bool"]: + return "passthrough" + if column.dtype in ["O"]: return OneHotEncoder(sparse=False) raise ValueError() return ColumnTransformer( [(col, transformer_for_column(data[col]), [col]) for col in data.columns], - remainder='drop' + remainder="drop", ).fit(data) def _convert_dataframe_schema(self, data): def type_for_column(column): - if column.dtype in ['float64', 'float32']: - # onnx does not really support float64 (DoubleTensorType does not work with TreeEnsembleRegressor) + if column.dtype in ["float64", "float32"]: return onnxtypes.FloatTensorType([None, 1]) - if column.dtype in ['int64']: + if column.dtype in ["int64"]: return onnxtypes.Int64TensorType([None, 1]) - if column.dtype in ['bool']: + if column.dtype in ["bool"]: return onnxtypes.BooleanTensorType([None, 1]) - if column.dtype in ['O']: + if column.dtype in ["O"]: return onnxtypes.StringTensorType([None, 1]) raise ValueError() @@ -78,7 +78,7 @@ def test_xgboost_10_skl_missing(self): def test_xgboost_10_skl_zero(self): try: - self.common_test_xgboost_10_skl(0., True) + self.common_test_xgboost_10_skl(0.0, True) except RuntimeError as e: assert "Cannot convert a XGBoost model where missing values" in str(e) @@ -92,53 +92,62 @@ def common_test_xgboost_10_skl(self, missing, replace=False): for col in data: dtype = data[col].dtype - if dtype in ['float64', 'float32']: - data[col].fillna(0., inplace=True) - if dtype in ['int64']: + if dtype in ["float64", "float32"]: + data[col].fillna(0.0, inplace=True) + if dtype in ["int64"]: data[col].fillna(0, inplace=True) - elif dtype in ['O']: - data[col].fillna('N/A', inplace=True) + elif dtype in ["O"]: + data[col].fillna("N/A", inplace=True) - data['pclass'] = data['pclass'] * float(1) - full_df = data.drop('survived', axis=1) - full_labels = data['survived'] + data["pclass"] = data["pclass"] * float(1) + full_df = data.drop("survived", axis=1) + full_labels = data["survived"] train_df, test_df, train_labels, test_labels = train_test_split( - full_df, full_labels, test_size=.2, random_state=11) + full_df, full_labels, test_size=0.2, random_state=11 + ) col_transformer = self._column_tranformer_fitted_from_df(full_df) param_distributions = { "colsample_bytree": 0.5, "gamma": 0.2, - 'learning_rate': 0.3, - 'max_depth': 2, - 'min_child_weight': 1., - 'n_estimators': 1, - 'missing': missing, + "learning_rate": 0.3, + "max_depth": 2, + "min_child_weight": 1.0, + "n_estimators": 1, + "missing": missing, } - regressor = XGBRegressor(verbose=0, objective='reg:squarederror', - **param_distributions) + regressor = XGBRegressor( + verbose=0, objective="reg:squarederror", **param_distributions + ) regressor.fit(col_transformer.transform(train_df), train_labels) - model = Pipeline(steps=[('preprocessor', col_transformer), - ('regressor', regressor)]) + model = Pipeline( + steps=[("preprocessor", col_transformer), ("regressor", regressor)] + ) update_registered_converter( - XGBRegressor, 'XGBRegressor', + XGBRegressor, + "XGBRegressor", calculate_linear_regressor_output_shapes, - convert_xgb) + convert_xgb, + ) # last step input_xgb = model.steps[0][-1].transform(test_df[:5]).astype(np.float32) if replace: input_xgb[input_xgb[:, :] == missing] = np.nan - onnx_last = convert_sklearn(model.steps[1][-1], - initial_types=[('X', FloatTensorType(shape=[None, input_xgb.shape[1]]))], - target_opset={'': TARGET_OPSET, 'ai.onnx.ml': TARGET_OPSET_ML}) - session = rt.InferenceSession(onnx_last.SerializeToString(), providers=["CPUExecutionProvider"]) + onnx_last = convert_sklearn( + model.steps[1][-1], + initial_types=[("X", FloatTensorType(shape=[None, input_xgb.shape[1]]))], + target_opset={"": TARGET_OPSET, "ai.onnx.ml": TARGET_OPSET_ML}, + ) + session = rt.InferenceSession( + onnx_last.SerializeToString(), providers=["CPUExecutionProvider"] + ) pred_skl = model.steps[1][-1].predict(input_xgb).ravel() - pred_onx = session.run(None, {'X': input_xgb})[0].ravel() + pred_onx = session.run(None, {"X": input_xgb})[0].ravel() assert_almost_equal(pred_skl, pred_onx) diff --git a/tests/xgboost/test_xgboost_unpickle_06.py b/tests/xgboost/test_xgboost_unpickle_06.py index f29cbe2c..5761b551 100644 --- a/tests/xgboost/test_xgboost_unpickle_06.py +++ b/tests/xgboost/test_xgboost_unpickle_06.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -""" -Tests scikit-linear converter. -""" -import sys import os import packaging.version as pv import unittest @@ -19,17 +15,21 @@ class TestXGBoostUnpickle06(unittest.TestCase): - - @unittest.skipIf(pv.Version(xgboost.__version__) >= pv.Version('1.0'), - reason="compatibility break with pickle in 1.0") + @unittest.skipIf( + pv.Version(xgboost.__version__) >= pv.Version("1.0"), + reason="compatibility break with pickle in 1.0", + ) def test_xgboost_unpickle_06(self): # Unpickle a model trained with an old version of xgboost. this = os.path.dirname(__file__) with open(os.path.join(this, "xgboost10day.pickle.dat"), "rb") as f: xgb = pickle.load(f) - conv_model = convert_xgboost(xgb, initial_types=[('features', FloatTensorType(['None', 10000]))], - target_opset=TARGET_OPSET) + conv_model = convert_xgboost( + xgb, + initial_types=[("features", FloatTensorType(["None", 10000]))], + target_opset=TARGET_OPSET, + ) assert conv_model is not None From b268fe75cba90fa8e59818c9fe60a7f65a44cd13 Mon Sep 17 00:00:00 2001 From: matthias_oh Date: Mon, 20 Mar 2023 15:22:34 +0800 Subject: [PATCH 5/5] Fix: buildInputDictSimple dataType Stringtype() --- onnxmltools/convert/sparkml/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxmltools/convert/sparkml/utils.py b/onnxmltools/convert/sparkml/utils.py index d1589ba0..5fe8e469 100644 --- a/onnxmltools/convert/sparkml/utils.py +++ b/onnxmltools/convert/sparkml/utils.py @@ -44,7 +44,7 @@ def buildInputDictSimple(dataframe): result = {} for field in dataframe.schema.fields: - if str(field.dataType) == "StringType": + if str(field.dataType) == 'StringType' or str(field.dataType) == 'StringType()': result[field.name] = dataframe.select(field.name).toPandas().values else: result[field.name] = (