diff --git a/coremltools/_deps/__init__.py b/coremltools/_deps/__init__.py index 0410e7833..31ce1eec2 100644 --- a/coremltools/_deps/__init__.py +++ b/coremltools/_deps/__init__.py @@ -9,19 +9,28 @@ """ from distutils.version import StrictVersion as _StrictVersion import logging as _logging +from packaging import version import platform as _platform import re as _re import sys as _sys -from packaging import version -def __get_version(version): +def _get_version(version): # matching 1.6.1, and 1.6.1rc, 1.6.1.dev version_regex = r"^\d+\.\d+\.\d+" version = _re.search(version_regex, str(version)).group(0) return _StrictVersion(version) +def _warn_if_above_max_supported_version(package_name, package_version, max_supported_version): + if _get_version(package_version) > _StrictVersion(max_supported_version): + _logging.warning( + "%s version %s has not been tested with coremltools. You may run into unexpected errors. " + "%s %s is the most recent version that has been tested." + % (package_name, package_version, package_name, max_supported_version) + ) + + # --------------------------------------------------------------------------------------- _IS_MACOS = _sys.platform == "darwin" @@ -77,8 +86,10 @@ def __get_sklearn_version(version): # --------------------------------------------------------------------------------------- _HAS_XGBOOST = True +_XGBOOST_MAX_VERSION = "1.4.2" try: import xgboost + _warn_if_above_max_supported_version("XGBoost", xgboost.__version__, _XGBOOST_MAX_VERSION) except: _HAS_XGBOOST = False @@ -89,12 +100,12 @@ def __get_sklearn_version(version): _TF_1_MIN_VERSION = "1.12.0" _TF_1_MAX_VERSION = "1.15.0" _TF_2_MIN_VERSION = "2.1.0" -_TF_2_MAX_VERSION = "2.3.1" +_TF_2_MAX_VERSION = "2.5.0" try: import tensorflow - tf_ver = __get_version(tensorflow.__version__) + tf_ver = _get_version(tensorflow.__version__) # TensorFlow if tf_ver < _StrictVersion("2.0.0"): @@ -112,11 +123,7 @@ def __get_sklearn_version(version): ) % (tensorflow.__version__, _TF_1_MIN_VERSION) ) - elif tf_ver > _StrictVersion(_TF_1_MAX_VERSION): - _logging.warning( - "TensorFlow version %s detected. Last version known to be fully compatible is %s ." - % (tensorflow.__version__, _TF_1_MAX_VERSION) - ) + _warn_if_above_max_supported_version("TensorFlow", tensorflow.__version__, _TF_1_MAX_VERSION) elif _HAS_TF_2: if tf_ver < _StrictVersion(_TF_2_MIN_VERSION): _logging.warn( @@ -126,11 +133,7 @@ def __get_sklearn_version(version): ) % (tensorflow.__version__, _TF_2_MIN_VERSION) ) - elif tf_ver > _StrictVersion(_TF_2_MAX_VERSION): - _logging.warning( - "TensorFlow version %s detected. Last version known to be fully compatible is %s ." - % (tensorflow.__version__, _TF_2_MAX_VERSION) - ) + _warn_if_above_max_supported_version("TensorFlow", tensorflow.__version__, _TF_2_MAX_VERSION) except: _HAS_TF = False @@ -168,7 +171,7 @@ def __get_sklearn_version(version): sys.stderr = stderr import tensorflow - k_ver = __get_version(keras.__version__) + k_ver = _get_version(keras.__version__) # keras 1 version too old if k_ver < _StrictVersion(_KERAS_MIN_VERSION): @@ -186,7 +189,8 @@ def __get_sklearn_version(version): _HAS_KERAS_TF = False _logging.warning( ( - "Keras version %s detected. Last version known to be fully compatible of Keras is %s ." + "Keras version %s has not been tested with coremltools. You may run into unexpected errors. " + "Keras %s is the most recent version that has been tested." ) % (keras.__version__, _KERAS_MAX_VERSION) ) @@ -214,8 +218,10 @@ def __get_sklearn_version(version): # --------------------------------------------------------------------------------------- _HAS_TORCH = True +_TORCH_MAX_VERSION = "1.9.1" try: import torch + _warn_if_above_max_supported_version("Torch", torch.__version__, _TORCH_MAX_VERSION) except: _HAS_TORCH = False MSG_TORCH_NOT_FOUND = "PyTorch not found." diff --git a/coremltools/converters/_converters_entry.py b/coremltools/converters/_converters_entry.py index 030c5d4d5..7f3ea77a6 100644 --- a/coremltools/converters/_converters_entry.py +++ b/coremltools/converters/_converters_entry.py @@ -46,7 +46,9 @@ def convert( compute_precision=None, skip_model_load=False, compute_units=_ComputeUnit.ALL, - **kwargs + useCPUOnly=False, + package_dir=None, + debug=False, ): """ Convert a TensorFlow or PyTorch model to the Core ML model format as either @@ -230,6 +232,24 @@ def convert( - ``coremltools.ComputeUnit.CPU_AND_GPU``: Use both the CPU and GPU, but not the neural engine. + useCPUOnly: bool + Deprecated, to be removed in coremltools 6.0. Please use `compute_units` instead. + - if True, identical to setting compute_units to `coremltools.ComputeUnit.CPU_ONLY`` + - if False, identical to setting compute_units to `coremltools.ComputeUnit.ALL`` + + package_dir : str + Post conversion, the model is compiled to form the MLModel object ready for prediction. + This requires a temporary directory to hold the mlmodelc archive. + - if not None, must be a path to a directory that is used for + temporarily storing the compiled model assets. If None, a temporary directory is created. + + debug : bool + This flag should generally be False except for debugging purposes + Setting this flag to True: + - For Torch conversion, it will print the list of supported and unsupported ops + found in the model if conversion fails due to an unsupported op. + - For Tensorflow conversion, it will cause to display extra logging and visualizations + Returns ------- model : ``coremltools.models.MLModel`` or ``coremltools.converters.mil.Program`` @@ -284,9 +304,9 @@ def convert( exact_source = _determine_source(model, source, outputs) exact_target = _determine_target(convert_to, minimum_deployment_target) _validate_inputs(model, exact_source, inputs, outputs, classifier_config, compute_precision, - exact_target, **kwargs) + exact_target) - if "useCPUOnly" in kwargs and kwargs["useCPUOnly"]: + if useCPUOnly: warnings.warn('The "useCPUOnly" parameter is deprecated and will be removed in 6.0. ' 'Use the compute_units parameter: "compute_units=coremotools.ComputeUnits.CPU_ONLY".') compute_units = _ComputeUnit.CPU_ONLY @@ -313,7 +333,8 @@ def convert( transforms=tuple(transforms), skip_model_load=skip_model_load, compute_units=compute_units, - **kwargs + package_dir=package_dir, + debug=debug, ) if exact_target == 'milinternal': @@ -344,8 +365,7 @@ def _check_deployment_target(minimum_deployment_target): ) raise TypeError(msg.format(minimum_deployment_target)) -def _validate_inputs(model, exact_source, inputs, outputs, classifier_config, compute_precision, convert_to, - **kwargs): +def _validate_inputs(model, exact_source, inputs, outputs, classifier_config, compute_precision, convert_to): """ Validate and process model, inputs, outputs, classifier_config based on `exact_source` (which cannot be `auto`) @@ -399,10 +419,6 @@ def raise_if_duplicated(input_list): raise ValueError("Input should be a list of TensorType or ImageType") elif exact_source == "pytorch": - if "example_inputs" in kwargs: - msg = 'Unexpected argument "example_inputs" found' - raise ValueError(msg) - if inputs is None: msg = 'Expected argument for pytorch "inputs" not provided' raise ValueError(msg) diff --git a/coremltools/converters/_profile_utils.py b/coremltools/converters/_profile_utils.py index 578efbf9d..1f59c4a27 100644 --- a/coremltools/converters/_profile_utils.py +++ b/coremltools/converters/_profile_utils.py @@ -1,3 +1,8 @@ +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + import os import time diff --git a/coremltools/converters/mil/__init__.py b/coremltools/converters/mil/__init__.py index e36bfb823..4b484545f 100644 --- a/coremltools/converters/mil/__init__.py +++ b/coremltools/converters/mil/__init__.py @@ -3,7 +3,58 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from .mil import * +# This import should be pruned rdar://84519338 +from .mil import ( + block, + Block, + BoolInputType, + BoolTensorInputType, + builder, + Builder, + curr_block, + DefaultInputs, + FloatInputType, + FloatTensorInputType, + Function, + get_existing_symbol, + get_new_symbol, + get_new_variadic_symbol, + input_type, + InputSpec, + IntInputType, + IntOrFloatInputType, + IntOrFloatOrBoolInputType, + IntTensorInputType, + InternalInputType, + InternalScalarOrTensorInputType, + InternalStringInputType, + InternalVar, + ListInputType, + ListOrScalarOrTensorInputType, + ListVar, + mil_list, + operation, + Operation, + ops, + Placeholder, + precondition, + program, + Program, + PyFunctionInputType, + register_op, + SPACES, + SUPPORT_FLOAT_TYPES, + SUPPORT_INT_TYPES, + ScalarOrTensorInputType, + StringInputType, + Symbol, + TensorInputType, + TupleInputType, + types, + var, + Var, + visitors +) from .frontend.torch import register_torch_op diff --git a/coremltools/converters/mil/backend/mil/__init__.py b/coremltools/converters/mil/backend/mil/__init__.py index 3256cf5b2..845cdd674 100644 --- a/coremltools/converters/mil/backend/mil/__init__.py +++ b/coremltools/converters/mil/backend/mil/__init__.py @@ -1,3 +1,6 @@ -# Copyright (c) 2020, Apple Inc. All rights reserved. +# Copyright (c) 2020, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from .load import load diff --git a/coremltools/converters/mil/backend/mil/helper.py b/coremltools/converters/mil/backend/mil/helper.py index 5214bf7ef..3437af8d3 100644 --- a/coremltools/converters/mil/backend/mil/helper.py +++ b/coremltools/converters/mil/backend/mil/helper.py @@ -1,3 +1,8 @@ +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + import numpy as np import os import re diff --git a/coremltools/converters/mil/backend/mil/load.py b/coremltools/converters/mil/backend/mil/load.py index 818d88b57..964457c29 100644 --- a/coremltools/converters/mil/backend/mil/load.py +++ b/coremltools/converters/mil/backend/mil/load.py @@ -4,47 +4,56 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import logging -import numpy as _np +import numpy as np import os -import tempfile -import shutil -from coremltools.converters.mil.backend.mil.helper import * -from coremltools.converters.mil.backend.backend_helper import _get_probability_var_for_classifier from .passes import mil_passes -import coremltools.proto.MIL_pb2 as pm -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil import Function +from coremltools import _SPECIFICATION_VERSION_IOS_15 +from coremltools.converters.mil.backend.mil.helper import ( + cast_to_framework_io_dtype, + create_file_value, + create_immediate_value, + create_list_scalarvalue, + create_scalar_value, + types_to_proto +) +from coremltools.converters.mil.backend.backend_helper import _get_probability_var_for_classifier +from coremltools.converters.mil.mil import ( + Builder as mb, + Function, + mil_list, + types +) from coremltools.converters.mil.backend.nn.load import _set_optional_inputs +from coremltools.converters.mil.input_types import ImageType, TensorType, EnumeratedShapes, RangeDim from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry from coremltools.converters.mil.mil.types.symbolic import ( any_symbolic, any_variadic, is_symbolic, ) +from coremltools.converters.mil.mil.types.type_mapping import types_int64 +from coremltools.libmilstoragepython import _BlobStorageWriter as BlobWriter +from coremltools.models.model import _WEIGHTS_FILE_NAME from coremltools.models.neural_network.flexible_shape_utils import ( - NeuralNetworkImageSize, - NeuralNetworkImageSizeRange, add_enumerated_image_sizes, add_multiarray_ndshape_enumeration, + NeuralNetworkImageSize, + NeuralNetworkImageSizeRange, set_multiarray_ndshape_range, - update_image_size_range, + update_image_size_range +) +from coremltools.proto import ( + FeatureTypes_pb2 as ft, + MIL_pb2 as pm, + Model_pb2 as ml ) -from coremltools.libmilstoragepython import _BlobStorageWriter as BlobWriter - -import coremltools.proto.Model_pb2 as ml -import coremltools.proto.FeatureTypes_pb2 as ft -from coremltools.converters.mil.input_types import ImageType, TensorType, EnumeratedShapes, RangeDim -from coremltools.models.model import _WEIGHTS_FILE_NAME -from coremltools.converters.mil.mil import Builder as mb -from coremltools.converters.mil.mil import mil_list -from coremltools import _SPECIFICATION_VERSION_IOS_15 def should_use_weight_file(val): return ( val is not None - and isinstance(val, (_np.ndarray, _np.generic)) + and isinstance(val, (np.ndarray, np.generic)) and val.size >= 10 and val.dtype in ['float16', 'float32'] ) @@ -97,7 +106,7 @@ def translate_generic_op(op, parameters, blob_writer, literal_params=[]): blocks = None if len(op.blocks) > 0: blocks = [create_block(b, parameters, blob_writer) \ - for b in op.blocks] + for b in op.blocks] op_type = op.op_type attr_dict = {} @@ -206,6 +215,9 @@ def _add_classify_op(prog, classifier_config): # add the classify op now with block: + # cast the int label to np.int64 + if isinstance(classes[0], int): + classes = [np.int64(x) for x in classes] classes_var = mb.const(val=mil_list(classes)) out = mb.classify(probabilities=probability_var, classes=classes_var) @@ -344,7 +356,7 @@ def load(prog, weights_dir, resume_on_errors=False, **kwargs): keytype, valtype = var.sym_type.T if types.is_str(keytype): output_feature_type.dictionaryType.stringKeyType.MergeFromString(b"") - elif (keytype == types_int64): + elif (keytype == types.int64): output_feature_type.dictionaryType.int64KeyType.MergeFromString(b"") else: raise ValueError("Dictionary key type not supported.") @@ -445,7 +457,6 @@ def load(prog, weights_dir, resume_on_errors=False, **kwargs): model, input_name, lower_bounds=lb, upper_bounds=ub ) - # Set optional inputs _set_optional_inputs(model, input_types) diff --git a/coremltools/converters/mil/backend/mil/passes/__init__.py b/coremltools/converters/mil/backend/mil/passes/__init__.py index d58e204f8..18f1a9071 100644 --- a/coremltools/converters/mil/backend/mil/passes/__init__.py +++ b/coremltools/converters/mil/backend/mil/passes/__init__.py @@ -3,20 +3,10 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -# Import all passes in this dir -from os.path import dirname, basename, isfile, join -import glob - -excluded_files = ["mil_passes.py", "__init__.py"] -modules = glob.glob(join(dirname(__file__), "*.py")) -pass_modules = [ - basename(f)[:-3] - for f in modules - if isfile(f) - and basename(f)[:1] != "_" # Follow python convention to hide _* files. - and basename(f)[:4] != "test" - and basename(f) not in excluded_files -] -__all__ = pass_modules - -from . import * # import everything in __all__ +from . import ( + adjust_io_to_supported_types, + fuse_activation_silu, + homogenize_input_dtypes, + insert_image_preprocessing_op, + sanitize_name_strings +) diff --git a/coremltools/converters/mil/backend/mil/passes/adjust_io_to_supported_types.py b/coremltools/converters/mil/backend/mil/passes/adjust_io_to_supported_types.py index 914f699a5..668d29548 100644 --- a/coremltools/converters/mil/backend/mil/passes/adjust_io_to_supported_types.py +++ b/coremltools/converters/mil/backend/mil/passes/adjust_io_to_supported_types.py @@ -5,15 +5,15 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil import Builder as _mb -from coremltools.converters.mil.mil import types as _types -from coremltools.converters.mil.mil.ops import defs as _ops -from coremltools.converters.mil.mil.passes.pass_registry import register_pass as _register_pass +from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil import types as types +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass import warnings as _warnings -@_register_pass(namespace="mil_backend") -def adjust_io_to_supported_types(prog): +@register_pass(namespace="mil_backend") +class adjust_io_to_supported_types(AbstractGraphPass): """ Converts all dTypes to types that are supported by the CoreML runtime. The runtime supports only fp16, fp32, int32, str, and bool variables. @@ -59,41 +59,43 @@ def adjust_io_to_supported_types(prog): is unchanged. """ - for name, func in prog.functions.items(): - _adjust_io_to_supported_types(func, name == "main") + def apply(self, prog): + for name, func in prog.functions.items(): + is_main_funtion = name == "main" + _adjust_io_to_supported_types(func, is_main_funtion) -__RUNTIME_SUPPORTED_TYPES = [_types.fp16, _types.fp32, _types.int32, _types.str, _types.bool] +__RUNTIME_SUPPORTED_TYPES = [types.fp16, types.fp32, types.int32, types.str, types.bool] ##### # Main Function ##### def _adjust_var_dtype_helper(var, dtype): - if (_types.is_scalar(var.sym_type)): + if (types.is_scalar(var.sym_type)): var._sym_type = dtype else: - var._sym_type = _types.tensor(dtype, var.sym_type.get_shape()) + var._sym_type = types.tensor(dtype, var.sym_type.get_shape()) def _adjust_main_inputs(func): first_op = func.operations[0] if len(func.operations) > 0 else None for input_name, input_var in func.inputs.items(): - if (_types.is_tensor(input_var.sym_type) or _types.is_scalar(input_var.sym_type)) \ - and input_var.dtype != _types.fp32 \ - and input_var.dtype != _types.int32: - input_dtype_str = _types.builtin_to_string(input_var.dtype) - if _types.is_int(input_var.dtype): + if (types.is_tensor(input_var.sym_type) or types.is_scalar(input_var.sym_type)) \ + and input_var.dtype != types.fp32 \ + and input_var.dtype != types.int32: + input_dtype_str = types.builtin_to_string(input_var.dtype) + if types.is_int(input_var.dtype): # Replace non-int32 input type with int32. _warnings.warn("Input" + input_var.name + " is of dType " + input_dtype_str +\ ". Only integer variables of bit width 32 are supported by the CoreML runtime. " +\ "This input will be assigned a dType of int32. " +\ "No cast will be inserted; the previous dtype will be replaced.") - _adjust_var_dtype_helper(input_var, _types.int32) - elif input_var.dtype == _types.fp64: + _adjust_var_dtype_helper(input_var, types.int32) + elif input_var.dtype == types.fp64: # Replace float64 input type with fp32. _warnings.warn("Input" + input_var.name + " is of dtype fp64. 64 bit float inputs are " +\ "not supported by ML program models. This input will be assigned a dType " +\ "of fp32. No cast will be inserted; the previous dtype will be replaced.") - _adjust_var_dtype_helper(input_var, _types.fp32) + _adjust_var_dtype_helper(input_var, types.fp32) else: # This is some other dType. Change the type to fp32 and add a cast. # This is only a limitation of main--other functions do not represent CoreML model inputs @@ -104,18 +106,18 @@ def _adjust_main_inputs(func): "fp32. A cast will be inserted at the beginning of the program to " +\ "convert the input to the originally defined dType.") with func: - casted_input_var = _mb.cast(x=input_var, dtype=input_dtype_str, before_op=first_op) + casted_input_var = mb.cast(x=input_var, dtype=input_dtype_str, before_op=first_op) func.replace_uses_of_var_after_op(anchor_op=casted_input_var.op, old_var=input_var, new_var=casted_input_var) - _adjust_var_dtype_helper(input_var, _types.fp32) + _adjust_var_dtype_helper(input_var, types.fp32) def _adjust_main_outputs(func): new_outputs = [] for output_var in func.outputs: output_type = output_var.sym_type - if (_types.is_tensor(output_type) or _types.is_scalar(output_type)) \ - and output_var.dtype != _types.fp32 \ - and output_var.dtype != _types.int32: - output_dtype_str = _types.builtin_to_string(output_var.dtype) + if (types.is_tensor(output_type) or types.is_scalar(output_type)) \ + and output_var.dtype != types.fp32 \ + and output_var.dtype != types.int32: + output_dtype_str = types.builtin_to_string(output_var.dtype) _warnings.warn("Output" + output_var.name + " is of dType " + output_dtype_str + ". The " +\ "CoreML runtime does not support outputs with this dType (only int32 and " +\ "fp32 are supported for outputs). This output will be assigned a dType " +\ @@ -126,7 +128,7 @@ def _adjust_main_outputs(func): output_var.set_name(output_var_name + "__pre__output__fp32__cast") # Convert the output to fp32, and add a cast. with func: - output_var = _mb.cast(x=output_var, dtype="fp32") + output_var = mb.cast(x=output_var, dtype="fp32") output_var.set_name(output_var_name) new_outputs.append(output_var) func.set_outputs(new_outputs) @@ -141,23 +143,23 @@ def _adjust_var(var): to the rules outlined in the top level pass comment (see adjust_io_to_supported_types). """ - if (_types.is_tensor(var.sym_type) or _types.is_scalar(var.sym_type)) \ + if (types.is_tensor(var.sym_type) or types.is_scalar(var.sym_type)) \ and var.dtype not in __RUNTIME_SUPPORTED_TYPES: - dtype_str = _types.builtin_to_string(var.dtype) - if _types.is_int(var.dtype): + dtype_str = types.builtin_to_string(var.dtype) + if types.is_int(var.dtype): # Replace non-int32 input type with int32. _warnings.warn("Input" + var.name + " is of dType " + dtype_str +\ ". Only integer variables of bit width 32 are supported by the CoreML runtime. " +\ "This input will be assigned a dType of int32. " +\ "No cast will be inserted; the previous dtype will be replaced.") - _adjust_var_dtype_helper(var, _types.int32) + _adjust_var_dtype_helper(var, types.int32) else: # This is some other unsupported dType. Change the input type to fp32. _warnings.warn("Var " + var.name + " is of dType " + dtype_str + ". The CoreML runtime " +\ "does not support this dType (only fp16, fp32, bool, and int32 are supported). " +\ "This input will be assigned a dType of fp32. No cast will be inserted; " +\ "the previous dtype will be replaced.") - _adjust_var_dtype_helper(var, _types.fp32) + _adjust_var_dtype_helper(var, types.fp32) def _adjust_func_inputs(func): @@ -193,7 +195,7 @@ def _adjust_ops(block): # If the output dtype or input dtype was previously adjusted, # the cast op must change or be removed in kind. if op.op_type == "cast": - output_type_str = _types.builtin_to_string(op.outputs[0].dtype) + output_type_str = types.builtin_to_string(op.outputs[0].dtype) if op.outputs[0].dtype == op.x.dtype: # The type of the input or output of this cast op was changed per the rules # defined in the top level comment for adjust_io_to_supported_types. @@ -217,7 +219,7 @@ def _adjust_ops(block): # This cast is meaningful, and the "dtype" param now differs from the output # type. Replace the dtype cast with a new cast op with a matching dtype param. with block: - new_cast_out = _mb.cast(x=op.x, dtype=output_type_str, before_op=op) + new_cast_out = mb.cast(x=op.x, dtype=output_type_str, before_op=op) block.replace_uses_of_var_after_op( anchor_op=op, old_var=op.outputs[0], new_var=new_cast_out ) diff --git a/coremltools/converters/mil/backend/mil/passes/fuse_activation_silu.py b/coremltools/converters/mil/backend/mil/passes/fuse_activation_silu.py index b22518bb4..2ce5826ee 100644 --- a/coremltools/converters/mil/backend/mil/passes/fuse_activation_silu.py +++ b/coremltools/converters/mil/backend/mil/passes/fuse_activation_silu.py @@ -4,9 +4,10 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb -def match_pattern(op): +def _match_pattern(op): if op.op_type == "sigmoid": # abort fusion if op output is also a block output if op.outputs[0] in op.enclosing_block.outputs: @@ -26,7 +27,7 @@ def match_pattern(op): return None -def try_to_transform(sigmoid_op, mul_op, block): +def _try_to_transform(sigmoid_op, mul_op, block): out_name = mul_op.outputs[0].name # create a new silu op x = mb.silu(x=sigmoid_op.x, name=out_name, before_op=sigmoid_op) @@ -38,20 +39,20 @@ def try_to_transform(sigmoid_op, mul_op, block): return True -def fuse_activation_silu_block(block): +def _fuse_activation_silu_block(block): fusion_status = False for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = fuse_activation_silu_block(b) + block_changed = _fuse_activation_silu_block(b) if len(op.blocks) > 0: continue - mul_op = match_pattern(op) + mul_op = _match_pattern(op) if mul_op is not None: with block: - fusion_status = try_to_transform(op, mul_op, block) + fusion_status = _try_to_transform(op, mul_op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status @@ -59,7 +60,7 @@ def fuse_activation_silu_block(block): @register_pass(namespace="mil_backend") -def fuse_activation_silu(prog): +class fuse_activation_silu(AbstractGraphPass): """ Fold x * sigmoid(x) into silu(x) @@ -72,7 +73,8 @@ def fuse_activation_silu(prog): %3 = silu(%0) ... """ - for f_name, f in prog.functions.items(): - block_changed = True - while block_changed: - block_changed = fuse_activation_silu_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_activation_silu_block(f) diff --git a/coremltools/converters/mil/backend/mil/passes/homogenize_input_dtypes.py b/coremltools/converters/mil/backend/mil/passes/homogenize_input_dtypes.py index 35381b3da..8617d12e6 100644 --- a/coremltools/converters/mil/backend/mil/passes/homogenize_input_dtypes.py +++ b/coremltools/converters/mil/backend/mil/passes/homogenize_input_dtypes.py @@ -7,6 +7,7 @@ from coremltools.converters.mil.mil.ops.defs import elementwise_binary, matmul from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.types import promote_dtypes, builtin_to_string from coremltools.converters.mil.mil import Builder as mb @@ -23,18 +24,15 @@ def _get_input_params(op): return params return None - -def _of_same_dtype(dtype1, dtype2): +def _is_same_dtype(dtype1, dtype2): return (dtype1 is dtype2) or (builtin_to_string(dtype1) == builtin_to_string(dtype2)) - def _promoted_var(op, var, promoted_dtype): x = mb.cast( x=var, dtype=builtin_to_string(promoted_dtype), name=var.name + "_promoted", before_op=op ) return x - def _homogenize_input_dtypes_block(block): for op in list(block.operations): for b in op.blocks: @@ -47,7 +45,7 @@ def _homogenize_input_dtypes_block(block): promoted_dtype = promote_dtypes([var.dtype for var in input_vars]) for i,var in enumerate(input_vars): - if not _of_same_dtype(var.dtype, promoted_dtype): + if not _is_same_dtype(var.dtype, promoted_dtype): has_mixed_dtypes = True with block: input_vars[i] = _promoted_var(op, var, promoted_dtype) @@ -59,7 +57,10 @@ def _homogenize_input_dtypes_block(block): {k: v for k, v in op.inputs.items() if k not in new_inputs} ) with block: - new_output = getattr(mb, op.op_type)(**new_inputs) + # create a new op with the promoted input vars + new_op_class = getattr(mb,op.op_type) + new_output = new_op_class(**new_inputs) + op.enclosing_block.replace_uses_of_var_after_op( anchor_op=op, old_var=op.outputs[0], new_var=new_output, no_check_var_types=True, # Has to set no_check_var_types=True because Matmul PyMIL type inference doesn't enforce same dtypes for x & y @@ -67,9 +68,8 @@ def _homogenize_input_dtypes_block(block): ) block.remove_ops([op]) - @register_pass(namespace="mil_backend") -def homogenize_input_dtypes(prog): +class homogenize_input_dtypes(AbstractGraphPass): """ If inputs to an op, doesn't have same dtypes for some parameters, explicit cast operations are injected to ensure inputs to that op have same promoted dtype. @@ -77,8 +77,9 @@ def homogenize_input_dtypes(prog): - Only ops specified in dict _SUPPORTED_OPS as its keys, can be affected by this pass - Only the named inputs of ops specified in dict _SUPPORTED_OPS as values, are promoted to match dtypes """ - for f_name, f in prog.functions.items(): - _homogenize_input_dtypes_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _homogenize_input_dtypes_block(f) - for op in f.operations: - op.type_value_inference(overwrite_output=True) + for op in f.operations: + op.type_value_inference(overwrite_output=True) diff --git a/coremltools/converters/mil/backend/mil/passes/insert_image_preprocessing_op.py b/coremltools/converters/mil/backend/mil/passes/insert_image_preprocessing_op.py index fb9565655..f73ab7f8d 100644 --- a/coremltools/converters/mil/backend/mil/passes/insert_image_preprocessing_op.py +++ b/coremltools/converters/mil/backend/mil/passes/insert_image_preprocessing_op.py @@ -6,6 +6,7 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.input_types import ImageType # import mil internal ops to add it to the builder from coremltools.converters.mil.mil.ops import defs as _ops @@ -15,13 +16,14 @@ import numpy as np @register_pass(namespace="mil_backend") -def insert_image_preprocessing_ops(prog): +class insert_image_preprocessing_ops(AbstractGraphPass): """ Insert preprocessing ops, right after the input if its of type Image """ - for f_name, f in prog.functions.items(): - if f_name == 'main': - _insert_image_preprocessing_ops(f, prog) + def apply(self, prog): + for f_name, f in prog.functions.items(): + if f_name == 'main': + _insert_image_preprocessing_ops(f, prog) def _insert_image_preprocessing_ops(block, prog): diff --git a/coremltools/converters/mil/backend/mil/passes/mil_passes.py b/coremltools/converters/mil/backend/mil/passes/mil_passes.py index 9b260968f..c5c6ced8a 100644 --- a/coremltools/converters/mil/backend/mil/passes/mil_passes.py +++ b/coremltools/converters/mil/backend/mil/passes/mil_passes.py @@ -3,9 +3,10 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import logging as _logging + from coremltools.converters.mil.backend.nn.passes.nn_passes import nn_backend_passes from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY -import logging as _logging def mil_backend_passes(prog): diff --git a/coremltools/converters/mil/backend/mil/passes/sanitize_name_strings.py b/coremltools/converters/mil/backend/mil/passes/sanitize_name_strings.py index 816e94a6f..24d65375a 100644 --- a/coremltools/converters/mil/backend/mil/passes/sanitize_name_strings.py +++ b/coremltools/converters/mil/backend/mil/passes/sanitize_name_strings.py @@ -6,20 +6,21 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.passes.name_sanitization_utils import NameSanitizer, sanitize_block @register_pass(namespace="mil_backend") -def sanitize_name_strings(prog): +class sanitize_name_strings(AbstractGraphPass): """ Sanitize the names of vars and ops to make sure that they are of the format as described in the NameSanitizer class, i.e. of the format [a-zA-Z_][a-zA-Z0-9_]* """ + def apply(self, prog): + sanitizer_vars = NameSanitizer(prefix="var_") + sanitizer_ops = NameSanitizer(prefix="op_") - sanitizer_vars = NameSanitizer(prefix="var_") - sanitizer_ops = NameSanitizer(prefix="op_") - - for _, f in prog.functions.items(): - sanitize_block(f, sanitizer_vars, sanitizer_ops, prog.main_input_types) + for f in prog.functions.values(): + sanitize_block(f, sanitizer_vars, sanitizer_ops, prog.main_input_types) diff --git a/coremltools/converters/mil/backend/nn/op_mapping.py b/coremltools/converters/mil/backend/nn/op_mapping.py index 75e53738a..6866f9a5e 100644 --- a/coremltools/converters/mil/backend/nn/op_mapping.py +++ b/coremltools/converters/mil/backend/nn/op_mapping.py @@ -5,6 +5,9 @@ import numpy as _np import logging as _logging +from tqdm import tqdm as _tqdm + +from .mil_to_nn_mapping_registry import MIL_TO_NN_MAPPING_REGISTRY, register_mil_to_nn_mapping from coremltools.models import neural_network as neural_network from coremltools.proto import NeuralNetwork_pb2 from coremltools.converters.mil.mil.types.symbolic import ( @@ -18,8 +21,6 @@ from coremltools.models.neural_network.quantization_utils import ( _convert_array_to_nbit_quantized_bytes, ) -from tqdm import tqdm as _tqdm -from .mil_to_nn_mapping_registry import * def convert_ops(const_context, builder, ops, outputs): @@ -431,12 +432,6 @@ def conv_helper(const_context, builder, op): has_bias = op.bias is not None groups = op.groups.val - rank_factor = 1 - if is_conv2d: - rank_factor = 2 - elif is_conv3d: - rank_factor = 3 - strides = op.strides.val.tolist() dilations = op.dilations.val.tolist() if is_conv1d: @@ -640,7 +635,6 @@ def _add_elementwise_binary( return add_const(const_context, builder, op.y.name, op.y.val) - if mode in {"add", "multiply", "max", "min"} and op.x.shape == op.y.shape: builder.add_elementwise( name=name, @@ -1045,8 +1039,7 @@ def slice_by_index(const_context, builder, op): for i in range(rank): if (not begin_mask[i] and begin[i] != 0) or \ (not end_mask[i] and end[i] != op.x.shape[i]): - slice_dim.append(i) - + slice_dim.append(i) if len(slice_dim) == 1 and not squeeze_mask[slice_dim[0]]: dim = slice_dim[0] - rank @@ -1117,7 +1110,6 @@ def slice_by_size(const_context, builder, op): # The static case if op.begin.val is not None and op.size.val is not None: - shape = op.x.shape begin = op.begin.val size = op.size.val rank = op.x.rank @@ -2381,7 +2373,7 @@ def pad(const_context, builder, op): pad = pad[-4:] left, right = pad[2], pad[3] top, bottom = pad[0], pad[1] - layer = builder.add_padding( + builder.add_padding( name=op.name, left=left, right=right, @@ -2480,7 +2472,7 @@ def layer_norm(const_context, builder, op): # - reshape back to (X1, X2) / (X0, X1, X2) # Otherwise, we express the layer_norm as primitive operations if rank in [2, 3] and len(axes) == 1 and axes[0] == rank - 1 and input_shape.count(-1) < 2 \ - and input_shape[-1] != -1 and input_shape[-2] != -1: + and input_shape[-1] != -1 and input_shape[-2] != -1: reshaped_shape = input_shape[:] # Insert a singleton dimension in the 'height' position @@ -2671,11 +2663,6 @@ def conv_transpose(const_context, builder, op): else: weight = _np.transpose(weight, [2, 3, 0, 1]) - # Adjust for Deconv1D case - # CoreML maps Deconv1D into Deconv2D - # Hence, adjust width dimension attributes by setting to 1 for 1D case - rank_factor = 1 if is_conv_transpose_1d else 2 - strides = op.strides.val.tolist() dilations = op.dilations.val.tolist() @@ -3019,7 +3006,6 @@ def while_loop(const_context, builder, op): # Also assume all outputs are different from loop inputs (i.e., no loop # invariant.) - #for vx_in, vx_out in zip(block.inputs, block.outputs[1:]): for vx_in, vx_out in zip(body_block.inputs, body_block.outputs): if vx_in.name == vx_out.name: msg = "Loop invariant var {} detected in block {}" @@ -3094,12 +3080,6 @@ def stack(const_context, builder, op): @register_mil_to_nn_mapping def split(const_context, builder, op): - split_sizes = None - if op.split_sizes is not None: - if op.split_sizes.val is None: - raise ValueError('Non-const split_sizes unsupported in NN') - split_sizes = op.split_sizes.val.tolist() - split = op.sizes split = [size for size in split if size != 0] has_equal_splits = all([size == split[0] for size in split]) @@ -3271,7 +3251,7 @@ def make_list(const_context, builder, op): # set the dynamic dimensions to 1 for initialization # Ex: op.elem_shape = [i0, 128] will result in [1, 128] elem_shape = [1 if isinstance(dim_var.val, str) else - dim_var.val for dim_var in op.elem_shape] + dim_var.val for dim_var in op.elem_shape] if size is not None: array_size = size if size > 0 else 1 @@ -3294,7 +3274,7 @@ def make_list(const_context, builder, op): # Concatenate list length of the input, should be a constant vector of size 1) with element shape node_arr_shape_name = op.name + "_arr_shape" - layer = builder.add_concat_nd( + builder.add_concat_nd( name=node_arr_shape_name, input_names=[op.init_length.name, node_es_name], output_name=node_arr_shape_name, @@ -3614,4 +3594,3 @@ def list_length(const_context, builder, op): def _const_symbolic(const_context, builder, op): # do nothing pass - diff --git a/coremltools/converters/mil/backend/nn/passes/__init__.py b/coremltools/converters/mil/backend/nn/passes/__init__.py index 12e97024f..cffa980aa 100644 --- a/coremltools/converters/mil/backend/nn/passes/__init__.py +++ b/coremltools/converters/mil/backend/nn/passes/__init__.py @@ -3,23 +3,11 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -# Import all passes in this dir -from os.path import dirname, basename, isfile, join -import glob - -excluded_files = [ - "__init__.py", - "nn_passes.py", -] -modules = glob.glob(join(dirname(__file__), "*.py")) -pass_modules = [ - basename(f)[:-3] - for f in modules - if isfile(f) - and basename(f)[:1] != "_" # Follow python convention to hide _* files. - and basename(f)[:4] != "test" - and basename(f) not in excluded_files -] -__all__ = pass_modules - -from . import * # import everything in __all__ +from . import ( + alert_return_type_cast, + commingle_loop_vars, + handle_return_inputs_as_outputs, + handle_return_unused_inputs, + handle_unused_inputs, + mlmodel_passes +) diff --git a/coremltools/converters/mil/backend/nn/passes/alert_return_type_cast.py b/coremltools/converters/mil/backend/nn/passes/alert_return_type_cast.py index 86a102f37..343e72259 100644 --- a/coremltools/converters/mil/backend/nn/passes/alert_return_type_cast.py +++ b/coremltools/converters/mil/backend/nn/passes/alert_return_type_cast.py @@ -7,12 +7,13 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Var, types import logging @register_pass(namespace="nn_backend") -def alert_return_type_cast(prog): +class alert_return_type_cast(AbstractGraphPass): """ prog: Program @@ -38,12 +39,13 @@ def alert_return_type_cast(prog): # # Comment: This pass should do more proper casting as backend supports more types. """ - for f_name, f in prog.functions.items(): - for v in f.outputs: - if isinstance(v, Var) and v.dtype != types.fp32: - msg = ( - "Output var {} of type {} in function {} is " + "cast to type fp32" - ) - logging.warning( - msg.format(v.name, types.builtin_to_string(v.dtype), f_name) - ) + def apply(self, prog): + for f_name, f in prog.functions.items(): + for v in f.outputs: + if isinstance(v, Var) and v.dtype != types.fp32: + msg = ( + "Output var {} of type {} in function {} is " + "cast to type fp32" + ) + logging.warning( + msg.format(v.name, types.builtin_to_string(v.dtype), f_name) + ) diff --git a/coremltools/converters/mil/backend/nn/passes/commingle_loop_vars.py b/coremltools/converters/mil/backend/nn/passes/commingle_loop_vars.py index ee26863dc..b82767d9d 100644 --- a/coremltools/converters/mil/backend/nn/passes/commingle_loop_vars.py +++ b/coremltools/converters/mil/backend/nn/passes/commingle_loop_vars.py @@ -7,12 +7,12 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass - -def commingle_loop_vars_block(block): +def _commingle_loop_vars_block(block): for op in list(block.operations): for b in op.blocks: - commingle_loop_vars_block(b) + _commingle_loop_vars_block(b) if op.op_type != "while_loop": continue @@ -30,9 +30,8 @@ def commingle_loop_vars_block(block): # replace block inputs block._block_inputs = op.outputs - @register_pass(namespace="nn_backend") -def commingle_loop_vars(prog): +class commingle_loop_vars(AbstractGraphPass): """ prog: Program @@ -71,5 +70,6 @@ def commingle_loop_vars(prog): # Comment: The resulting program is no longer SSA (multiple assignments on # %loop:0). """ - for f_name, f in prog.functions.items(): - commingle_loop_vars_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _commingle_loop_vars_block(f) diff --git a/coremltools/converters/mil/backend/nn/passes/handle_return_inputs_as_outputs.py b/coremltools/converters/mil/backend/nn/passes/handle_return_inputs_as_outputs.py index 9aa0b069d..232315c3c 100644 --- a/coremltools/converters/mil/backend/nn/passes/handle_return_inputs_as_outputs.py +++ b/coremltools/converters/mil/backend/nn/passes/handle_return_inputs_as_outputs.py @@ -8,9 +8,9 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass - -def handle_return_inputs_as_outputs_func(f): +def _handle_return_inputs_as_outputs_func(f): returned_inputs = [] for v_name, v in f.inputs.items(): if v not in f.outputs: @@ -26,9 +26,8 @@ def handle_return_inputs_as_outputs_func(f): anchor_op=res.op, old_var=v, new_var=res ) - @register_pass(namespace="nn_backend") -def handle_return_inputs_as_outputs(prog): +class handle_return_inputs_as_outputs(AbstractGraphPass): """ prog: Program @@ -60,5 +59,6 @@ def handle_return_inputs_as_outputs(prog): # where identity is applied twice since NN layer cannot have # input name == output name """ - for f_name, f in prog.functions.items(): - handle_return_inputs_as_outputs_func(f) + def apply(self, prog): + for f in prog.functions.values(): + _handle_return_inputs_as_outputs_func(f) diff --git a/coremltools/converters/mil/backend/nn/passes/handle_return_unused_inputs.py b/coremltools/converters/mil/backend/nn/passes/handle_return_unused_inputs.py index 0e6dd65d5..9b741eee4 100644 --- a/coremltools/converters/mil/backend/nn/passes/handle_return_unused_inputs.py +++ b/coremltools/converters/mil/backend/nn/passes/handle_return_unused_inputs.py @@ -8,14 +8,11 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +def _handle_return_unused_inputs_func(f): -def handle_return_unused_inputs_func(f): - returned_unused_inputs = [] - for v_name, v in f.inputs.items(): - if v not in f.outputs: - continue - returned_unused_inputs.append(v) + returned_unused_inputs = filter(lambda x: x in f.outputs, list(f.inputs.values())) with f: for v in returned_unused_inputs: @@ -26,9 +23,8 @@ def handle_return_unused_inputs_func(f): anchor_op=res.op, old_var=v, new_var=res ) - @register_pass(namespace="nn_backend") -def handle_return_unused_inputs(prog): +class handle_return_unused_inputs(AbstractGraphPass): """ prog: Program @@ -60,5 +56,6 @@ def handle_return_unused_inputs(prog): # where identity is applied twice since NN layer cannot have # input name == output name """ - for f_name, f in prog.functions.items(): - handle_return_unused_inputs_func(f) + def apply(self, prog): + for f in prog.functions.values(): + _handle_return_unused_inputs_func(f) diff --git a/coremltools/converters/mil/backend/nn/passes/handle_unused_inputs.py b/coremltools/converters/mil/backend/nn/passes/handle_unused_inputs.py index 9ac1eff93..a49922d7d 100644 --- a/coremltools/converters/mil/backend/nn/passes/handle_unused_inputs.py +++ b/coremltools/converters/mil/backend/nn/passes/handle_unused_inputs.py @@ -8,19 +8,18 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass - -def handle_unused_inputs_func(f): +def _handle_unused_inputs_func(f): unused_inputs = [v for v_name, v in f.inputs.items() if len(v.child_ops) == 0] with f: for v in unused_inputs: - # copy twice since NN layer cannot have input name == output name + # copy the input v_tmp = mb.identity(x=v, name=v.name + "_tmp") - @register_pass(namespace="nn_backend") -def handle_unused_inputs(prog): +class handle_unused_inputs(AbstractGraphPass): """ prog: Program @@ -47,5 +46,6 @@ def handle_unused_inputs(prog): # } -> (%shape_0_const) # } """ - for f_name, f in prog.functions.items(): - handle_unused_inputs_func(f) + def apply(self, prog): + for f in prog.functions.values(): + _handle_unused_inputs_func(f) diff --git a/coremltools/converters/mil/backend/nn/passes/nn_passes.py b/coremltools/converters/mil/backend/nn/passes/nn_passes.py index 3e9f21a9e..83bde56db 100644 --- a/coremltools/converters/mil/backend/nn/passes/nn_passes.py +++ b/coremltools/converters/mil/backend/nn/passes/nn_passes.py @@ -3,9 +3,10 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY import logging +from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY + def nn_backend_passes(prog): passes = [ diff --git a/coremltools/converters/mil/converter.py b/coremltools/converters/mil/converter.py index 69076ecae..1d4c573ab 100644 --- a/coremltools/converters/mil/converter.py +++ b/coremltools/converters/mil/converter.py @@ -7,6 +7,7 @@ from coremltools.converters._profile_utils import _profile from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.quantization_passes import AbstractQuantizationPass +from coremltools.converters.mil.mil.types.symbolic import k_used_symbols, k_num_internal_syms from coremltools.models import MLModel from coremltools.models.model import _MODEL_FILE_NAME, _WEIGHTS_DIR_NAME from coremltools.models.utils import _MLMODEL_EXTENSION, _MLPACKAGE_EXTENSION @@ -133,6 +134,12 @@ def _reset_conversion_state(): # which is used to generate unique op names in the mil builder class. mb.name_count.clear() + # Clear "k_used_symbols" dict, and the int counter "k_num_internal_syms" that are used to track symbolic names + global k_used_symbols + global k_num_internal_syms + k_used_symbols.clear() + k_num_internal_syms = 0 + @_profile def mil_convert( model, diff --git a/coremltools/converters/mil/experimental/passes/generic_conv_batchnorm_fusion.py b/coremltools/converters/mil/experimental/passes/generic_conv_batchnorm_fusion.py index d17cbb2a0..8b40aa2ee 100644 --- a/coremltools/converters/mil/experimental/passes/generic_conv_batchnorm_fusion.py +++ b/coremltools/converters/mil/experimental/passes/generic_conv_batchnorm_fusion.py @@ -34,18 +34,19 @@ arbitrary_mean= np.random.rand(arbitrary_cout) arbitrary_variance = np.random.rand(arbitrary_cout) -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) -def conv_batchnorm(x): - conv = mb.conv(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") - batch_norm = mb.batch_norm(x=conv, mean=arbitrary_mean, variance=arbitrary_variance, name="batchnorm") - return batch_norm - - -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) -def conv_transpose_batchorm(x): - conv = mb.conv_transpose(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") - batch_norm = mb.batch_norm(x=conv, mean=arbitrary_mean, variance=arbitrary_variance, name="batchnorm") - return batch_norm +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) + def conv_batchnorm(x): + conv = mb.conv(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") + batch_norm = mb.batch_norm(x=conv, mean=arbitrary_mean, variance=arbitrary_variance, name="batchnorm") + return batch_norm + +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) + def conv_transpose_batchorm(x): + conv = mb.conv_transpose(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") + batch_norm = mb.batch_norm(x=conv, mean=arbitrary_mean, variance=arbitrary_variance, name="batchnorm") + return batch_norm def var_constraints(pattern): diff --git a/coremltools/converters/mil/experimental/passes/generic_conv_scale_fusion.py b/coremltools/converters/mil/experimental/passes/generic_conv_scale_fusion.py index 65dcc391d..d0da48e00 100644 --- a/coremltools/converters/mil/experimental/passes/generic_conv_scale_fusion.py +++ b/coremltools/converters/mil/experimental/passes/generic_conv_scale_fusion.py @@ -36,32 +36,33 @@ arbitrary_input = (3, arbitrary_cin, 224, 224) arbitrary_weight = np.random.rand(arbitrary_cout, arbitrary_cin, 10, 10) -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) -def conv_scale_mul(x): - conv = mb.conv(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") - mul = mb.mul(x=conv, y=arbitrary_scalar, name="scale") - return mul - - -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) -def conv_transpose_scale_mul(x): - conv = mb.conv_transpose(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") - mul = mb.mul(x=conv, y=arbitrary_scalar, name="scale") - return mul - +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) + def conv_scale_mul(x): + conv = mb.conv(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") + mul = mb.mul(x=conv, y=arbitrary_scalar, name="scale") + return mul -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) -def conv_scale_div(x): - conv = mb.conv(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") - real_div = mb.real_div(x=conv, y=arbitrary_scalar, name="scale") - return real_div +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) + def conv_transpose_scale_mul(x): + conv = mb.conv_transpose(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") + mul = mb.mul(x=conv, y=arbitrary_scalar, name="scale") + return mul +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) + def conv_scale_div(x): + conv = mb.conv(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") + real_div = mb.real_div(x=conv, y=arbitrary_scalar, name="scale") + return real_div -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) -def conv_transpose_scale_div(x): - conv = mb.conv_transpose(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") - real_div = mb.real_div(x=conv, y=arbitrary_scalar, name="scale") - return real_div +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_input)]) + def conv_transpose_scale_div(x): + conv = mb.conv_transpose(x=x, weight=arbitrary_weight, pad_type="valid", name="conv") + real_div = mb.real_div(x=conv, y=arbitrary_scalar, name="scale") + return real_div def _cin_cout(pattern): diff --git a/coremltools/converters/mil/experimental/passes/generic_gelu_tanh_approximation_fusion.py b/coremltools/converters/mil/experimental/passes/generic_gelu_tanh_approximation_fusion.py index 1f6aacb86..9671262f8 100644 --- a/coremltools/converters/mil/experimental/passes/generic_gelu_tanh_approximation_fusion.py +++ b/coremltools/converters/mil/experimental/passes/generic_gelu_tanh_approximation_fusion.py @@ -11,55 +11,57 @@ from coremltools.converters.mil.mil.passes.helper import _check_var_scalar_value from coremltools.converters.mil.experimental.passes.generic_pass_infrastructure import register_generic_pass -@mb.program(input_specs=[mb.TensorSpec(shape=([1, 1024, 4096])), ]) -def gelu_to_detect_1(x): - # MIL operation takes named inputs (instead of positional inputs). - # Here `name` argument is MANDATORY. - pow = mb.pow(x=x, y=3.0, name="pow") - mul_1 = mb.mul(x=0.044714998453855515, y=pow, name="mul_1") - add = mb.add(x=x, y=mul_1, name="add") - mul_2 = mb.mul(x=0.7978845834732056, y=add, name="mul_2") - tanh = mb.tanh(x=mul_2, name="tanh") - add_1 = mb.add(x=1.0, y=tanh, name="add_1") - mul = mb.mul(x=0.5, y=add_1, name="mul") - mul_3 = mb.mul(x=mul, y=x, name="mul_3") - return mul_3 -""" -y = x * (0.5 * (tanh(((.0447)x^3 + x ) * sqrt(2/pi)) + 1)) - - -[...] -----> pow (3) ----> mul (.044715) ---> add -----> mul (sqrt(2/pi)) ---> tanh ----> add (1) ----> mul (0.5) -----> mul ---> [...] - | ^ ^ - | | | - |------------------------------------------------------------------------------------------------------------------------ - -""" - -# In this pattern, 0.5 is first multiplied with the input which is then multiplied with the tanh term. -# In pattern1, 0.5 is first multiplied with the tanh term, and then multiplied with input -@mb.program(input_specs=[mb.TensorSpec(shape=([1, 1024, 4096])), ]) -def gelu_to_detect_2(x): - pow = mb.pow(x=x, y=3.0, name ="pow") - mul_1 = mb.mul(x=0.044714998453855515, y=pow, name="mul_1") - add = mb.add(x=x, y=mul_1, name="add") - mul_2 = mb.mul(x=0.7978845834732056, y=add, name="mul_2") - tanh = mb.tanh(x=mul_2, name="tanh") - add_1 = mb.add(x=1.0, y=tanh, name="add_1") - mul = mb.mul(x = 0.5, y=x, name="mul") - mul_3 = mb.mul(x=mul, y=add_1, name="mul_3") - return mul_3 - -""" -y = (0.5 * x) * (tanh(((.0447)x^3 + x ) * sqrt(2/pi)) + 1) - - --------------------------------------------------------------------------------------------------------- - ^ | - | V - [...] -----> mul(0.5) pow (3) ----> mul (.044715) ---> add -----> mul (sqrt(2/pi)) ---> tanh ----> add (1) -----> mul ---> [...] - | ^ ^ - | | | - |------------------------------------------------------------ -""" +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=([1, 1024, 4096])), ]) + def gelu_to_detect_1(x): + # MIL operation takes named inputs (instead of positional inputs). + # Here `name` argument is MANDATORY. + pow = mb.pow(x=x, y=3.0, name="pow") + mul_1 = mb.mul(x=0.044714998453855515, y=pow, name="mul_1") + add = mb.add(x=x, y=mul_1, name="add") + mul_2 = mb.mul(x=0.7978845834732056, y=add, name="mul_2") + tanh = mb.tanh(x=mul_2, name="tanh") + add_1 = mb.add(x=1.0, y=tanh, name="add_1") + mul = mb.mul(x=0.5, y=add_1, name="mul") + mul_3 = mb.mul(x=mul, y=x, name="mul_3") + return mul_3 + """ + y = x * (0.5 * (tanh(((.0447)x^3 + x ) * sqrt(2/pi)) + 1)) + + + [...] -----> pow (3) ----> mul (.044715) ---> add -----> mul (sqrt(2/pi)) ---> tanh ----> add (1) ----> mul (0.5) -----> mul ---> [...] + | ^ ^ + | | | + |------------------------------------------------------------------------------------------------------------------------ + + """ + +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + # In this pattern, 0.5 is first multiplied with the input which is then multiplied with the tanh term. + # In pattern1, 0.5 is first multiplied with the tanh term, and then multiplied with input + @mb.program(input_specs=[mb.TensorSpec(shape=([1, 1024, 4096])), ]) + def gelu_to_detect_2(x): + pow = mb.pow(x=x, y=3.0, name ="pow") + mul_1 = mb.mul(x=0.044714998453855515, y=pow, name="mul_1") + add = mb.add(x=x, y=mul_1, name="add") + mul_2 = mb.mul(x=0.7978845834732056, y=add, name="mul_2") + tanh = mb.tanh(x=mul_2, name="tanh") + add_1 = mb.add(x=1.0, y=tanh, name="add_1") + mul = mb.mul(x = 0.5, y=x, name="mul") + mul_3 = mb.mul(x=mul, y=add_1, name="mul_3") + return mul_3 + + """ + y = (0.5 * x) * (tanh(((.0447)x^3 + x ) * sqrt(2/pi)) + 1) + + --------------------------------------------------------------------------------------------------------- + ^ | + | V + [...] -----> mul(0.5) pow (3) ----> mul (.044715) ---> add -----> mul (sqrt(2/pi)) ---> tanh ----> add (1) -----> mul ---> [...] + | ^ ^ + | | | + |------------------------------------------------------------ + """ def var_constraints(pattern): passed = True diff --git a/coremltools/converters/mil/experimental/passes/generic_layernorm_instancenorm_pattern_fusion.py b/coremltools/converters/mil/experimental/passes/generic_layernorm_instancenorm_pattern_fusion.py index 39b41092d..959f7f157 100644 --- a/coremltools/converters/mil/experimental/passes/generic_layernorm_instancenorm_pattern_fusion.py +++ b/coremltools/converters/mil/experimental/passes/generic_layernorm_instancenorm_pattern_fusion.py @@ -14,7 +14,8 @@ from coremltools.converters.mil.experimental.passes.generic_pass_infrastructure import register_generic_pass from coremltools.converters.mil.mil import get_new_symbol -shape = (get_new_symbol(), get_new_symbol(), get_new_symbol(), get_new_symbol()) +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + shape = (get_new_symbol(), get_new_symbol(), get_new_symbol(), get_new_symbol()) def _check_reduce_op(reduce_op, mode="reduce_mean") -> bool: """ @@ -36,179 +37,181 @@ def _check_reduce_op(reduce_op, mode="reduce_mean") -> bool: return False return True +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=shape)]) + def instancenorm_or_layernorm(x): + """ + Identify the pattern: + + y = gamma * (x - mean) / sqrt(variance + epsilon) + beta + + y = x * [gamma * rsqrt(variance + eps)] + (beta - mean * [gamma * rsqrt(variance + eps)]) + + x --> main_reduce --> sub --> square --> reduce_mean_2 --> add(epsilon) --> rsqrt + | | ^ | + | | | V + |----------------------- mul (gamma) + | | | + | | --------|--------- + | | | | + | | | V + | |------------------------------------------------------------------> mul_3 + | | | + | V | + |----------------------------------------------------------------> mul_2 | + | V + | sub (beta) --> add_2 --> [...] + | ^ + |------------------------------- + + This pattern corresponds to either layer_norm or instance_norm. + + It is instance_norm if all of the following are true: + - input is rank 4 + - axes of reduce_mean is [-2, -1] or [-3, -2] + (when [-3, -2], a channel first to channel last transpose would be inserted) + - gamma and beta are rank 1, after squeeze + + It is layer_norm if all of the following are true: + - axes is either [-1] or [-1, -2] or [-1, -2, -3] and so on + - rank of gamma and beta is equal to the length of the axes + """ + main_reduce = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") + sub = mb.sub(x=x, y=main_reduce, name="sub") + square = mb.square(x=sub, name="square") + reduce_mean_2 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True, name="reduce_mean_2") + add_epsilon = mb.add(x=reduce_mean_2, y=1e-5, name="add_epsilon") + rsqrt = mb.rsqrt(x=add_epsilon, epsilon=1e-12, name="rsqrt") + mul_gamma = mb.mul(x=rsqrt, y=np.random.rand(1, 5, 1, 1), name="mul_gamma") + mul_2 = mb.mul(x=x, y=mul_gamma, name="mul_2") + mul_3 = mb.mul(x=main_reduce, y=mul_gamma, name="mul_3") + sub_beta = mb.sub(x=np.random.rand(1, 5, 1, 1), y=mul_3, name="sub_beta") + add_2 = mb.add(x=sub_beta, y=mul_2, name="add_2") + return add_2 -@mb.program(input_specs=[mb.TensorSpec(shape=shape)]) -def instancenorm_or_layernorm(x): - """ - Identify the pattern: - - y = gamma * (x - mean) / sqrt(variance + epsilon) + beta - - y = x * [gamma * rsqrt(variance + eps)] + (beta - mean * [gamma * rsqrt(variance + eps)]) - - x --> main_reduce --> sub --> square --> reduce_mean_2 --> add(epsilon) --> rsqrt - | | ^ | - | | | V - |----------------------- mul (gamma) - | | | - | | --------|--------- - | | | | - | | | V - | |------------------------------------------------------------------> mul_3 - | | | - | V | - |----------------------------------------------------------------> mul_2 | - | V - | sub (beta) --> add_2 --> [...] - | ^ - |------------------------------- - - This pattern corresponds to either layer_norm or instance_norm. - - It is instance_norm if all of the following are true: - - input is rank 4 - - axes of reduce_mean is [-2, -1] or [-3, -2] - (when [-3, -2], a channel first to channel last transpose would be inserted) - - gamma and beta are rank 1, after squeeze - - It is layer_norm if all of the following are true: - - axes is either [-1] or [-1, -2] or [-1, -2, -3] and so on - - rank of gamma and beta is equal to the length of the axes - """ - main_reduce = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") - sub = mb.sub(x=x, y=main_reduce, name="sub") - square = mb.square(x=sub, name="square") - reduce_mean_2 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True, name="reduce_mean_2") - add_epsilon = mb.add(x=reduce_mean_2, y=1e-5, name="add_epsilon") - rsqrt = mb.rsqrt(x=add_epsilon, epsilon=1e-12, name="rsqrt") - mul_gamma = mb.mul(x=rsqrt, y=np.random.rand(1, 5, 1, 1), name="mul_gamma") - mul_2 = mb.mul(x=x, y=mul_gamma, name="mul_2") - mul_3 = mb.mul(x=main_reduce, y=mul_gamma, name="mul_3") - sub_beta = mb.sub(x=np.random.rand(1, 5, 1, 1), y=mul_3, name="sub_beta") - add_2 = mb.add(x=sub_beta, y=mul_2, name="add_2") - return add_2 - - -@mb.program(input_specs=[mb.TensorSpec(shape=shape)]) -def instancenorm_2(x): - """ - Identify the pattern: - y = (x - mean) / pow(variance + epsilon) * gamma + beta - - This pattern corresponds to, should be fused as instance_norm. - All of the following must be satisty: - 1) Input is rank 4 tensor - 2) Reduce operates on spatial dimensions axes=[-2, -1], or axes=[-3, -2] (a - channel first to channel last transpose would be inserted in such case) - 3) Gamma and beta are both shape (C,) after squeeze, where C is number of channels - - - |----> sub0 ----------| const (0.5) - | ^ | | - | | V V - x ---> main_reduce square --> mean1 --> add_eps ---> pow const_gamma const_beta - | | | | | - | V V V V - |----> sub1 --------------------------------------> real_div --> mul_gamma --> add_beta --> ... - """ +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=shape)]) + def instancenorm_2(x): + """ + Identify the pattern: + y = (x - mean) / pow(variance + epsilon) * gamma + beta + + This pattern corresponds to, should be fused as instance_norm. + All of the following must be satisty: + 1) Input is rank 4 tensor + 2) Reduce operates on spatial dimensions axes=[-2, -1], or axes=[-3, -2] (a + channel first to channel last transpose would be inserted in such case) + 3) Gamma and beta are both shape (C,) after squeeze, where C is number of channels + + + |----> sub0 ----------| const (0.5) + | ^ | | + | | V V + x ---> main_reduce square --> mean1 --> add_eps ---> pow const_gamma const_beta + | | | | | + | V V V V + |----> sub1 --------------------------------------> real_div --> mul_gamma --> add_beta --> ... + """ + + main_reduce = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") + sub0 = mb.sub(x=x, y=main_reduce, name="sub0") + sub1 = mb.sub(x=x, y=main_reduce, name="sub1") + square = mb.square(x=sub0, name="square") + mean1 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True, name="mean1") + add_epsilon = mb.add(x=mean1, y=1e-5, name="add_epsilon") + pow = mb.pow(x=add_epsilon, y=0.5, name="pow") + real_div = mb.real_div(x=sub1, y=pow, name="real_div") + mul_gamma = mb.mul(x=np.random.rand(1, 5, 1, 1), y=real_div, name="mul_gamma") + add_beta = mb.add(x=np.random.rand(1, 5, 1, 1), y=mul_gamma, name="add_beta") + return add_beta - main_reduce = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") - sub0 = mb.sub(x=x, y=main_reduce, name="sub0") - sub1 = mb.sub(x=x, y=main_reduce, name="sub1") - square = mb.square(x=sub0, name="square") - mean1 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True, name="mean1") - add_epsilon = mb.add(x=mean1, y=1e-5, name="add_epsilon") - pow = mb.pow(x=add_epsilon, y=0.5, name="pow") - real_div = mb.real_div(x=sub1, y=pow, name="real_div") - mul_gamma = mb.mul(x=np.random.rand(1, 5, 1, 1), y=real_div, name="mul_gamma") - add_beta = mb.add(x=np.random.rand(1, 5, 1, 1), y=mul_gamma, name="add_beta") - return add_beta - - -@mb.program(input_specs=[mb.TensorSpec(shape=shape)]) -def instancenorm_3(x): - """ - Detect InstanceNorm pattern in TensorFlow-Addons. - - This pattern corresponds to, should be fused as instance_norm. - All of the following must be satisty: - 1) Input is rank 4 tensor - 2) Reduce operates on spatial dimensions axes=[-2, -1], or axes=[-3, -2] (a - channel first to channel last transpose would be inserted in such case) - 3) Gamma and beta are absent. Default values for gamma and beta would be used. - - |-------------------------------------------------------| - | | - | V - x --> main_reduce square --> mean1 --> add_eps --> rsqrt --> mul2 --> mul_sub - | | ^ | | - | V | | | - | --> sub -----------| | | - | V V - |--------------------------------------------------> mul1 -------------> add --> ... - """ - main_reduce = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") - sub = mb.sub(x=x, y=main_reduce, name="sub") - square = mb.square(x=sub, name="square") - mean1 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True, name="mean1") - add_epsilon = mb.add(x=mean1, y=1e-5, name="add_epsilon") # epsilon - rsqrt = mb.rsqrt(x=add_epsilon, name="rsqrt") - mul1 = mb.mul(x=rsqrt, y=x, name="mul1") - mul2 = mb.mul(x=main_reduce, y=rsqrt, name="mul2") - mul_sub = mb.mul(x=mul2, y=-1, name="mul_sub") - add = mb.add(x=mul1, y=mul_sub, name="add") - return add - - -@mb.program(input_specs=[mb.TensorSpec(shape=shape)]) -def instancenorm_4(x): - """ - Identify the pattern: - y = x * [gamma * rsqrt(variance + eps)] + (beta - mean * [gamma * rsqrt(variance + eps)]) - - This pattern corresponds to, should be fused as instance_norm. - All of the following must be satisty: - 1) Input is rank 4 tensor - 2) Reduce operates on spatial dimensions axes=[-2, -1], or axes=[-3, -2] (a - channel first to channel last transpose would be inserted in such case) - 3) Gamma and beta are both shape (C,) after squeeze, where C is number of channels - - |-----------| - | V - |------> mul_square1 -------------> sum1 -----> mul_mean1 - | | - | V - x --> main_reduce --> mul_mean ==> mul_square --> sub_variance --> add_eps --> rsqrt - | | | - | | V - | | mul_gamma - | | | - | | |----------------| - | | | V - | |--------------------------------------------+-------------> mul2 - | V | - |------------------------------------------------------------------> mul1 | - | V - | sub_beta --> add --> [...] - | ^ - |---------------------------| - """ - mul_square1 = mb.mul(x=x, y=x, name="mul_square1") - main_reduce = mb.reduce_sum(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") - mul_mean = mb.mul(x=main_reduce, y=3.3333334e-05, name="mul_mean") # dummy value here - mul_square = mb.mul(x=mul_mean, y=mul_mean, name="mul_square") - sum1 = mb.reduce_sum(x=mul_square1, axes=[2, 3], keep_dims=True, name="sum1") - mul_mean1 = mb.mul(x=sum1, y=8.333333e-06, name="mul_mean1") # dummy value here - sub_variance = mb.sub(x=mul_mean1, y=mul_square, name="sub_variance") - add_epsilon = mb.add(x=sub_variance, y=1e-5, name="add_epsilon") # epsilon - rsqrt = mb.rsqrt(x=add_epsilon, name="rsqrt") - mul_gamma = mb.mul(x=rsqrt, y=np.random.rand(1, 5, 1, 1), name="mul_gamma") - mul1 = mb.mul(x=mul_gamma, y=x, name="mul1") - mul2 = mb.mul(x=mul_mean, y=mul_gamma, name="mul2") - sub_beta = mb.sub(x=np.random.rand(1, 5, 1, 1), y=mul2, name="sub_beta") - add = mb.add(x=mul1, y=sub_beta, name="add") - return add +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=shape)]) + def instancenorm_3(x): + """ + Detect InstanceNorm pattern in TensorFlow-Addons. + + This pattern corresponds to, should be fused as instance_norm. + All of the following must be satisty: + 1) Input is rank 4 tensor + 2) Reduce operates on spatial dimensions axes=[-2, -1], or axes=[-3, -2] (a + channel first to channel last transpose would be inserted in such case) + 3) Gamma and beta are absent. Default values for gamma and beta would be used. + + |-------------------------------------------------------| + | | + | V + x --> main_reduce square --> mean1 --> add_eps --> rsqrt --> mul2 --> mul_sub + | | ^ | | + | V | | | + | --> sub -----------| | | + | V V + |--------------------------------------------------> mul1 -------------> add --> ... + """ + + main_reduce = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") + sub = mb.sub(x=x, y=main_reduce, name="sub") + square = mb.square(x=sub, name="square") + mean1 = mb.reduce_mean(x=square, axes=[2, 3], keep_dims=True, name="mean1") + add_epsilon = mb.add(x=mean1, y=1e-5, name="add_epsilon") # epsilon + rsqrt = mb.rsqrt(x=add_epsilon, name="rsqrt") + mul1 = mb.mul(x=rsqrt, y=x, name="mul1") + mul2 = mb.mul(x=main_reduce, y=rsqrt, name="mul2") + mul_sub = mb.mul(x=mul2, y=-1, name="mul_sub") + add = mb.add(x=mul1, y=mul_sub, name="add") + return add + + +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=shape)]) + def instancenorm_4(x): + """ + Identify the pattern: + y = x * [gamma * rsqrt(variance + eps)] + (beta - mean * [gamma * rsqrt(variance + eps)]) + + This pattern corresponds to, should be fused as instance_norm. + All of the following must be satisty: + 1) Input is rank 4 tensor + 2) Reduce operates on spatial dimensions axes=[-2, -1], or axes=[-3, -2] (a + channel first to channel last transpose would be inserted in such case) + 3) Gamma and beta are both shape (C,) after squeeze, where C is number of channels + + |-----------| + | V + |------> mul_square1 -------------> sum1 -----> mul_mean1 + | | + | V + x --> main_reduce --> mul_mean ==> mul_square --> sub_variance --> add_eps --> rsqrt + | | | + | | V + | | mul_gamma + | | | + | | |----------------| + | | | V + | |--------------------------------------------+-------------> mul2 + | V | + |------------------------------------------------------------------> mul1 | + | V + | sub_beta --> add --> [...] + | ^ + |---------------------------| + """ + mul_square1 = mb.mul(x=x, y=x, name="mul_square1") + main_reduce = mb.reduce_sum(x=x, axes=[2, 3], keep_dims=True, name="main_reduce") + mul_mean = mb.mul(x=main_reduce, y=3.3333334e-05, name="mul_mean") # dummy value here + mul_square = mb.mul(x=mul_mean, y=mul_mean, name="mul_square") + sum1 = mb.reduce_sum(x=mul_square1, axes=[2, 3], keep_dims=True, name="sum1") + mul_mean1 = mb.mul(x=sum1, y=8.333333e-06, name="mul_mean1") # dummy value here + sub_variance = mb.sub(x=mul_mean1, y=mul_square, name="sub_variance") + add_epsilon = mb.add(x=sub_variance, y=1e-5, name="add_epsilon") # epsilon + rsqrt = mb.rsqrt(x=add_epsilon, name="rsqrt") + mul_gamma = mb.mul(x=rsqrt, y=np.random.rand(1, 5, 1, 1), name="mul_gamma") + mul1 = mb.mul(x=mul_gamma, y=x, name="mul1") + mul2 = mb.mul(x=mul_mean, y=mul_gamma, name="mul2") + sub_beta = mb.sub(x=np.random.rand(1, 5, 1, 1), y=mul2, name="sub_beta") + add = mb.add(x=mul1, y=sub_beta, name="add") + return add def instancenorm_1_constraints(pattern): passed = True diff --git a/coremltools/converters/mil/experimental/passes/generic_linear_bias_fusion.py b/coremltools/converters/mil/experimental/passes/generic_linear_bias_fusion.py index d0f18e9bc..afe6ada7e 100644 --- a/coremltools/converters/mil/experimental/passes/generic_linear_bias_fusion.py +++ b/coremltools/converters/mil/experimental/passes/generic_linear_bias_fusion.py @@ -12,45 +12,48 @@ from coremltools.converters.mil.experimental.passes.generic_pass_infrastructure import register_generic_pass from coremltools.converters.mil.mil import get_new_symbol -arbitrary_shape = (get_new_symbol(), get_new_symbol()) -np.random.seed() -arbitrary_weight = np.random.rand(4,3) -arbitrary_bias = np.random.rand(4) - -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_shape)]) -def pattern_add(x): - """ - Original: - % 4 = linear(x= % 1, weight = % 2, bias = % 3) # %2 is a rank-2 const tensor (weight) - # %3 is a rank-1 const tensor (bias) - ... - % 6 = add(x= % 4, y = % 5) # %5 is a const tensor with same shape as %3 - - Result: - % 8 = linear(x= % 1, weight = % 2, bias = % 7) # where %7 is a new const tensor with value - # %7 = %3 + %6 - """ - linear = mb.linear(x=x, weight=arbitrary_weight, bias=arbitrary_bias, name="linear") - add_or_sub = mb.add(x=linear, y=arbitrary_bias, name="add_or_sub") - return add_or_sub - -@mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_shape)]) -def pattern_sub(x): - """ - Original: - %4 = linear(x=%1, weight=%2, bias=%3) # %2 is a rank-2 const tensor (weight) - # %3 is a rank-1 const tensor (bias) - ... - %6 = sub(x=%5, y=%4) # %5 is a const tensor with a broacasable shape with %3. - i.e. if %3 has shape (Dout), %5 could be (1, Dout). - - Result: - %9 = linear(x=%1, weight=%7, bias=%8) # where %7 is a new const tensor with value %7 = -%2 - # %8 = %5 - %3 - """ - linear = mb.linear(x=x, weight=arbitrary_weight, bias=arbitrary_bias, name="linear") - add_or_sub = mb.sub(x=linear, y=arbitrary_bias, name="add_or_sub") - return add_or_sub +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + arbitrary_shape = (get_new_symbol(), get_new_symbol()) + np.random.seed() + arbitrary_weight = np.random.rand(4,3) + arbitrary_bias = np.random.rand(4) + +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_shape)]) + def pattern_add(x): + """ + Original: + % 4 = linear(x= % 1, weight = % 2, bias = % 3) # %2 is a rank-2 const tensor (weight) + # %3 is a rank-1 const tensor (bias) + ... + % 6 = add(x= % 4, y = % 5) # %5 is a const tensor with same shape as %3 + + Result: + % 8 = linear(x= % 1, weight = % 2, bias = % 7) # where %7 is a new const tensor with value + # %7 = %3 + %6 + """ + linear = mb.linear(x=x, weight=arbitrary_weight, bias=arbitrary_bias, name="linear") + add_or_sub = mb.add(x=linear, y=arbitrary_bias, name="add_or_sub") + return add_or_sub + +if os.getenv("ENABLE_EXPERIMENTAL_PASSES") == "1": + @mb.program(input_specs=[mb.TensorSpec(shape=arbitrary_shape)]) + def pattern_sub(x): + """ + Original: + %4 = linear(x=%1, weight=%2, bias=%3) # %2 is a rank-2 const tensor (weight) + # %3 is a rank-1 const tensor (bias) + ... + %6 = sub(x=%5, y=%4) # %5 is a const tensor with a broacasable shape with %3. + i.e. if %3 has shape (Dout), %5 could be (1, Dout). + + Result: + %9 = linear(x=%1, weight=%7, bias=%8) # where %7 is a new const tensor with value %7 = -%2 + # %8 = %5 - %3 + """ + linear = mb.linear(x=x, weight=arbitrary_weight, bias=arbitrary_bias, name="linear") + add_or_sub = mb.sub(x=linear, y=arbitrary_bias, name="add_or_sub") + return add_or_sub def var_constraints(pattern): diff --git a/coremltools/converters/mil/frontend/_utils.py b/coremltools/converters/mil/frontend/_utils.py index 73164aebd..dec3891d8 100644 --- a/coremltools/converters/mil/frontend/_utils.py +++ b/coremltools/converters/mil/frontend/_utils.py @@ -3,8 +3,9 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil import Builder as mb, types from coremltools.converters.mil.mil.ops.defs._utils import parse_einsum_equation +from coremltools.converters.mil.mil.types.symbolic import any_symbolic, is_symbolic def _reverse_input_einsum_eq(equation): @@ -106,4 +107,41 @@ def _swap(a, b): "Einsum unsupported equation format: ", equation ) - return x \ No newline at end of file + return x + + +def is_symbolic_dim_in_prog(prog): + ''' + Takes in a MIL program object, checks if any of the tensors in it contain a symbolic dimension. + Returns true if it does. + + :param prog: coremltools.converters.mil.Program + :return: bool + ''' + def _does_block_contain_symbolic_shape(block): + for op in block.operations: + for b in op.blocks: + if _does_block_contain_symbolic_shape(b): + return True + for out in op.outputs: + if types.is_tensor(out.sym_type): + shape = out.sym_type.get_shape() + if any_symbolic(shape): + return True + elif types.is_scalar(out.sym_type) or types.is_str(out.sym_type): + if is_symbolic(out.val): + return True + elif types.is_list(out.sym_type): + if types.is_tensor(out.elem_type): + if any_symbolic(out.elem_type.get_shape()): + return True + else: + raise NotImplementedError("\'{}\' type in a list not handled".format(out.elem_type)) + else: + raise NotImplementedError("\'{}\' type is not handled".format(out.sym_type)) + return False + + for f in prog.functions.values(): + if _does_block_contain_symbolic_shape(f): + return True + return False diff --git a/coremltools/converters/mil/frontend/tensorflow/__init__.py b/coremltools/converters/mil/frontend/tensorflow/__init__.py index 19cca2c37..4a8a9dc05 100644 --- a/coremltools/converters/mil/frontend/tensorflow/__init__.py +++ b/coremltools/converters/mil/frontend/tensorflow/__init__.py @@ -16,6 +16,13 @@ register_tf_op = None if _HAS_TF_1: - from .ops import * # register all - from .dialect_ops import * # register tf extension ops + # Importing these causes them to register their ops + from . import ops + + from .dialect_ops import ( + tf_make_list, + TfLSTMBase, + tf_lstm_block_cell, + tf_lstm_block + ) from .tf_op_registry import register_tf_op diff --git a/coremltools/converters/mil/frontend/tensorflow/basic_graph_ops.py b/coremltools/converters/mil/frontend/tensorflow/basic_graph_ops.py index 5359554b4..ed8b3e06e 100644 --- a/coremltools/converters/mil/frontend/tensorflow/basic_graph_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/basic_graph_ops.py @@ -1,12 +1,9 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - def connect_edge(g, source, dest): if isinstance(source, str): source = g[source] diff --git a/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py b/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py index bdad2e8b2..7e9f34fb3 100644 --- a/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py @@ -3,9 +3,16 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil import Operation -from coremltools.converters.mil.mil.input_type import * +from coremltools.converters.mil.mil import types, Operation +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + DefaultInputs, + FloatInputType, + InputSpec, + IntInputType, + StringInputType, + TensorInputType, +) from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry register_op = SSAOpRegistry.register_op diff --git a/coremltools/converters/mil/frontend/tensorflow/load.py b/coremltools/converters/mil/frontend/tensorflow/load.py index 59613416b..18077f5dc 100644 --- a/coremltools/converters/mil/frontend/tensorflow/load.py +++ b/coremltools/converters/mil/frontend/tensorflow/load.py @@ -1,27 +1,34 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - +from distutils.version import StrictVersion as _StrictVersion +import gc import logging import os -import gc - +from tempfile import mktemp import tensorflow as tf +from tqdm import tqdm as _tqdm -from tempfile import mktemp from .basic_graph_ops import fill_outputs from .converter import TFConverter -from .tf_graph_pass import * # pylint: disable=unused-wildcard-import,wildcard-import +from .tf_graph_pass import ( + cond_to_where, + constant_propagation, + delete_asserts, + delete_disconnected_nodes, + functionalize_loops, + fuse_dilation_conv, + insert_get_tuple, + quantization_pass, + remove_variable_nodes, + tensor_array_resource_removal +) from .tfssa import NetworkEnsemble, SSAFunction from .parsed_tf_node import ParsedTFNode from coremltools.converters._profile_utils import _profile -from tqdm import tqdm as _tqdm -from distutils.version import StrictVersion as _StrictVersion -from coremltools._deps import __get_version as _get_version +from coremltools._deps import _get_version class TFLoader: diff --git a/coremltools/converters/mil/frontend/tensorflow/ops.py b/coremltools/converters/mil/frontend/tensorflow/ops.py index b4c354fe4..48ed72668 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/ops.py @@ -6,13 +6,13 @@ import logging as _logging import numpy as _np -from coremltools.converters.mil.mil import Builder as mb -from coremltools.converters.mil.mil.ops.defs._utils import broadcast_shapes +from .._utils import build_einsum_mil from .convert_utils import convert_graph from .tf_op_registry import register_tf_op -from coremltools.converters.mil.mil import types +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.mil.ops.defs._utils import broadcast_shapes from coremltools.converters.mil.mil.types.symbolic import is_symbolic, any_symbolic -from .._utils import build_einsum_mil + def _adjust_min_max(min, max, num_bits=8): if (min <= max) and (max <= 0): diff --git a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/__init__.py b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/__init__.py index a659bb2d3..f1215aa4c 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/__init__.py +++ b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/__init__.py @@ -3,23 +3,9 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -# Import all passes in this dir -from os.path import dirname, basename, isfile, join -import glob - -excluded_files = [ - "__init__.py", - "tf_passes.py", -] -modules = glob.glob(join(dirname(__file__), "*.py")) -pass_modules = [ - basename(f)[:-3] - for f in modules - if isfile(f) - and basename(f)[:1] != "_" # Follow python convention to hide _* files. - and basename(f)[:4] != "test" - and basename(f) not in excluded_files -] -__all__ = pass_modules - -from . import * # import everything in __all__ +from . import ( + backfill_make_list_elem_type, + expand_tf_lstm, + tf_lstm_to_core_lstm, + tf_passes +) diff --git a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/backfill_make_list_elem_type.py b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/backfill_make_list_elem_type.py index d259a7364..22b10682b 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/backfill_make_list_elem_type.py +++ b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/backfill_make_list_elem_type.py @@ -1,20 +1,17 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.passes.pass_registry import register_pass -from coremltools.converters.mil.mil import Builder as mb -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil.var import ListVar from coremltools.converters.mil.mil.types.symbolic import is_symbolic +from coremltools.converters.mil.mil.var import ListVar @register_pass(namespace="tensorflow") -def backfill_make_list_elem_type(prog): +class backfill_make_list_elem_type(AbstractGraphPass): """ TF's TensorArrayV3 (represented as make_list in mil) doesn't necessarily contain elem shape/type, which is known when write is performed. We @@ -24,26 +21,26 @@ def backfill_make_list_elem_type(prog): prog: Program """ - for f_name, f in prog.functions.items(): - backfill_make_list_elem_type_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _backfill_make_list_elem_type_block(f) -def backfill_make_list_elem_type_block(block): +def _backfill_make_list_elem_type_block(block): # shallow copy hides changes on f.operations during the loop - for op in block.operations[:]: + for op in block.operations: for b in op.blocks: - backfill_make_list_elem_type_block(b) + _backfill_make_list_elem_type_block(b) if op.op_type != "tf_make_list": continue - # op is `make_list` if op.outputs[0].elem_type != types.unknown: # elem_type of the list is known continue list_var = op.outputs[0] - elem_type = infer_elem_type(list_var) # types.tensor + elem_type = _infer_elem_type(list_var) # types.tensor if elem_type is None: msg = ( "No list_write or list_scatter op to infer make_list " @@ -71,7 +68,7 @@ def backfill_make_list_elem_type_block(block): block.remove_ops([op]) -def infer_elem_type(list_var): +def _infer_elem_type(list_var): """ Returns types.tensor. None if failed to infer element type. Example: @@ -105,7 +102,7 @@ def infer_elem_type(list_var): block = o.blocks[0] # the corresponding Var in body block block_var = block.inputs[idx] - elem_type = infer_elem_type(block_var) + elem_type = _infer_elem_type(block_var) if elem_type is not None: def _set_types_for_block_inputs(block): diff --git a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/expand_tf_lstm.py b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/expand_tf_lstm.py index a6bf77993..4c253d58c 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/expand_tf_lstm.py +++ b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/expand_tf_lstm.py @@ -7,14 +7,14 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import types import numpy as np import logging - @register_pass(namespace="tensorflow") -def expand_tf_lstm(prog): +class expand_tf_lstm(AbstractGraphPass): """ Expand tf_lstm_block_cell to fine-grained SSA ops following: @@ -36,30 +36,30 @@ def expand_tf_lstm(prog): prog: Program """ - for f_name, f in prog.functions.items(): - expand_tf_lstm_helper(f) + def apply(self, prog): + for f in prog.functions.values(): + _expand_tf_lstm_helper(f) -def expand_tf_lstm_helper(block): +def _expand_tf_lstm_helper(block): # shallow copy hides changes on f.operations during the loop for op in block.operations[:]: for b in op.blocks: - expand_tf_lstm_helper(b) + _expand_tf_lstm_helper(b) if op.op_type == "tf_lstm_block_cell": - expand_tf_lstm_block_cell(op) + _expand_tf_lstm_block_cell(op) logging.info("Expanding {} (op_type: {})".format(op.name, op.op_type)) if op.op_type == "tf_lstm_block": # only cs, h are supported for now. Can be easily extended to other outputs at performance hit. i, cs, f, o, ci, co, h = op.outputs if all( - [ - len(ov.child_ops) <= 0 and len(ov.consuming_blocks) <= 0 - for ov in [i, f, o, ci, co] - ] + map(lambda x: len(x.child_ops) == 0 and len(x.consuming_blocks) == 0, + (i, f, o, ci, co) + ) ): - expand_tf_lstm_block(op) + _expand_tf_lstm_block(op) logging.info("Expanding {} (op_type: {})".format(op.name, op.op_type)) @@ -81,6 +81,7 @@ def _lstm_cell_builder(op, x, h_prev, cs_prev, before_op=None): if op.forget_bias.val != 0: f = mb.add(x=f, y=forget_bias, before_op=before_op) + # note that .* means Hadamard product # i = sigmoid(cs_prev .* wci + i) # f = sigmoid(cs_prev .* wcf + f) if op.use_peephole.val: @@ -98,8 +99,6 @@ def _lstm_cell_builder(op, x, h_prev, cs_prev, before_op=None): i = mb.sigmoid(x=pre_i, before_op=before_op) f = mb.sigmoid(x=pre_f, before_op=before_op) - - # ci = tanh(ci) ci = mb.tanh(x=ci, before_op=before_op) # cs = ci .* i + cs_prev .* f @@ -120,8 +119,6 @@ def _lstm_cell_builder(op, x, h_prev, cs_prev, before_op=None): else: pre_o = o o = mb.sigmoid(x=pre_o, before_op=before_op) - - # co = tanh(cs) co = mb.tanh(x=cs, before_op=before_op) # h = co .* o @@ -130,7 +127,7 @@ def _lstm_cell_builder(op, x, h_prev, cs_prev, before_op=None): return [i, cs, f, o, ci, co, h] -def expand_tf_lstm_block_cell(op): +def _expand_tf_lstm_block_cell(op): if op.op_type != "tf_lstm_block_cell": raise ValueError() @@ -152,7 +149,7 @@ def expand_tf_lstm_block_cell(op): block.remove_ops([op]) -def expand_tf_lstm_block(op): +def _expand_tf_lstm_block(op): if op.op_type != "tf_lstm_block": raise ValueError() @@ -161,20 +158,7 @@ def expand_tf_lstm_block(op): h_prev = op.h_prev # [b, hidden_dim] cs_prev = op.c_prev # [b, hidden_dim] - # Allocate two lists: cs & h - x_shape = mb.shape(x=x, before_op=op) - length = mb.slice_by_index(x=x_shape, begin=[0], end=[1], before_op=op) - h_shape = mb.shape(x=h_prev, before_op=op) - list_shape = mb.concat(values=[length, h_shape], axis=0, before_op=op) - cs_list = mb.fill(shape=list_shape, before_op=op) - h_list = mb.fill(shape=list_shape, before_op=op) - - # append initial state at index 0 - cs_prev = mb.expand_dims(x=cs_prev, axes=[0], before_op=op) - cs_list = mb.concat(values=[cs_prev, cs_list], axis=0, before_op=op) - h_prev = mb.expand_dims(x=h_prev, axes=[0], before_op=op) - h_list = mb.concat(values=[h_prev, h_list], axis=0, before_op=op) - + # cond and body function gor the while_loop def cond(i, cs_list, h_list): return mb.less(x=i, y=length) @@ -193,6 +177,20 @@ def body(i, cs_list, h_list): mb.scatter(data=h_list, indices=counter, updates=h), ) + # Allocate two lists: cs & h + x_shape = mb.shape(x=x, before_op=op) + length = mb.slice_by_index(x=x_shape, begin=[0], end=[1], before_op=op) + h_shape = mb.shape(x=h_prev, before_op=op) + list_shape = mb.concat(values=[length, h_shape], axis=0, before_op=op) + cs_list = mb.fill(shape=list_shape, before_op=op) + h_list = mb.fill(shape=list_shape, before_op=op) + + # append initial state at index 0 + cs_prev = mb.expand_dims(x=cs_prev, axes=[0], before_op=op) + cs_list = mb.concat(values=[cs_prev, cs_list], axis=0, before_op=op) + h_prev = mb.expand_dims(x=h_prev, axes=[0], before_op=op) + h_list = mb.concat(values=[h_prev, h_list], axis=0, before_op=op) + _, cs_list, h_list = mb.while_loop( _cond=cond, _body=body, loop_vars=([0], cs_list, h_list), before_op=op ) diff --git a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/tf_lstm_to_core_lstm.py b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/tf_lstm_to_core_lstm.py index 1228763f7..e4ed522cf 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ssa_passes/tf_lstm_to_core_lstm.py +++ b/coremltools/converters/mil/frontend/tensorflow/ssa_passes/tf_lstm_to_core_lstm.py @@ -7,6 +7,7 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import types import numpy as np @@ -14,7 +15,7 @@ @register_pass(namespace="tensorflow") -def tf_lstm_to_core_lstm(prog): +class tf_lstm_to_core_lstm(AbstractGraphPass): """ Try to map TF dialect ops `tf_lstm_block` and `tf_lstm_block_cell` to `lstm` in the core op set if compatible. They are compatible if all of the @@ -36,24 +37,25 @@ def tf_lstm_to_core_lstm(prog): prog: Program """ - for f_name, f in prog.functions.items(): - tf_lstm_to_core_lstm_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _tf_lstm_to_core_lstm_block(f) -def tf_lstm_to_core_lstm_block(block): +def _tf_lstm_to_core_lstm_block(block): # shallow copy hides changes on f.operations during the loop - for op in block.operations[:]: + for op in block.operations: for b in op.blocks: - tf_lstm_to_core_lstm_block(b) + _tf_lstm_to_core_lstm_block(b) if op.op_type in ["tf_lstm_block_cell", "tf_lstm_block"]: - if try_replace_with_core_lstm(op): + if _try_replace_with_core_lstm(op): logging.info("Successfully map {} to lstm".format(op.op_type)) else: logging.info("Unable to map {} to lstm".format(op.op_type)) -def try_replace_with_core_lstm(op): +def _try_replace_with_core_lstm(op): """ Inputs: @@ -63,6 +65,13 @@ def try_replace_with_core_lstm(op): True if op can be represented by mb.lstm op in SSA. False otherwise """ + def _check_unsupported_outputs(unsupported_outputs): + for ov in unsupported_outputs: + if len(ov.child_ops) > 0 or len(ov.consuming_blocks) > 0: + return False + return True + + if op.op_type == "tf_lstm_block_cell": batch = op.x.shape[0] else: # tf_lstm_block @@ -76,16 +85,12 @@ def try_replace_with_core_lstm(op): i, cs, f, o, ci, co, h = op.outputs if op.op_type == "tf_lstm_block_cell": unsupported_outputs = [i, f, o, ci, co] # only cs, h are supported - for ov in unsupported_outputs: - if len(ov.child_ops) > 0 or len(ov.consuming_blocks) > 0: - return False else: # tf_lstm_block unsupported_outputs = [i, cs, f, o, ci, co] # only h is supported - for ov in unsupported_outputs: - if len(ov.child_ops) > 0 or len(ov.consuming_blocks) > 0: - return False - # op is compatible with lstm + if not _check_unsupported_outputs(unsupported_outputs): + return False + # op is compatible with lstm hidden_dim = op.c_prev.shape[1] mb_peep = None @@ -94,7 +99,8 @@ def try_replace_with_core_lstm(op): [op.weight_peep_i.val, op.weight_peep_f.val, op.weight_peep_o.val] ) - # weights. TF1 W is icfo. Need to convert to ifoc + # Set weights. The layout of the weight in TF1 is icfo (input, cell, forget, output gate). + # Need to convert to ifoc for coreml tf_w = op.weight.val # [input_dim+hidden_dim, 4*hidden_dim] in icfo layout tf_w_i, tf_w_c, tf_w_f, tf_w_o = np.split(tf_w, 4, axis=1) w = np.concatenate([tf_w_i, tf_w_f, tf_w_o, tf_w_c], axis=1) diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_composite_ops.py b/coremltools/converters/mil/frontend/tensorflow/test/test_composite_ops.py index 2c15e7079..bf400b36d 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_composite_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_composite_ops.py @@ -3,8 +3,12 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +import itertools +import numpy as np +import pytest + +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import random_gen from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import ( make_tf_graph, TensorFlowBaseTest @@ -20,9 +24,10 @@ from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import ( _TF_OPS_REGISTRY, ) -from coremltools.converters.mil.mil.ops.defs._op_reqs import * from coremltools.converters.mil.mil import Builder as mb +tf = pytest.importorskip("tensorflow") + class TestCompositeOp(TensorFlowBaseTest): @pytest.fixture(scope="class") diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py b/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py index 52d7b6e6a..ec0ac6778 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py @@ -2,14 +2,20 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import numpy as np +import pytest +from coremltools._deps import MSG_TF1_NOT_FOUND from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +from coremltools.converters.mil.testing_reqs import backends from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import ( - make_tf_graph, tf_graph_to_mlmodel ) +if testing_reqs._HAS_TF_1: + from coremltools.converters.mil.testing_reqs import tf + # Custom Op imports from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import register_tf_op @@ -20,8 +26,20 @@ from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import ( _TF_OPS_REGISTRY, ) -from coremltools.converters.mil.mil.ops.defs._op_reqs import * -from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.testing_utils import random_gen +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil import ( + Builder as mb, + Operation, + types, +) +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + DefaultInputs, + InputSpec, + IntInputType, + TensorInputType, +) from coremltools.converters.mil.mil.types.symbolic import is_symbolic @@ -251,4 +269,3 @@ def test_tf(self, use_cpu_only, backend, rank, k): assert ( True == layers[-1].custom.parameters["sorted"].boolValue ), "Incorrect parameter value for Sorted" - diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_graphs.py b/coremltools/converters/mil/frontend/tensorflow/test/test_graphs.py index 659f8933b..3edbc0fbd 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_graphs.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_graphs.py @@ -3,11 +3,13 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import numpy as np +import pytest + from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import ( make_tf_graph, - layer_counts, TensorFlowBaseTest ) backends = testing_reqs.backends diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py b/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py index 8859faf98..ccf8a3c46 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py @@ -3,8 +3,16 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import numpy as np +import math +import os +import pytest +import shutil +import tempfile + from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +from coremltools.converters.mil.testing_utils import random_gen from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import ( make_tf_graph, layer_counts, @@ -15,10 +23,6 @@ ) from coremltools.models.utils import _macos_version -import math -import tempfile -import shutil - backends = testing_reqs.backends tf = pytest.importorskip("tensorflow") @@ -670,6 +674,8 @@ class TestCond(TensorFlowBaseTest): "use_cpu_only, backend", itertools.product([True, False], backends,) ) def test_cond_naive(self, use_cpu_only, backend): + if (backend[0] == "mlprogram" and backend[1] == "fp16"): + pytest.xfail("rdar://83626929 (Reenable CondTests)") @make_tf_graph([(1,), (1,)]) def build_model(x, y): return tf.cond(tf.constant(True), lambda: x + y, lambda: x * y) @@ -874,8 +880,8 @@ def build_model(x,y): input_values = [np.array([[1],[2]], dtype=np.float32),np.array([[1],[2]], dtype=np.float32)] input_dict = dict(zip(inputs, input_values)) TensorFlowBaseTest.run_compare_tf(model, input_dict, outputs, - use_cpu_only=use_cpu_only, - backend=backend) + use_cpu_only=use_cpu_only, + backend=backend) @pytest.mark.parametrize( "use_cpu_only, backend", itertools.product([True, False], backends) @@ -1729,7 +1735,6 @@ def test_conv3d_transpose( pass w_shape = (kD, kH, kW, C_out, C_in) - x_input = np.random.randn(*input_shape) @make_tf_graph([tf_input_shape]) def build_model(x): @@ -5655,7 +5660,6 @@ def test_tf_no_variable( # _lstm_block_cell allows fine-grained control of W, peephole etc from tensorflow.contrib.rnn.python.ops.lstm_ops import _lstm_block_cell - actual_len, padded_len = 3, 4 input_dim, hidden_dim = 2, 3 x_shape = (batch, input_dim) init_h = np.random.rand(batch, hidden_dim).astype(np.float32) @@ -5701,7 +5705,6 @@ def test_tf_no_variable( itertools.product([True, False], backends, [1, 2],), ) def test_tf_lstm_block_cell(self, use_cpu_only, backend, batch): - actual_len, padded_len = 3, 4 input_dim, hidden_dim = 2, 3 # [timelen, batch_size, num_inputs] x_shape = (batch, input_dim) diff --git a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/constant_propagation.py b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/constant_propagation.py index 3d34fd682..b506f388a 100644 --- a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/constant_propagation.py +++ b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/constant_propagation.py @@ -5,16 +5,17 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +from distutils.version import StrictVersion as _StrictVersion +import gc import logging import tensorflow as tf -import gc + +from ..basic_graph_ops import const_determined_nodes from .delete_constant import delete_unnecessary_constant_nodes -from ..basic_graph_ops import const_determined_nodes, delete_node, disconnect_edge +from coremltools._deps import _get_version +from coremltools.converters._profile_utils import _profile from coremltools.converters.mil.mil import types from coremltools.converters.mil.mil.types.type_mapping import numpy_val_to_builtin_val -from coremltools.converters._profile_utils import _profile -from distutils.version import StrictVersion as _StrictVersion -from coremltools._deps import __get_version as _get_version def _get_const_nodes(fn): diff --git a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/delete_constant.py b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/delete_constant.py index 3cbdf66a0..a4d230ffd 100644 --- a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/delete_constant.py +++ b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/delete_constant.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be diff --git a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/functionalize_loops.py b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/functionalize_loops.py index cfd452ef7..6428ffa8c 100644 --- a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/functionalize_loops.py +++ b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/functionalize_loops.py @@ -1,12 +1,18 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from ..parsed_tf_node import ParsedTFNode -from ..basic_graph_ops import * # pylint: disable=unused-wildcard-import,wildcard-import +from ..basic_graph_ops import ( + connect_dests, + connect_edge, + connect_sources, + delete_node, + disconnect_edge, + replace_dest, + replace_source +) from ..tfssa import SSAFunction from .visitors import ( FindAllReachableNodes, @@ -311,7 +317,6 @@ def functionalize_loops(self, tfssa, function_to_functionalize): # connect constant enters to come from function # connect constant enters to exit for idx, enter in enumerate(self.constant_enters): - body_connected = False for output in list(g[enter].outputs): if output not in self.cond and output not in self.body: cond_intersection = ( diff --git a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/visitors.py b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/visitors.py index 7790a855a..a02f8517f 100644 --- a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/visitors.py +++ b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/visitors.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be diff --git a/coremltools/converters/mil/frontend/tensorflow2/__init__.py b/coremltools/converters/mil/frontend/tensorflow2/__init__.py index 4cca3ab2a..6cf9ed730 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/__init__.py +++ b/coremltools/converters/mil/frontend/tensorflow2/__init__.py @@ -6,7 +6,9 @@ from ....._deps import _HAS_TF_2 if _HAS_TF_2: - from .ops import * # register all TF2 ops + # importing these causes all its imports to be registered + from . import ops + from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import ( register_tf_op, ) diff --git a/coremltools/converters/mil/frontend/tensorflow2/load.py b/coremltools/converters/mil/frontend/tensorflow2/load.py index b98fc7545..dd32a2eba 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/load.py +++ b/coremltools/converters/mil/frontend/tensorflow2/load.py @@ -1,25 +1,40 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - +from distutils.version import StrictVersion as _StrictVersion import logging as _logging import os.path as _os_path +from tqdm import tqdm as _tqdm import tensorflow as _tf +from tensorflow.lite.python.util import get_grappler_config as _get_grappler_config +from tensorflow.lite.python.util import ( + run_graph_optimizations as _run_graph_optimizations, +) +from tensorflow.python.framework import dtypes as _dtypes +from tensorflow.python.framework.convert_to_constants import ( + convert_variables_to_constants_v2 as _convert_variables_to_constants_v2, +) +from tensorflow.python.framework.function_def_to_graph import ( + function_def_to_graph as _function_def_to_graph, +) +from tensorflow.python.keras.saving import saving_utils as _saving_utils +from tensorflow.python.eager import context + +from .converter import TF2Converter +from coremltools._deps import _get_version from coremltools.converters.mil.frontend.tensorflow.basic_graph_ops import fill_outputs from coremltools.converters.mil.frontend.tensorflow.load import TFLoader from coremltools.converters.mil.frontend.tensorflow.parsed_tf_node import ParsedTFNode from coremltools.converters.mil.frontend.tensorflow.tf_graph_pass import ( constant_propagation, - remove_variable_nodes, - tensor_array_resource_removal, - insert_get_tuple, delete_disconnected_nodes, fuse_dilation_conv, + insert_get_tuple, + remove_variable_nodes, + tensor_array_resource_removal, ) from coremltools.converters.mil.frontend.tensorflow.tfssa import ( NetworkEnsemble, @@ -29,24 +44,7 @@ flatten_sub_graph_namespaces, rewrite_control_flow_functions, ) -from coremltools._deps import __get_version as _get_version -from tensorflow.lite.python.util import get_grappler_config as _get_grappler_config -from tensorflow.lite.python.util import ( - run_graph_optimizations as _run_graph_optimizations, -) -from tensorflow.python.framework import dtypes as _dtypes -from tensorflow.python.framework.convert_to_constants import ( - convert_variables_to_constants_v2 as _convert_variables_to_constants_v2, -) -from tensorflow.python.framework.function_def_to_graph import ( - function_def_to_graph as _function_def_to_graph, -) -from tensorflow.python.keras.saving import saving_utils as _saving_utils -from tqdm import tqdm as _tqdm -from tensorflow.python.eager import context -from .converter import TF2Converter -from distutils.version import StrictVersion as _StrictVersion class TF2Loader(TFLoader): def __init__(self, model, debug=False, **kwargs): @@ -57,7 +55,7 @@ def __init__(self, model, debug=False, **kwargs): ---------- model: Model created with TensorFlow 2.x One of the following model format: - - TensorFlow tf.keras.Model object or HDF5 (.h5) file path + - TensorFlow tf.keras.Model object or HDF5 (.h5 or .hdf5) file path - TensorFlow SavedModel directory path - TensorFlow list of concrete functions(s) debug: bool, optional. Defaults to False. @@ -91,7 +89,8 @@ def _graph_def_from_model(self, outputs=None): raise ValueError( 'Input model "{}" does not exist'.format(self.model) ) - elif _os_path.isfile(self.model) and self.model.endswith(".h5"): + elif _os_path.isfile(self.model) \ + and (self.model.endswith(".h5") or self.model.endswith(".hdf5")): cfs = self._concrete_fn_from_tf_keras_or_h5(self.model) elif _os_path.isdir(self.model): saved_model = _tf.saved_model.load(self.model) diff --git a/coremltools/converters/mil/frontend/tensorflow2/ops.py b/coremltools/converters/mil/frontend/tensorflow2/ops.py index 85dec35ef..15d548dcf 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/ops.py +++ b/coremltools/converters/mil/frontend/tensorflow2/ops.py @@ -5,18 +5,21 @@ import numpy as _np +from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.frontend.tensorflow.convert_utils import convert_graph from coremltools.converters.mil.frontend.tensorflow.ops import ( _transpose_NHWC_to_NCHW, _transpose_NCHW_to_NHWC, ) - +from coremltools.converters.mil.frontend.tensorflow.tf_op_registry import register_tf_op +from coremltools.converters.mil.mil.types import builtin_to_string from coremltools.converters.mil.mil.types.symbolic import any_symbolic # TF 2.x now imports and registers all TF 1.x op against the new registry # (separated from TF 1.x registry). Overwrite might needed in case the op # semantics are different between TF 1.x and TF 2.x.< -from coremltools.converters.mil.frontend.tensorflow.ops import * -from coremltools.converters.mil.frontend.tensorflow.dialect_ops import * +from coremltools.converters.mil.frontend.tensorflow import ops +from coremltools.converters.mil.frontend.tensorflow import dialect_ops @register_tf_op(override=True, tf_alias=["FusedBatchNorm"]) @@ -136,14 +139,14 @@ def TensorListFromTensor(context, node): value = context[node.inputs[0]] element_shape = context[node.inputs[1]] element_dtype = node.attr.get("element_dtype") - dtype_str = types.builtin_to_string(element_dtype) + dtype_str = builtin_to_string(element_dtype) length = mb.shape(x=value) length = mb.slice_by_index(x=length, begin=[0], end=[1], squeeze_mask=[True]) if element_shape is not None and all(_np.atleast_1d(element_shape.val) != -1): ls = mb.make_list(init_length=length, - elem_shape=tuple(element_shape.val.tolist()), dtype=dtype_str) + elem_shape=tuple(element_shape.val.tolist()), dtype=dtype_str) else: ls = mb.tf_make_list(init_length=length, dtype=dtype_str) @@ -175,29 +178,26 @@ def TensorListLength(context, node): context.add(node.name, length) -@register_tf_op -def TensorListResize(context, node): - # skip here as the list will be dynamically resized when - # necessary in downstream list_write or list_scatter ops - Identity(context, node) - - @register_tf_op def TensorListReserve(context, node): element_shape = context[node.inputs[0]] num_elements = context[node.inputs[1]] element_dtype = node.attr.get("element_dtype") - dtype = types.builtin_to_string(element_dtype) + dtype = builtin_to_string(element_dtype) if element_shape is not None and all(_np.atleast_1d(element_shape.val) != -1): ls = mb.make_list( init_length=num_elements, elem_shape=tuple(element_shape.val.tolist()), + dynamic_length=num_elements.val is None, dtype=dtype, name=node.name, ) else: - ls = mb.tf_make_list(init_length=num_elements, dtype=dtype, name=node.name) + ls = mb.tf_make_list(init_length=num_elements, + dtype=dtype, + dynamic_length=num_elements.val is None, + name=node.name) context.add(node.name, ls) diff --git a/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/__init__.py b/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/__init__.py index a659bb2d3..91ca84e43 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/__init__.py +++ b/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/__init__.py @@ -3,23 +3,4 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -# Import all passes in this dir -from os.path import dirname, basename, isfile, join -import glob - -excluded_files = [ - "__init__.py", - "tf_passes.py", -] -modules = glob.glob(join(dirname(__file__), "*.py")) -pass_modules = [ - basename(f)[:-3] - for f in modules - if isfile(f) - and basename(f)[:1] != "_" # Follow python convention to hide _* files. - and basename(f)[:4] != "test" - and basename(f) not in excluded_files -] -__all__ = pass_modules - -from . import * # import everything in __all__ +from . import remove_vacuous_cond diff --git a/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/remove_vacuous_cond.py b/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/remove_vacuous_cond.py index 9707f14ea..6a96e89e5 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/remove_vacuous_cond.py +++ b/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/remove_vacuous_cond.py @@ -8,14 +8,14 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass import logging - -def remove_vacuous_cond_block(block): +def _remove_vacuous_cond_block(block): num_changes = 0 for op in list(block.operations): for b in op.blocks: - num_changes += remove_vacuous_cond_block(b) + num_changes += _remove_vacuous_cond_block(b) if op.op_type != "cond": continue @@ -79,9 +79,8 @@ def remove_vacuous_cond_block(block): return num_changes - @register_pass(namespace="tensorflow2") -def remove_vacuous_cond(prog): +class remove_vacuous_cond(AbstractGraphPass): """ Remove cond op and it's sub-graphs that produces identity on both then and else branch. One example use case is the TensorListReverse op, in Core ML, @@ -114,7 +113,8 @@ def remove_vacuous_cond(prog): } -> (%cond_0) } """ - for f_name, f in prog.functions.items(): - num_changes = remove_vacuous_cond_block(f) - msg = "remove_vacuous_cond: changed {} ops in function '{}'" - logging.info(msg.format(num_changes, f_name)) + def apply(self, prog): + for f_name, f in prog.functions.items(): + num_changes = _remove_vacuous_cond_block(f) + msg = "remove_vacuous_cond: changed {} ops in function '{}'" + logging.info(msg.format(num_changes, f_name)) diff --git a/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/tf_passes.py b/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/tf_passes.py index 57f30828a..b1c241aa4 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/tf_passes.py +++ b/coremltools/converters/mil/frontend/tensorflow2/ssa_passes/tf_passes.py @@ -3,9 +3,10 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY import logging +from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY + def tensorflow_passes(prog): passes = [ diff --git a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_composite_ops.py b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_composite_ops.py index 00f7c420d..6119f6005 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_composite_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_composite_ops.py @@ -3,6 +3,8 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import pytest + from coremltools.converters.mil import testing_reqs from coremltools.converters.mil.frontend.tensorflow.test import ( testing_utils as tf_testing_utils, @@ -10,7 +12,6 @@ from coremltools.converters.mil.frontend.tensorflow2.test.testing_utils import ( make_tf2_graph as make_tf_graph ) -from coremltools.converters.mil.testing_reqs import * tf = pytest.importorskip("tensorflow", minversion="2.1.0") diff --git a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops.py b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops.py index 9f976126f..638d9d0a1 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops.py @@ -5,12 +5,14 @@ import itertools import numpy as np +import pytest from coremltools._deps import version_lt, version_ge from coremltools.converters.mil import testing_reqs from coremltools.converters.mil.frontend.tensorflow.test import ( testing_utils as tf_testing_utils, ) +from coremltools.converters.mil.testing_utils import random_gen from coremltools.converters.mil.frontend.tensorflow2.test.testing_utils import ( TensorFlow2BaseTest @@ -22,7 +24,6 @@ from coremltools.converters.mil.frontend.tensorflow2.test.testing_utils import ( make_tf2_graph as make_tf_graph ) -from coremltools.converters.mil.testing_reqs import * tf = pytest.importorskip("tensorflow", minversion="2.1.0") diff --git a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops_tf_keras.py b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops_tf_keras.py index 2421d2bbd..1b76c0984 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops_tf_keras.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops_tf_keras.py @@ -3,20 +3,25 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -import random +from distutils.version import StrictVersion as _StrictVersion import itertools -import pytest import numpy as np -from coremltools._deps import __get_version as _get_version +import pytest +import random + + +import coremltools as ct +from coremltools._deps import _get_version from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * from coremltools.converters.mil.frontend.tensorflow2.test.testing_utils import ( TensorFlow2BaseTest ) from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import ( TensorFlowBaseTest ) -from distutils.version import StrictVersion as _StrictVersion +from coremltools.converters.mil.testing_utils import random_gen +from ..._utils import is_symbolic_dim_in_prog + TensorFlowBaseTest.run_compare_tf_keras = \ TensorFlow2BaseTest.run_compare_tf_keras backends = testing_reqs.backends @@ -1258,6 +1263,61 @@ def test_lstm_dynamic_batch(self, use_cpu_only, backend): backend=backend, ) + @pytest.mark.parametrize( + "use_cpu_only, backend", itertools.product([True], backends) + ) + def test_lstm_conversion_static_shapes(self, use_cpu_only, backend): + ''' + Test that intermediate tensor shapes are populated correctly by the converter. + That is, there are no symbolic dimensions in the shapes, when conversion is + performed with a fixed input shape, irrespective of the shape used in the source model definition. + ''' + def _get_keras_simple_lstm_model(input_shape): + input = tf.keras.Input(batch_input_shape=input_shape) + output = tf.keras.layers.LSTM(5)(input) + keras_model = tf.keras.Model(inputs=input, outputs=output) + return keras_model + + def _test_for_symbolic_shapes(keras_input_shape, input_shape_for_conversion, are_symbols_expected): + keras_model = _get_keras_simple_lstm_model(keras_input_shape) + res = TensorFlowBaseTest.run_compare_tf_keras( + keras_model, + [ + random_gen((1, 32, 10), -1, 1), + ], + inputs_for_conversion=[ct.TensorType(shape=input_shape_for_conversion)], + use_cpu_only=use_cpu_only, + backend=backend, + ) + coreml_model = res[1] + mil_prog = coreml_model._get_mil_internal() + assert is_symbolic_dim_in_prog(mil_prog) == are_symbols_expected + + + _test_for_symbolic_shapes(keras_input_shape=(1, 32, 10), + input_shape_for_conversion=(1, 32, 10), + are_symbols_expected=False) + + _test_for_symbolic_shapes(keras_input_shape=(None, 32, 10), + input_shape_for_conversion=(1, 32, 10), + are_symbols_expected=False) + + _test_for_symbolic_shapes(keras_input_shape=(None, None, 10), + input_shape_for_conversion=(1, 32, 10), + are_symbols_expected=False) + + _test_for_symbolic_shapes(keras_input_shape=(None, 32, 10), + input_shape_for_conversion=(ct.RangeDim(1, 10), 32, 10), + are_symbols_expected=True) + + if backend[0] != "mlprogram": + # FIX ME: model load fails if backend is "mlprogram". rdar://84862138 + _test_for_symbolic_shapes(keras_input_shape=(None, None, 10), + input_shape_for_conversion=(ct.RangeDim(1, 10), ct.RangeDim(16, 64), 10), + are_symbols_expected=True) + + + class TestRepeatVector(TensorFlowBaseTest): @pytest.mark.parametrize( "use_cpu_only, backend, n", diff --git a/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py b/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py index 51559128d..41506fe1d 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py @@ -182,6 +182,7 @@ def run_compare_tf2( def run_compare_tf_keras( model, input_values, + inputs_for_conversion=None, use_cpu_only=False, frontend_only=False, frontend="tensorflow", @@ -196,6 +197,8 @@ def run_compare_tf_keras( TensorFlow 2.x model annotated with @tf.function. input_values: list of np.array List of input values in the same order as the input signature. + inputs_for_conversion: list of coremltools.TensorType() or coremltools.ImageType() objects + Defaults to None. It is passed as is to the "inputs" argument of the converter. use_cpu_only: bool If true, use CPU only for prediction, otherwise, use GPU also. frontend_only: bool @@ -209,7 +212,7 @@ def run_compare_tf_keras( rtol: float The relative tolerance parameter. """ - mlmodel = ct_convert(model, source=frontend, convert_to=backend) + mlmodel = ct_convert(model, inputs=inputs_for_conversion, source=frontend, convert_to=backend) # assumes conversion preserve the i/o names proto = mlmodel.get_spec() @@ -278,10 +281,12 @@ def run_compare_tf2(model, return tuple(alist) @staticmethod - def run_compare_tf_keras(model, input_values, use_cpu_only=False, + def run_compare_tf_keras(model, input_values, inputs_for_conversion=None, use_cpu_only=False, frontend_only=False, frontend="tensorflow", backend=("neuralnetwork", "fp32"), atol=1e-04, rtol=1e-05): - res = run_compare_tf_keras(model, input_values, use_cpu_only=use_cpu_only, + res = run_compare_tf_keras(model, input_values, + inputs_for_conversion=inputs_for_conversion, + use_cpu_only=use_cpu_only, frontend_only=frontend_only, frontend=frontend, backend=backend, atol=atol, rtol=rtol) diff --git a/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/__init__.py b/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/__init__.py index 3a0cf529d..c230ec68d 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/__init__.py +++ b/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/__init__.py @@ -3,4 +3,7 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from .rewrite_control_flow_functions import * +from .rewrite_control_flow_functions import ( + flatten_sub_graph_namespaces, + rewrite_control_flow_functions +) diff --git a/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/rewrite_control_flow_functions.py b/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/rewrite_control_flow_functions.py index 463aec01d..f3a0ea8d8 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/rewrite_control_flow_functions.py +++ b/coremltools/converters/mil/frontend/tensorflow2/tf_graph_pass/rewrite_control_flow_functions.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be diff --git a/coremltools/converters/mil/frontend/torch/__init__.py b/coremltools/converters/mil/frontend/torch/__init__.py index 984f69ab6..2b096c5c3 100644 --- a/coremltools/converters/mil/frontend/torch/__init__.py +++ b/coremltools/converters/mil/frontend/torch/__init__.py @@ -10,4 +10,8 @@ if _HAS_TORCH: from .load import load from .torch_op_registry import register_torch_op - from .dialect_ops import * + from .dialect_ops import ( + torch_tensor_assign, + torch_upsample_bilinear, + torch_upsample_nearest_neighbor + ) diff --git a/coremltools/converters/mil/frontend/torch/converter.py b/coremltools/converters/mil/frontend/torch/converter.py index 33f2fd7dd..6a4f54904 100644 --- a/coremltools/converters/mil/frontend/torch/converter.py +++ b/coremltools/converters/mil/frontend/torch/converter.py @@ -3,26 +3,29 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - +from collections import OrderedDict import logging as _logging import torch as _torch -from coremltools._deps import version_lt -from coremltools.converters.mil.input_types import InputType, ImageType -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil import Builder as mb +from coremltools._deps import version_lt +from coremltools.converters.mil.input_types import ImageType from coremltools.converters.mil.mil import ( - Placeholder, + Builder as mb, Function, - Program, - get_new_symbol, + types, + Program ) -from coremltools.converters.mil.mil import Var -from .internal_graph import * -from .ops import * +from .internal_graph import InternalTorchIRGraph +from .ops import convert_nodes from .torch_op_registry import _TORCH_OPS_REGISTRY -from .torchir_passes import * +from .torchir_passes import ( + flatten_graph_input_values, + flatten_graph_output_values, + generate_tensor_assignment_ops, + transform_inplace_ops, + remove_getattr_nodes +) from .ssa_passes.torch_passes import torch_passes torch_to_mil_types = { @@ -145,6 +148,7 @@ def __init__( self.output_names = outputs self.context = TranscriptionContext() raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript) + self.params_dict = params_dict self.graph = InternalTorchIRGraph( raw_graph, params_dict, self.inputs, cut_at_symbols ) @@ -152,6 +156,7 @@ def __init__( transform_inplace_ops, flatten_graph_input_values, flatten_graph_output_values, + remove_getattr_nodes, generate_tensor_assignment_ops, ] for p in passes: @@ -193,6 +198,11 @@ def check_ops(self): the set for which no conversion function is registered.""" return TorchConverter._check_ops(self.graph) + def convert_const(self): + for name, val in self.graph.params.items(): + const = mb.const(val=val, name=name) + self.context.add(const) + def convert(self): _logging.info("Converting graph.") @@ -222,9 +232,8 @@ def convert(self): self.graph.inputs.keys(), ssa_func_inputs.keys() ): self.context.add(ssa_func.inputs[users_name], torch_name=internal_name) - for name, val in self.graph.params.items(): - const = mb.const(val=val, name=name) - self.context.add(const) + + self.convert_const() # Add the rest of the operations convert_nodes(self.context, self.graph) @@ -250,19 +259,144 @@ def convert(self): self.torch_passes(prog) return prog + def _jit_pass_lower_graph(graph, torchscript): + """ + This graph pass does a similar thing as _torch._C._jit_pass_lower_graph does. + It does two things: + 1. Rename getattr nodes which produce a torch tensor to match the keys in torch model's state_dict + 2. Construct the params_dict, with the keys similar to state_dict + + To be more specific, this graph pass traces down series of GetAttr ops, and rename the final node to match the torch model state_dict. + It also replaces the node inputs by the first created tensor node with the same name. + + Example: + Input graph: + graph(%self.1 : __torch__.torch.nn.modules.Sequential, %input.1 : Tensor): + %2 : prim::GetAttr[name="linear"](%self.1) + %3 : prim::GetAttr[name="weight"](%2) + %4 : prim::GetAttr[name="bias"](%2) + %5 : prim::GetAttr[name="bias"](%2) # duplicated node + %6 : conv(%input.1, %3, %4) + %7 : add(%input.1, %5) + return (%6, %7) + + Output graph: + graph(%self.1 : __torch__.torch.nn.modules.Sequential, %input.1 : Tensor): + %2 : prim::GetAttr[name="linear"](%self.1) + %linear.weight : prim::GetAttr[name="weight"](%2) + %linear.bias : prim::GetAttr[name="bias"](%2) + %5 : prim::GetAttr[name="bias"](%2) # duplicated node, it is not used now + %6 : conv(%input.1, %linear.weight, %linear.bias) + %7 : add(%input.1, %linear.bias) # the second input is replaced + return (%6, %7) + + And a dictionary {"linear.weight": ..., "linear.bias": ...} is returned, to record the parameters values. + Note that, those GetAttr nodes are still in the torch ir graph, but they would be removed in a latter + graph pass in the coreml torch internal graph + + """ + + """ + Each getattr node corresponds to a torch object in the torch IR, + it could be either: + 1. torch.nn.modules: submodule in a torch model. For instance, a linear layer in a MLP network. + 2. torch.Tensor: torch model parameters. For instance, weight for a conv layer. + 3. torch._C.ScriptObject: quantized torch model parameters. + For example, in the graph above, %2 is pointing to the __torch__.torch.nn.modules.Sequential.linear torch submodule. + node_to_module_map tracks these mapping. + + node_to_prefic_map track the name for each module, + for example, %2 has the prefix name linear and %3 is linear.weight. + These names are also keys in the state_dict + """ + node_to_module_map = {} + node_to_prefix_map = {} + first_node_with_prefix = {} + replace_input = {} + + base_module_node = list(graph.inputs())[0] + node_to_module_map[base_module_node] = torchscript + node_to_prefix_map[base_module_node] = "" + + """ + params_dict will be contructed in this graph pass. It contains all const tensors needed for the graph computation. + And the value is validated against the state_dict if the key is presented in both dictionaries. + In some rare cases, state_dict lacks parameters / buffers, so we still need to go through the while graph ourselves. + """ + params_dict = {} + state_dict = torchscript.state_dict(keep_vars=True) + + def _check_is_tensor(node, module): + if not isinstance(module, _torch.Tensor): + return False + assert str(node.output().type()) == "Tensor" + return True + + def _check_is_quantized_tensor(node, module): + if not isinstance(module, _torch._C.ScriptObject): + return False + # There are three quantized parameters currently supported in Torch: + # ref: https://github.com/pytorch/pytorch/blob/master/torch/csrc/jit/passes/lower_graph.cpp + supported_quantizes_types = ["LinearPackedParamsBase", "Conv2dPackedParamsBase", "Conv3dPackedParamsBase"] + assert node.output().type().name() in supported_quantizes_types + return True + + def _get_tensor(module): + return module + + def _get_quantized_tensor(module): + return tuple(list(module.__getstate__())[:-1]) + + def _lower_graph_block(graph): + for node in list(graph.nodes()): + + for block in node.blocks(): + _lower_graph_block(block) + + for idx, _input in enumerate(list(node.inputs())): + if _input in replace_input: + node.replaceInput(idx, replace_input[_input]) + + kind = node.kind().split("::")[1].lower() + if kind != "getattr": + continue + + _input = node.input() + _output = node.output() + attr_name = getattr(node, node.kindOf("name"))("name") + + module = getattr(node_to_module_map[_input], attr_name) + node_to_module_map[_output] = module + + input_prefix = node_to_prefix_map[_input] + prefix = input_prefix + '.' + attr_name if input_prefix != "" else attr_name + node_to_prefix_map[_output] = prefix + + is_tensor = _check_is_tensor(node, module) + is_quantized_tensor = _check_is_quantized_tensor(node, module) + + if is_tensor or is_quantized_tensor: + if is_tensor and prefix in state_dict: + assert _torch.equal(module, state_dict[prefix]), "tensor value not consistent between torch ir and state_dict" + if prefix in params_dict: + assert _torch.equal(module, params_dict[prefix]) + replace_input[_output] = first_node_with_prefix[prefix] + else: + params_dict[prefix] = _get_tensor(module) if is_tensor else _get_quantized_tensor(module) + first_node_with_prefix[prefix] = _output + _output.setDebugName(prefix) + + _lower_graph_block(graph) + + return graph, params_dict + + @staticmethod def _expand_and_optimize_ir(torchscript): """Given a torch.jit.ScriptModule, convert it to a optimized torch._C.Graph and dict of model parameter's names to tensors. """ - - # Recursively replaces all attribute accesses with the sub-graphs of - # those modules. The resulting graph will be self-contained and will - # not reference into other modules. Params will contain the "trainable" - # inputs to the graph. - graph, params = _torch._C._jit_pass_lower_graph( - torchscript.forward.graph, torchscript._c - ) + graph = torchscript.forward.graph # From PyTorch code: Inline function and method calls. _torch._C._jit_pass_inline(graph) @@ -315,8 +449,7 @@ def _expand_and_optimize_ir(torchscript): # NOTE: Don't need another DCE, it's included in constant propagation. _torch._C._jit_pass_lint(graph) - input_and_param_names = [val.debugName() for val in graph.inputs()] - param_names = input_and_param_names[len(input_and_param_names) - len(params):] - params_dict = dict(zip(param_names, params)) + # Get the params_dict and rename the getattr nodes in the graph + graph, params_dict = TorchConverter._jit_pass_lower_graph(graph, torchscript) return graph, params_dict diff --git a/coremltools/converters/mil/frontend/torch/dialect_ops.py b/coremltools/converters/mil/frontend/torch/dialect_ops.py index a9a6c5896..e097f6956 100644 --- a/coremltools/converters/mil/frontend/torch/dialect_ops.py +++ b/coremltools/converters/mil/frontend/torch/dialect_ops.py @@ -3,13 +3,20 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil import Operation -from coremltools.converters.mil.mil.input_type import * -from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry -from coremltools.converters.mil.mil.types.symbolic import is_symbolic, is_compatible_symbolic_vector -from coremltools.converters.mil.mil import get_new_symbol +from coremltools.converters.mil.mil import get_new_symbol, Operation, types +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + BoolTensorInputType, + DefaultInputs, + InputSpec, + IntOrFloatInputType, + IntTensorInputType, + TensorInputType +) from coremltools.converters.mil.mil.ops.defs.tensor_transformation import _solve_slice_by_index_shape +from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry +from coremltools.converters.mil.mil.types.symbolic import is_compatible_symbolic_vector + register_op = SSAOpRegistry.register_op diff --git a/coremltools/converters/mil/frontend/torch/internal_graph.py b/coremltools/converters/mil/frontend/torch/internal_graph.py index 097519f89..d6155c8a9 100644 --- a/coremltools/converters/mil/frontend/torch/internal_graph.py +++ b/coremltools/converters/mil/frontend/torch/internal_graph.py @@ -250,7 +250,9 @@ def __init__( self.params[name] = value # Add inputs - for index, _input in enumerate(islice(raw_graph.inputs(), len(input_values))): + # The first element of the raw_graph.inputs() is the 'self' of the module, which is not used. + graph_inputs = list(raw_graph.inputs())[1:] + for index, _input in enumerate(islice(graph_inputs, len(input_values))): name = _input.debugName() value = input_values[index] self.inputs[name] = value diff --git a/coremltools/converters/mil/frontend/torch/load.py b/coremltools/converters/mil/frontend/torch/load.py index a7ad76302..d85ccf8b7 100644 --- a/coremltools/converters/mil/frontend/torch/load.py +++ b/coremltools/converters/mil/frontend/torch/load.py @@ -6,8 +6,8 @@ import logging as _logging import os.path as _os_path - import torch as _torch + from .converter import TorchConverter, torch_to_mil_types from coremltools.converters.mil.input_types import InputType, TensorType from coremltools.converters.mil.mil import Program, types @@ -43,45 +43,11 @@ def load(model_spec, debug=False, **kwargs): only. """ torchscript = _torchscript_from_model(model_spec) - - def _convert_to_inputtype(inputs): - input_type = [] - for _input in inputs: - if isinstance(_input, (list, tuple)): - input_type.append(_convert_to_inputtype(_input)) - elif isinstance(_input, InputType): - input_type.append(_input) - elif isinstance(_input, _torch.Tensor): - input_type.append( - TensorType( - shape=_input.shape, dtype=torch_to_mil_types[_input.dtype] - ) - ) - else: - raise ValueError( - "Unknown type {} for conversion to InputType.".format(type(_input)) - ) - return input_type - - inputs = _convert_to_inputtype(kwargs["inputs"]) + inputs = _convert_to_torch_inputtype(kwargs["inputs"]) outputs = kwargs.get("outputs", None) cut_at_symbols = kwargs.get("cut_at_symbols", None) converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols) - - try: - prog = converter.convert() - except RuntimeError as e: - if debug and "convert function" in str(e): - implemented, missing = converter.check_ops() - print("the following model ops are IMPLEMENTED:") - print("\n".join([" " + str(x) for x in sorted(implemented)])) - print("the following model ops are MISSING:") - print("\n".join([" " + str(x) for x in sorted(missing)])) - raise e - except Exception as e: - raise e - - return prog + return _perform_torch_convert(converter, debug) def _torchscript_from_model(model_spec): @@ -96,3 +62,37 @@ def _torchscript_from_model(model_spec): type(model_spec) ) ) + +def _convert_to_torch_inputtype(inputs): + input_type = [] + for _input in inputs: + if isinstance(_input, (list, tuple)): + input_type.append(_convert_to_torch_inputtype(_input)) + elif isinstance(_input, InputType): + input_type.append(_input) + elif isinstance(_input, _torch.Tensor): + input_type.append( + TensorType( + shape=_input.shape, dtype=torch_to_mil_types[_input.dtype] + ) + ) + else: + raise ValueError( + "Unknown type {} for conversion to InputType.".format(type(_input)) + ) + return input_type + +def _perform_torch_convert(converter, debug): + try: + prog = converter.convert() + except RuntimeError as e: + if debug and "convert function" in str(e): + implemented, missing = converter.check_ops() + print("the following model ops are IMPLEMENTED:") + print("\n".join([" " + str(x) for x in sorted(implemented)])) + print("the following model ops are MISSING:") + print("\n".join([" " + str(x) for x in sorted(missing)])) + raise e + + return prog + diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index c653e2a3c..787d1ea55 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -3,24 +3,24 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -import logging as _logging - import builtins +from collections.abc import Iterable +import logging as _logging import math as _math import numbers import numpy as _np +import torch from tqdm import tqdm as _tqdm -from collections.abc import Iterable -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil import ( + Builder as mb, + types, + Symbol +) from coremltools.converters.mil.mil.var import Var, ListVar -from coremltools.converters.mil.mil import Placeholder, Symbol -from .internal_graph import * from .torch_op_registry import _TORCH_OPS_REGISTRY, register_torch_op -from coremltools.converters.mil.mil.types.symbolic import any_symbolic, is_symbolic +from coremltools.converters.mil.mil.types.symbolic import any_symbolic, is_symbolic, is_compatible_symbolic_vector from coremltools.converters.mil.mil.types import is_bool -from coremltools.converters.mil.mil.ops.defs._utils import broadcast_shapes from .._utils import build_einsum_mil # The pytorch args for many of the below ops were sourced from @@ -52,22 +52,28 @@ def _value_at(x, idx): def convert_nodes(context, graph): - """Iterate over the nodes of a graph or block and convert to MIL. + """ + Iterate over the nodes of a graph or block and convert to MIL. - Arguments: - context: A TranscriptionContext object to pull node inputs and - assign node outputs. - graph: An InternalTorchIRGraph or InternalTorchIRBlock object. + Arguments: + context: A TranscriptionContext object to pull node inputs and + assign node outputs. + graph: An InternalTorchIRGraph or InternalTorchIRBlock object. """ for node in _tqdm(graph.nodes, desc="Converting Frontend ==> MIL Ops", unit=" ops"): - _add_op = _TORCH_OPS_REGISTRY.get(node.kind, None) + op_lookup = node.kind + if op_lookup.endswith("_"): + # This is an "in place" op. + # Look up the standard op instead by removing underscore. + op_lookup = op_lookup[:-1] + add_op = _TORCH_OPS_REGISTRY.get(op_lookup, None) + _logging.info("Converting op {} : {}".format(node.name, node.kind)) - if _add_op is None: + if add_op is None: raise RuntimeError( "PyTorch convert function for op '{}' not implemented.".format(node.kind) ) - else: - _add_op(context, node) + add_op(context, node) # We've generated all the outputs the graph needs, terminate conversion. if _all_outputs_present(context, graph): @@ -205,12 +211,12 @@ def _construct_constant(val, name): # Pytorch uses inf if val is not None and isinstance(val, numbers.Number) \ - and _np.isinf(val): - if val < 0: # neg inf - # most negative number in fp32 - val = -3.4e+38 - else: # positive inf - val = 3.4e+38 + and _np.isinf(val): + if val < 0: # neg inf + # most negative number in fp32 + val = -3.4e+38 + else: # positive inf + val = 3.4e+38 if val is None: return None else: @@ -315,7 +321,7 @@ def grid_sampler(context, node): context.add(x) -@register_torch_op(torch_alias=["silu_"]) +@register_torch_op() def silu(context, node): inputs = _get_inputs(context, node, expected=1) x = mb.silu(x=inputs[0], name=node.name) @@ -333,6 +339,37 @@ def constant(context, node): const = _construct_constant(val, name) context.add(const, torch_name=name) +@register_torch_op +def frobenius_norm(context, node): + x, dim, keep_dims = _get_inputs(context, node, expected=3) + result = mb.reduce_l2_norm(x=x, axes=dim, keep_dims=keep_dims, name=node.name) + context.add(result) + +@register_torch_op +def norm(context, node): + VALUE_CLOSE_TO_INFINITY = 1e+38 + + x, num, dim, keep_dims = _get_inputs(context, node, expected=4) + assert x is not None and num is not None and dim is not None and keep_dims is not None + if num.val is None: + raise RuntimeError("Dynamic 'p' values for 'norm' layers are not supported.") + assert num.val != 0 + + if num.val == 1: + temp = mb.reduce_l1_norm(x=x, axes=dim, keep_dims=keep_dims, name=node.name) + elif num.val == 2: + temp = mb.reduce_l2_norm(x=x, axes=dim, keep_dims=keep_dims, name=node.name) + elif num.val > VALUE_CLOSE_TO_INFINITY: + temp = mb.reduce_max(x=x, axes=dim, keep_dims=keep_dims, name=node.name) + elif num.val < -VALUE_CLOSE_TO_INFINITY: + temp = mb.reduce_min(x=x, axes=dim, keep_dims=keep_dims, name=node.name) + else: + # sum(abs(x)**num)**(1./num) + temp = mb.abs(x=x) + temp = mb.pow(x=temp, y=num) + temp = mb.reduce_sum(x=temp, axes=dim, keep_dims=keep_dims) + temp = mb.pow(x=temp, y=1./num.val, name=node.name) + context.add(temp) def _array_construct(context, node, array_type): assert len(node.outputs) == 1 @@ -408,7 +445,7 @@ def gt(context, node): context.add(greater) -@register_torch_op(torch_alias=["t", "transpose_"]) +@register_torch_op(torch_alias=["t"]) def transpose(context, node): assert len(node.outputs) == 1 inputs = _get_inputs(context, node) @@ -688,7 +725,7 @@ def softsign(context, node): res = mb.softsign(x=inputs[0], name=node.name) context.add(res) -@register_torch_op(torch_alias=["relu_"]) +@register_torch_op() def relu(context, node): inputs = _get_inputs(context, node, expected=1) @@ -727,7 +764,7 @@ def elu(context, node): res = mb.elu(x=inputs[0], alpha = inputs[1], name=node.name) context.add(res) -@register_torch_op(torch_alias=["leaky_relu_"]) +@register_torch_op() def leaky_relu(context, node): inputs = _get_inputs(context, node, expected=2) @@ -928,8 +965,8 @@ def mul(context, node): context.add(res) -@register_torch_op(torch_alias=["pow"]) -def pow_(context, node): +@register_torch_op() +def pow(context, node): inputs = _get_inputs(context, node, expected=2) res = mb.pow(x=inputs[0], y=inputs[1], name=node.name) @@ -1080,7 +1117,6 @@ def adaptive_avg_pool2d(context, node): pad_type = "valid" # Need to explicity state L-R, T-B pad pad = [0, 0, 0, 0] - dilation = [1, 1] kernel_sizes = [ ind - s * (outd - 1) for ind, outd, s in zip(_input.shape[-2:], output_size, strides) @@ -1145,7 +1181,6 @@ def adaptive_max_pool2d(context, node): pad_type = "valid" # Need to explicity state L-R, T-B pad pad = [0, 0, 0, 0] - dilation = [1, 1] kernel_sizes = [ ind - s * (outd - 1) for ind, outd, s in zip(_input.shape[-2:], output_size, strides) @@ -1404,7 +1439,7 @@ def embedding(context, node): context.add(gather) -@register_torch_op(torch_alias=["hardtanh_"]) +@register_torch_op() def hardtanh(context, node): inputs = _get_inputs(context, node, expected=3) _input = inputs[0] @@ -1832,7 +1867,7 @@ def _add_simple_rnn(context, node, activation): (2) h0: (num_layers, B, H) ''' _input = inputs[0] - h0 = inputs[1] + h0 = inputs[1] weights_list = inputs[2] has_bias = inputs[3].val num_layers = inputs[4].val @@ -2250,8 +2285,8 @@ def _is_float_value(x, threshold=0.001): if (is_h_float or is_w_float) and not align_corners: msg = "recompute_scale_factor = False, align_corners = False with float output size is " + \ - "not supported for the upsample op {}".format(node.name) - raise NotImplementedError("") + "not supported for the upsample op {}".format(node.name) + raise NotImplementedError(msg) elif isinstance(output_size, list)and output_size[0].val is None and output_size[1].val is None: # the input shape is dynamic and recompute_scale_factor = True @@ -2621,7 +2656,7 @@ def _get_slice_params(context, data, inputs): num_of_slice_set = len(inputs) // 2 for i in range(num_of_slice_set): - if inputs[2*i + 1] == None: + if inputs[2*i + 1] is None: # This is pure index select idx = context[inputs[2*i]].val begin[i] = idx @@ -2666,7 +2701,7 @@ def _get_slice_params(context, data, inputs): context.add(updated_x) @register_torch_op -def index_put_(context, node): +def index_put(context, node): inputs = _get_inputs(context, node, expected=4) x = inputs[0] indices = inputs[1] @@ -2712,9 +2747,12 @@ def index(context, node): The true value indicates whether the element should be selected. The output b is a 1-D vector with shape (N), where N is the number of elements satisfying condition > 0.1 """ - if len(indices) == 1 and indices[0] is not None and indices[0].sym_type.get_primitive() == types.bool: + if (len(indices) == 1 and + indices[0] is not None and + indices[0].sym_type.get_primitive() == types.bool and + indices[0].shape == x.shape): + indices = indices[0] - assert indices.shape == x.shape, "indices shape must equal to input shape for index selection." x_reshape = mb.reshape(x=x, shape=[-1]) indices = mb.cast(x=indices, dtype="int32") indices_reshape = mb.reshape(x=indices, shape=[-1]) @@ -2760,6 +2798,16 @@ def index(context, node): b has shape (1,2,4) + EX # 4: + a = torch.rand(4,5) + index_1 = [True, True, False, False] + index_2 = [False, True, True, False, False] + b = a[index_1, index_2] + + indices is a list [[True, True, False, False], [False, True, True, False, False]] + + b has shape (2, ) + Note that, in pytorch, the indices can be broadcasable. And it is NOT supported right now. """ @@ -2778,17 +2826,23 @@ def index(context, node): context.add(x) return + # convert all indices to int type + for i, indice in enumerate(valid_indices): + if indice is not None and types.is_bool(indice.dtype): + indice = mb.non_zero(x=indice) + indice = mb.squeeze(x=indice, axes=[1]) + valid_indices[i] = indice + # For the single index axis case, we can use mb.gather directly if len(indices_axes) == 1: axis = indices_axes[0] - x = mb.gather(x=x, indices=indices[axis], axis=axis, name=node.name) + x = mb.gather(x=x, indices=valid_indices[0], axis=axis, name=node.name) context.add(x) return # For multiple index axes case, we now assume that all the index have equal shape - index_length = valid_indices[0].shape for index in valid_indices: - if index.shape != valid_indices[0].shape: + if not is_compatible_symbolic_vector(index.shape, valid_indices[0].shape): raise NotImplementedError("Broadcasable tensor index not supported.") # First stack the index together @@ -2824,6 +2878,8 @@ def ones(context, node): # device = inputs[3] unused # requires_grad = inputs[4] unused # out = inputs[5] unused + if isinstance(size, list): + size = mb.concat(values=size, axis=0) fill = mb.fill(shape=size, value=1.0, name=node.name) context.add(fill) @@ -2840,12 +2896,23 @@ def ones_like(context, node): fill = mb.fill(shape=size, value=1.0, name=node.name) context.add(fill) +@register_torch_op +def full(context, node): + inputs = _get_inputs(context, node) + size = inputs[0] + val = inputs[1].val + assert val is not None + if isinstance(size, list): + size = mb.concat(values=size, axis=0) + full = mb.fill(shape=size, value=val, name=node.name) + context.add(full) + def _avg_pool(context, node, inputs): x = inputs[0] kernel_sizes = inputs[1] strides = inputs[2] - if strides.op.op_type == "const" and (not list(strides.val)): + if strides.op.op_type == "const" and (not list(strides.val)): strides = mb.const(val=kernel_sizes.val, name=strides.name) pad_type = "custom" # Need to explicitly state L-R, T-B pad @@ -2967,7 +3034,7 @@ def sigmoid(context, node): res = mb.sigmoid(x=inputs[0], name=node.name) context.add(res) -@register_torch_op(torch_alias=["hardsigmoid_"]) +@register_torch_op() def hardsigmoid(context, node): inputs = _get_inputs(context, node, expected=1) @@ -3147,7 +3214,7 @@ def implicittensortonum(context, node): inputs = _get_inputs(context, node, expected=1) _input = inputs[0] - if _input.shape == (): #already a scalar + if _input.shape == (): # already a scalar context.add(_input, node.name) else: assert _input.shape == (1,) @@ -3176,12 +3243,20 @@ def constantchunk(context, node): context.add(val, name) -def _expand(context, name, tensor, shape): - if tensor.shape == () and len(shape) > 0: - tensor = mb.expand_dims(x=tensor, axes=list(range(len(shape)))) - reps = [ds if ds > 0 and ts == 1 else 1 for ts, ds in zip(tensor.shape, shape)] +def _broadcast(name, tensor, shape): + if len(shape) > tensor.rank: + new_dims = len(shape) - tensor.rank + tensor = mb.expand_dims(x=tensor, axes=list(range(new_dims))) + + reps = [] + for ts, ds in zip(tensor.shape, shape): + if not is_symbolic(ts) and not is_symbolic(ds) and ds > 0 and ts == 1: + reps.append(ds) + else: + reps.append(1) + res = mb.tile(x=tensor, reps=reps, name=name) - context.add(res) + return res @register_torch_op @@ -3192,7 +3267,8 @@ def expand(context, node): x = inputs[0] shape = inputs[1].val - _expand(context, node.name, x, shape) + res = _broadcast(node.name, x, shape) + context.add(res) @register_torch_op @@ -3202,8 +3278,8 @@ def expand_as(context, node): x = inputs[0] other = inputs[1] - _expand(context, node.name, x, other.shape) - + res = _broadcast(node.name, x, other.shape) + context.add(res) @register_torch_op def arange(context, node): @@ -3236,7 +3312,7 @@ def arange(context, node): context.add(res) -@register_torch_op(torch_alias=["masked_fill_"]) +@register_torch_op() def masked_fill(context, node): inputs = _get_inputs(context, node, expected=3) x = inputs[0] @@ -3957,3 +4033,44 @@ def replication_pad2d(context, node): pad = _np.pad(pad_flipped, (len(x.shape) * 2 - len(pad_flipped), 0)) context.add(mb.pad(x=x, pad=pad, mode='replicate'), node.name) +def _broadcast_tensors(tensors): + + def _solve_broadcast_shape(shapes): + rank = _np.max([len(shape) for shape in shapes]) + shapes = [[1]*(rank - len(shape)) + shape for shape in shapes] + result_shape = [] + for i in range(rank): + dims = [shapes[j][i] for j in range(len(tensors))] + if any_symbolic(dims): + raise NotImplementedError("Only static shaped inputs are supported for torch.broadcast_tensors conversion.") + result_shape.append(_np.max(dims)) + return result_shape + + if len(tensors) == 1: + return tensors + + # solve the broadcast shape + input_shapes = [list(x.shape) for x in tensors] + broadcast_shape = _solve_broadcast_shape(input_shapes) + + # do the broadcasting + results = [] + for tensor in tensors: + name = tensor.name + "_after_broadcast" + results.append(_broadcast(name, tensor, broadcast_shape)) + return results + +@register_torch_op +def broadcast_tensors(context, node): + inputs = _get_inputs(context, node) + context.add(_broadcast_tensors(inputs[0]), node.name) + +@register_torch_op +def scatter_add(context, node): + inputs = _get_inputs(context, node) + data = inputs[0] + axis = inputs[1].val + indices = inputs[2] + updates = inputs[3] + result = mb.scatter_along_axis(data=data, indices=indices, updates=updates, axis=axis, mode="add", name=node.name) + context.add(result) diff --git a/coremltools/converters/mil/frontend/torch/ssa_passes/__init__.py b/coremltools/converters/mil/frontend/torch/ssa_passes/__init__.py index 194ad37bb..2dac14c23 100644 --- a/coremltools/converters/mil/frontend/torch/ssa_passes/__init__.py +++ b/coremltools/converters/mil/frontend/torch/ssa_passes/__init__.py @@ -3,23 +3,4 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -# Import all passes in this dir -from os.path import dirname, basename, isfile, join -import glob - -excluded_files = [ - "__init__.py", - "torch_passes.py", -] -modules = glob.glob(join(dirname(__file__), "*.py")) -pass_modules = [ - basename(f)[:-3] - for f in modules - if isfile(f) - and basename(f)[:1] != "_" # Follow python convention to hide _* files. - and basename(f)[:4] != "test" - and basename(f) not in excluded_files -] -__all__ = pass_modules - -from . import * # import everything in __all__ +from . import torch_tensor_assign_to_core, torch_upsample_to_core_upsample diff --git a/coremltools/converters/mil/frontend/torch/ssa_passes/torch_passes.py b/coremltools/converters/mil/frontend/torch/ssa_passes/torch_passes.py index d210b3d3a..5c935b263 100644 --- a/coremltools/converters/mil/frontend/torch/ssa_passes/torch_passes.py +++ b/coremltools/converters/mil/frontend/torch/ssa_passes/torch_passes.py @@ -3,9 +3,10 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY import logging + from coremltools.converters._profile_utils import _profile +from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY @_profile diff --git a/coremltools/converters/mil/frontend/torch/ssa_passes/torch_tensor_assign_to_core.py b/coremltools/converters/mil/frontend/torch/ssa_passes/torch_tensor_assign_to_core.py index 64b9d7203..7b42cc694 100644 --- a/coremltools/converters/mil/frontend/torch/ssa_passes/torch_tensor_assign_to_core.py +++ b/coremltools/converters/mil/frontend/torch/ssa_passes/torch_tensor_assign_to_core.py @@ -5,23 +5,24 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb @register_pass(namespace="torch") -def torch_tensor_assign_to_core(prog): +class torch_tensor_assign_to_core(AbstractGraphPass): """ Try to map Torch dialect ops `torch_tensor_assign` into core op set if compatible. - We only support the single index selection + whole dimension slicing + stride 1 + We only support the single index selection + whole dimension slicing + stride 1 - Example 1: - x[0] = 0 + Example 1: + x[0] = 0 - Example 2: - x[:,:,0,:] = [[[0],[0]]] + Example 2: + x[:,:,0,:] = [[[0],[0]]] - Currently, we tranform the torch_tensor_assign op into transposes + scatter + expand_dims + Currently, we tranform the torch_tensor_assign op into transposes + scatter + expand_dims - Given: + Given: %output = torch_tensor_assign(data=%data, updates=%updates) Result: @@ -31,9 +32,11 @@ def torch_tensor_assign_to_core(prog): %output = transpose(%scatter) ... - """ - for f in prog.functions.values(): - _torch_tensor_assign_to_core_block(f) + """ + def apply(self, prog): + for f in prog.functions.values(): + _torch_tensor_assign_to_core_block(f) + def _torch_tensor_assign_to_core_block(block): for op in block.operations[:]: @@ -41,48 +44,49 @@ def _torch_tensor_assign_to_core_block(block): _torch_tensor_assign_to_core_block(b) if op.op_type in ["torch_tensor_assign"]: - with block: - _transform_tensor_assign(op, block) + with block: + _transform_tensor_assign(op, block) + def _transform_tensor_assign(op, block): - begin = op.begin.val - end = op.end.val - strides = op.stride.val - - begin_mask = op.begin_mask.val - end_mask = op.end_mask.val - squeeze_mask = op.squeeze_mask.val - - # check for the pattern is supported - if any([stride != 1 for stride in strides]): - raise NotImplementedError("Only tensor assignment with stride 1 is supported.") - - if sum(squeeze_mask) != 1: - raise NotImplementedError("Only tensor assignment with exactly 1 pure dimension selection is supported") - - for i in range(len(squeeze_mask)): - if not squeeze_mask[i]: - if not (begin_mask[i] or begin[i] == 0) or not end_mask[i]: - raise NotImplementedError("Non supported tensor assignment pattern detected.") - - # put the select dimension in front - # for instance, x[:,0] = ... - # we transpose the tensor to make the assignment be x[0] = ... instead - data = op.data - updates = op.updates - out_name = op.outputs[0].name - - select_dim = squeeze_mask.tolist().index(True) - perm = [select_dim] + [i for i in range(len(squeeze_mask)) if i != select_dim] - data = mb.transpose(x=data, perm=perm, before_op=op) - updates = mb.expand_dims(x=updates, axes=[0], before_op=op) - select_idx = begin[select_dim] - data = mb.scatter(data=data, indices=[select_idx], updates=updates, axis=0, mode="update", before_op=op) - perm_back = [perm.index(i) for i in range(len(perm))] - data = mb.transpose(x=data, perm=perm_back, name=out_name, before_op=op) - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=data - ) - # Remove all the ops at once - block.remove_ops([op]) + begin = op.begin.val + end = op.end.val + strides = op.stride.val + + begin_mask = op.begin_mask.val + end_mask = op.end_mask.val + squeeze_mask = op.squeeze_mask.val + + # check for the pattern is supported + if any([stride != 1 for stride in strides]): + raise NotImplementedError("Only tensor assignment with stride 1 is supported.") + + if sum(squeeze_mask) != 1: + raise NotImplementedError("Only tensor assignment with exactly 1 pure dimension selection is supported") + + for i in range(len(squeeze_mask)): + if not squeeze_mask[i]: + if not (begin_mask[i] or begin[i] == 0) or not end_mask[i]: + raise NotImplementedError("Non supported tensor assignment pattern detected.") + + # put the select dimension in front + # for instance, x[:,0] = ... + # we transpose the tensor to make the assignment be x[0] = ... instead + data = op.data + updates = op.updates + out_name = op.outputs[0].name + + select_dim = squeeze_mask.tolist().index(True) + perm = [select_dim] + [i for i in range(len(squeeze_mask)) if i != select_dim] + data = mb.transpose(x=data, perm=perm, before_op=op) + updates = mb.expand_dims(x=updates, axes=[0], before_op=op) + select_idx = begin[select_dim] + data = mb.scatter(data=data, indices=[select_idx], updates=updates, axis=0, mode="update", before_op=op) + perm_back = [perm.index(i) for i in range(len(perm))] + data = mb.transpose(x=data, perm=perm_back, name=out_name, before_op=op) + + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=data + ) + # Remove all the ops at once + block.remove_ops([op]) diff --git a/coremltools/converters/mil/frontend/torch/ssa_passes/torch_upsample_to_core_upsample.py b/coremltools/converters/mil/frontend/torch/ssa_passes/torch_upsample_to_core_upsample.py index 69183ac70..a5ad224ad 100644 --- a/coremltools/converters/mil/frontend/torch/ssa_passes/torch_upsample_to_core_upsample.py +++ b/coremltools/converters/mil/frontend/torch/ssa_passes/torch_upsample_to_core_upsample.py @@ -3,14 +3,15 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import logging -from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil import Builder as mb -import logging +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.pass_registry import register_pass @register_pass(namespace="torch") -def torch_upsample_to_core_upsample(prog): +class torch_upsample_to_core_upsample(AbstractGraphPass): """ Try to map Torch dialect ops `torch_upsample_nearest_neighbor` or `torch_upsample_bilinear` to `upsample_nearest_neighbor` or `upsample_bilinear` in the core op set if compatible. @@ -19,22 +20,24 @@ def torch_upsample_to_core_upsample(prog): prog: Program """ - for f in prog.functions.values(): - torch_upsample_to_core_upsample_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _torch_upsample_to_core_upsample_block(f) -def torch_upsample_to_core_upsample_block(block): +def _torch_upsample_to_core_upsample_block(block): for op in block.operations[:]: for b in op.blocks: - torch_upsample_to_core_upsample_block(b) + _torch_upsample_to_core_upsample_block(b) if op.op_type in ["torch_upsample_nearest_neighbor", "torch_upsample_bilinear"]: - if try_replace_with_core_upsample(op): + if _try_replace_with_core_upsample(op): logging.info("Successfully map {} to core upsample".format(op.op_type)) else: raise ValueError("Unable to map {} to core upsample".format(op.op_type)) -def try_get_upsample_factor(output_size): + +def _try_get_upsample_factor(output_size): # output_size = [ # (torch.floor((input.size(i + 2).float() * torch.tensor(scale_factors[i], dtype=torch.float32)).float())) # for i in range(dim) @@ -66,7 +69,8 @@ def try_get_upsample_factor(output_size): # we successfully trace back the original scale factor return op.y.val -def try_replace_with_core_upsample(op): + +def _try_replace_with_core_upsample(op): """ Inputs: @@ -77,8 +81,8 @@ def try_replace_with_core_upsample(op): True if op can be represented by mb.upsample_nearest_neighbor or mb.upsample_bilinear op in SSA. False otherwise """ - scales_h = try_get_upsample_factor(op.output_height.op) - scales_w = try_get_upsample_factor(op.output_width.op) + scales_h = _try_get_upsample_factor(op.output_height.op) + scales_w = _try_get_upsample_factor(op.output_width.op) if scales_h is None or scales_w is None: return False diff --git a/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py b/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py index 570e9d9e2..24b59aed5 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py @@ -3,18 +3,16 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -import itertools - -import numpy as np import pytest import torch import torch.nn as nn -from .testing_utils import * +from .testing_utils import convert_to_mlmodel, TorchBaseTest + # Custom layer imports -from coremltools.converters.mil.mil.ops.defs._op_reqs import * +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from coremltools.converters.mil.frontend.torch.torch_op_registry import ( register_torch_op, ) @@ -22,7 +20,18 @@ _TORCH_OPS_REGISTRY as _TORCH_OPS_REG, ) from coremltools.converters.mil.frontend.torch.ops import _get_inputs -from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil import ( + Builder as mb, + Operation, + types +) +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + DefaultInputs, + InputSpec, + TensorInputType +) + # Log Converter supported Cosine Similarity conversion function default_cosine_similarity = _TORCH_OPS_REG.get("cosine_similarity", None) @@ -110,7 +119,6 @@ def type_inference(self): # For illustration purpose, assumming getting valid shape # Ideally, should consider transpose_?, ?_is_sparse parameters into consideration # for computing output shape - ret_shape = [x_shape[0], y_shape[1]] return types.tensor(x_type, [x_shape[0], y_shape[1]]) @register_torch_op() @@ -144,14 +152,14 @@ def forward(self, x, y): "SparseMatMul" == layers[-1].custom.className ), "Custom Layer class name mis-match" assert ( - False == layers[-1].custom.parameters["transpose_x"].boolValue + not layers[-1].custom.parameters["transpose_x"].boolValue ), "Incorrect parameter value k" assert ( - False == layers[-1].custom.parameters["transpose_y"].boolValue + not layers[-1].custom.parameters["transpose_y"].boolValue ), "Incorrect parameter value k" assert ( - True == layers[-1].custom.parameters["x_is_sparse"].boolValue + layers[-1].custom.parameters["x_is_sparse"].boolValue ), "Incorrect parameter value k" assert ( - True == layers[-1].custom.parameters["y_is_sparse"].boolValue + layers[-1].custom.parameters["y_is_sparse"].boolValue ), "Incorrect parameter value k" diff --git a/coremltools/converters/mil/frontend/torch/test/test_internal_graph.py b/coremltools/converters/mil/frontend/torch/test/test_internal_graph.py index caf69af07..8b0356fe9 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_internal_graph.py +++ b/coremltools/converters/mil/frontend/torch/test/test_internal_graph.py @@ -239,7 +239,7 @@ def test_pow(self, context): self._test_elementwise_binary( context, "Pow", - ops.pow_, + ops.pow, [test_input_1, test_input_2], 2, np.power(test_input_1, test_input_2), diff --git a/coremltools/converters/mil/frontend/torch/test/test_passes.py b/coremltools/converters/mil/frontend/torch/test/test_passes.py index 7d799dfd1..b669137b4 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_passes.py +++ b/coremltools/converters/mil/frontend/torch/test/test_passes.py @@ -1,11 +1,23 @@ -import itertools +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +from collections import OrderedDict import numpy as np import pytest import torch -from ..internal_graph import * -from ..torchir_passes import * +from ..internal_graph import ( + InternalTorchIRBlock, + InternalTorchIRGraph, + InternalTorchIRNode +) +from ..torchir_passes import ( + flatten_graph_input_values, + flatten_graph_output_values, + transform_inplace_ops +) def _build_flattening_test_graph(): diff --git a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py index 62ab6998d..7d8d90738 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -6,12 +6,15 @@ import sys import itertools import numpy as np +import pytest +import torch.nn as nn +from .testing_utils import contains_op, ModuleWrapper, TorchBaseTest +from coremltools import RangeDim from coremltools.models.utils import _python_version from coremltools.models.utils import _macos_version from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * -from .testing_utils import * + from coremltools import TensorType from coremltools._deps import version_lt @@ -192,6 +195,40 @@ def test_sort(self, shape, axis, descending, backend): TorchBaseTest.run_compare_torch(shape, model, backend=backend) +class TestNorms(TorchBaseTest): + @pytest.mark.parametrize( + "shape, backend, keepdim", + itertools.product( + COMMON_SHAPES, + backends, + [True, False] + ) + ) + def test_frobenius_norm(self, shape, backend, keepdim): + num_dims = len(shape) + for dim in range(-num_dims, num_dims): + model = ModuleWrapper( + function=torch.norm, kwargs={'keepdim': keepdim, 'dim': dim} + ) + TorchBaseTest.run_compare_torch(shape, model, backend=backend) + + @pytest.mark.parametrize( + "shape, backend, p, keepdim", + itertools.product( + COMMON_SHAPES, + backends, + [1, 2, -1, 3, np.inf, -np.inf], + [True, False] + ) + ) + def test_number_norm(self, shape, backend, p, keepdim): + for dim in (-1, 0, 1): + model = ModuleWrapper( + function=torch.norm, kwargs={'p': p, 'keepdim': keepdim, 'dim': dim} + ) + TorchBaseTest.run_compare_torch(shape, model, backend=backend, places=2) + + class TestBatchNorm(TorchBaseTest): @pytest.mark.parametrize( "num_features, eps, affine, backend", @@ -727,8 +764,7 @@ def test_cond(self, use_cpu_for_conversion, backend): pytest.skip("rdar://81169758 (Cond tests hang on mlprogram backend)") if backend[0] == "mlprogram" and not use_cpu_for_conversion: pytest.xfail("rdar://78343191 ((MIL GPU) Core ML Tools Unit Test failures [failure to load or Seg fault])") - in_features = 1 - out_features = 2 + class TestNet(nn.Module): def forward(self, x): if torch.squeeze(x) < 10.: @@ -827,7 +863,7 @@ def test_upsample_bilinear2d_with_output_size( "scales_h, scales_w, align_corners, recompute_scale_factor, backend", itertools.product( [2, 0.5, 4.1], [3, 0.5, 5.3], [True, False], [True, False], backends - ) + ) ) def test_upsample_bilinear2d_with_scales( self, scales_h, scales_w, align_corners, recompute_scale_factor, backend @@ -843,7 +879,7 @@ def _is_float_value(x, threshold=0.001): is_h_float = _is_float_value(output_h) is_w_float = _is_float_value(output_w) - if (is_h_float or is_w_float) and not align_corners and not recompute_scale_factor: + if (is_h_float or is_w_float) and not align_corners and not recompute_scale_factor: pytest.xfail("rdar://81124053 (Support recompute_scale_factor)") model = ModuleWrapper( @@ -1455,7 +1491,6 @@ def test_lstm( input_size = 4 hidden_size = 6 num_layers = 1 - bias = True class Encoder(torch.nn.Module): def __init__(self): @@ -1547,8 +1582,6 @@ def test_lstm( SEQUENCE_LENGTH = 3 BATCH_SIZE = 2 - num_directions = int(bidirectional) + 1 - # (seq_len, batch, input_size) if batch_first: _input = torch.rand(BATCH_SIZE, SEQUENCE_LENGTH, input_size) @@ -1579,6 +1612,101 @@ def forward(self, x): model = TestNet() self.run_compare_torch((1, 3, 16, 16), model, backend=backend) +class TestFull(TorchBaseTest): + @pytest.mark.parametrize( + "backend, rank", + itertools.product( + backends, + [1, 3], + ), + ) + def test_full_dynamic(self, backend, rank): + class FullDynamicModel(nn.Module): + def __init__(self): + super(FullDynamicModel, self).__init__() + + def forward(self, x): + if rank == 1: + h = x[0] + x = torch.zeros(h) + elif rank == 3: + h, w, d = x[0], x[1], x[2] + x = torch.zeros(h, w, d) + return torch.full(x.shape, fill_value=3.14) + + input_shape = np.random.randint(low=2, high=6, size=rank) + torch_in = torch.tensor(input_shape) + model = FullDynamicModel().eval() + torch_out = model(torch_in) + self.run_compare_torch(torch_in, model, expected_results=torch_out, + input_as_shape=False, backend=backend) + + @pytest.mark.parametrize("shape_val, backend", + itertools.product( + [ + [(1,), 0.], + [(2, 3), 3.1415], + [(1, 1, 2, 5, 1), -2.], + ], + backends, + ) + ) + def test_full_static(self, shape_val, backend): + shape, val = shape_val + class FullStaticModel(nn.Module): + def __init__(self): + super(FullStaticModel, self).__init__() + + def forward(self, x): + return torch.full(x.shape, fill_value=val) + + self.run_compare_torch(shape, FullStaticModel().eval(), backend=backend) + +class TestOnes(TorchBaseTest): + @pytest.mark.parametrize( + "backend, rank", + itertools.product( + backends, + [1, 3], + ), + ) + def test_ones_dynamic(self, backend, rank): + class OnesDynamicModel(nn.Module): + def __init__(self): + super(OnesDynamicModel, self).__init__() + + def forward(self, x): + if rank == 1: + h = x[0] + x = torch.zeros(h) + elif rank == 3: + h, w, d = x[0], x[1], x[2] + x = torch.zeros(h, w, d) + return torch.ones(x.shape) + + input_shape = np.random.randint(low=2, high=6, size=rank) + torch_in = torch.tensor(input_shape) + model = OnesDynamicModel().eval() + torch_out = model(torch_in) + self.run_compare_torch(torch_in, model, expected_results=torch_out, + input_as_shape=False, backend=backend) + + @pytest.mark.parametrize("shape, backend", + itertools.product( + [(1,), (2, 3), (1, 1, 2, 5, 1)], + backends, + ) + ) + def test_ones_static(self, shape, backend): + class OnesStaticModel(nn.Module): + def __init__(self): + super(OnesStaticModel, self).__init__() + + def forward(self, x): + return torch.ones(x.shape) + + self.run_compare_torch(shape, OnesStaticModel().eval(), backend=backend) + class TestTypeAs(TorchBaseTest): @pytest.mark.parametrize("backend, type", itertools.product( @@ -1666,7 +1794,7 @@ class TestExpand(TorchBaseTest): "backend, shapes", itertools.product( backends, - [[(2, 1), (2, 2)], [(3, 1), (-1, 4)], [(1, 3, 4, 4), (3, 3, 4, 4)]] + [[(2, 1), (2, 2)], [(3, 1), (-1, 4)], [(1, 3, 4, 4), (3, 3, 4, 4)], [(4,), (3, 4)], [(3, 2), (1, 2, -1, 2)]] ), ) def test_expand(self, backend, shapes): @@ -1684,7 +1812,7 @@ def forward(self, x): "backend, input_shapes", itertools.product( backends, - [[(2, 1), (2, 2)], [(3, 1), (3, 4)], [(1, 3, 4, 4), (3, 3, 4, 4)]] + [[(2, 1), (2, 2)], [(3, 1), (3, 4)], [(1, 3, 4, 4), (3, 3, 4, 4)], [(4,), (1, 3, 4)]] ), ) def test_expand_as(self, backend, input_shapes): @@ -2787,6 +2915,30 @@ def forward(self, x): shape, model, backend=backend, ) + @pytest.mark.parametrize( + "backend", + backends, + ) + def test_index_put_case_3(self, backend): + pytest.xfail("rdar://84892125 (Empty tensors handling for non_zero, tile and scatter_nd)") + class IndexPutModel(torch.nn.Module): + def __init__(self): + super(IndexPutModel, self).__init__() + + def forward(self, x, y): + mask = y > 1 + x[y > 1] = 0. + return x + + inputs = [ + torch.Tensor([1., 2., 3., 4., 5., 6]), + torch.Tensor([0., 0., 0., 0., 0., 0.]), + ] + model = IndexPutModel() + self.run_compare_torch( + inputs, model, backend=backend, input_as_shape=False, + ) + class TestIndex(TorchBaseTest): @pytest.mark.parametrize( "backend, shape", @@ -3060,6 +3212,63 @@ def forward(self, x): self.run_compare_torch( shape, model, backend=backend, ) + + @pytest.mark.parametrize( + "backend, shape", + itertools.product( + backends, + [ + (1, 2, 3), + (2, 3, 4, 5), + ] + ), + ) + def test_index_int_index_case_9(self, backend, shape): + # one axis is sliced through bool mask + class IndexModel(torch.nn.Module): + def __init__(self): + super(IndexModel, self).__init__() + + def forward(self, x): + if len(shape) == 3: + return x[:, [True, False], :] + + elif len(shape) == 4: + return x[[True, False], :, :, :] + + model = IndexModel() + self.run_compare_torch( + shape, model, backend=backend, + ) + + @pytest.mark.parametrize( + "backend, shape", + itertools.product( + backends, + [ + (1, 2, 3), + (2, 3, 4, 5), + ] + ), + ) + def test_index_int_index_case_10(self, backend, shape): + # multiple axes are sliced through bool masks + class IndexModel(torch.nn.Module): + def __init__(self): + super(IndexModel, self).__init__() + + def forward(self, x): + if len(shape) == 3: + return x[[True], [True, False], [False, True, False]] + + elif len(shape) == 4: + return x[[True, True], :, [True, True, False, False], [True, False, False, True, False]] + + model = IndexModel() + self.run_compare_torch( + shape, model, backend=backend, + ) + class TestPad(TorchBaseTest): @pytest.mark.parametrize( "backend, rank, mode", @@ -3155,3 +3364,89 @@ def forward(self, rows, cols): self.run_compare_torch( inputs, model, expected_results, input_as_shape=False, backend=backend, ) + +class TestSacatterAdd(TorchBaseTest): + @pytest.mark.parametrize( + "shapes_dims, backend", + itertools.product( + [ + [(10,), (0, -1)], + [(2, 3), (1, -1)], + [(2, 3, 4, 5), (0, -2)], + ], + backends + ), + ) + def test_scatter_add(self, shapes_dims, backend): + shapes, dims = shapes_dims + for dim in dims: + + class TestModel(nn.Module): + def __init__(self): + super(TestModel, self).__init__() + self.source = torch.rand(*(shapes)) + self.index = torch.randint(0, shapes[dim], size=shapes) + + def forward(self, x): + index = torch.tensor(self.index) + return x.scatter_add_(dim, self.index, self.source) + + self.run_compare_torch(shapes, TestModel().eval(), backend=backend) + +class TestBroadcastTensors(TorchBaseTest): + @pytest.mark.parametrize( + "shapes, backend", + itertools.product( + [(1,), (1, 2)], + backends + ), + ) + def test_one_tensor(self, shapes, backend): + class TestModel(nn.Module): + def __init__(self): + super(TestModel, self).__init__() + + def forward(self, a): + return torch.broadcast_tensors(a) + self.run_compare_torch(shapes, TestModel().eval(), backend=backend) + + @pytest.mark.parametrize( + "shapes, backend", + itertools.product( + [ + [(2,1), (1,3)], + [(5,1,4,1), (3,1,1)], + [(1,), (3,1,7)], + [(2,1), (4,3,2,1,)] + ], + backends + ), + ) + def test_two_tensors(self, shapes, backend): + class TestModel(nn.Module): + def __init__(self): + super(TestModel, self).__init__() + + def forward(self, a, b): + return torch.broadcast_tensors(a, b) + self.run_compare_torch(shapes, TestModel().eval(), backend=backend) + + @pytest.mark.parametrize( + "shapes, backend", + itertools.product( + [ + [(2,1), (1,3), (1,), (1,1)], + [(5,1,4,1), (3,1,1), (1,), (4,8)], + [(1,), (2,1), (3,2,1), (5,4,3,2,1)], + ], + backends + ), + ) + def test_four_tensors(self, shapes, backend): + class TestModel(nn.Module): + def __init__(self): + super(TestModel, self).__init__() + + def forward(self, a, b, c, d): + return torch.broadcast_tensors(a, b, c, d) + self.run_compare_torch(shapes, TestModel().eval(), backend=backend) diff --git a/coremltools/converters/mil/frontend/torch/test/testing_utils.py b/coremltools/converters/mil/frontend/torch/test/testing_utils.py index cf819a3d4..218411e9d 100644 --- a/coremltools/converters/mil/frontend/torch/test/testing_utils.py +++ b/coremltools/converters/mil/frontend/torch/test/testing_utils.py @@ -22,7 +22,6 @@ class ModuleWrapper(nn.Module): """ Helper class to transform torch function into torch nn module. This helps to keep the testing interface same for torch functional api. - """ def __init__(self, function, kwargs=None): super(ModuleWrapper, self).__init__() @@ -52,9 +51,7 @@ def _copy_input_data(input_data): def contains_op(torch, op_string): - if hasattr(torch, op_string): - return True - return False + return hasattr(torch, op_string) def convert_to_coreml_inputs(input_description, inputs): @@ -215,15 +212,18 @@ def run_compare_torch( Traces a model and runs a numerical test. Args: input_as_shape : If true generates random input data with shape. - expected_results : Expected result from running pytorch model. - + expected_results : Expected result from running pytorch model. converter_input_type: If not None, then pass it to the "inputs" argument to the ct.convert() call. """ model.eval() if input_as_shape: input_data = generate_input_data(input_data, rand_range) - model_spec = torch.jit.script(model) if use_scripting else trace_model( - model, _copy_input_data(input_data)) + + if use_scripting: + model_spec = torch.jit.script(model) + else: + model_spec = trace_model(model, _copy_input_data(input_data)) + model_spec, mlmodel, coreml_inputs, coreml_results = \ convert_and_compare( input_data, model_spec, expected_results=expected_results, diff --git a/coremltools/converters/mil/frontend/torch/torch_op_registry.py b/coremltools/converters/mil/frontend/torch/torch_op_registry.py index 8ecc04d57..64407d3d0 100644 --- a/coremltools/converters/mil/frontend/torch/torch_op_registry.py +++ b/coremltools/converters/mil/frontend/torch/torch_op_registry.py @@ -18,6 +18,10 @@ def register_torch_op(_func=None, torch_alias=None, override=False): e.g. Sort aliased with SortV1, SortV2 All provided alias operators must not be registered previously. + "In place" alias are looked up automatically and do not need to + be registered. PyTorch uses an underscore suffix to denote the + in place version, e.g. "sum_" is the in place version of "sum". + override: (Boolean) [Default=False] If True, overrides earlier registration i.e. specified operator and alias will start pointing to current conversion diff --git a/coremltools/converters/mil/frontend/torch/torchir_passes.py b/coremltools/converters/mil/frontend/torch/torchir_passes.py index ab3313904..0b9e23713 100644 --- a/coremltools/converters/mil/frontend/torch/torchir_passes.py +++ b/coremltools/converters/mil/frontend/torch/torchir_passes.py @@ -1,7 +1,12 @@ +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + from collections import OrderedDict import logging as _logging -from .internal_graph import * +from .internal_graph import InternalTorchIRNode, InternalTorchIRGraph def generate_tensor_assignment_ops(graph): """ @@ -122,6 +127,32 @@ def _construct_nodes_to_fuse_inputs(nodes_to_fuse): graph.outputs[idx] = _get_updated_name(output, updated_tensor_count) +def remove_getattr_nodes(graph): + """Remove the getattr nodes in the graph + """ + + getattr_nodes = [] + new_nodes = [] + + for node in graph.nodes: + + for block in node.blocks: + remove_getattr_nodes(block) + + if node.kind == "getattr": + getattr_nodes.append(node) + else: + new_nodes.append(node) + + # check the getattr nodes not in the outputs + for node in getattr_nodes: + if node.name in graph.outputs: + raise RuntimeError("{} should not be in the graph outputs.".format(node.name)) + + # remove the getattr nodes + graph.nodes = new_nodes + + def transform_inplace_ops(graph, name_remap_dict=None): # As we modify ops, we'll need to remap symbols. diff --git a/coremltools/converters/mil/mil/__init__.py b/coremltools/converters/mil/mil/__init__.py index a10b65d58..b51b488ba 100644 --- a/coremltools/converters/mil/mil/__init__.py +++ b/coremltools/converters/mil/mil/__init__.py @@ -5,10 +5,42 @@ SPACES = " " from .block import curr_block, Block, Function -from .input_type import * -from .operation import * -from .program import * -from .var import * +from .input_type import ( + BoolInputType, + BoolTensorInputType, + DefaultInputs, + FloatInputType, + FloatTensorInputType, + InputSpec, + IntInputType, + IntOrFloatInputType, + IntOrFloatOrBoolInputType, + IntTensorInputType, + InternalInputType, + InternalScalarOrTensorInputType, + InternalStringInputType, + InternalVar, + ListInputType, + ListOrScalarOrTensorInputType, + PyFunctionInputType, + SUPPORT_FLOAT_TYPES, + SUPPORT_INT_TYPES, + ScalarOrTensorInputType, + StringInputType, + TensorInputType, + TupleInputType +) +from .operation import mil_list, precondition, Operation +from .program import ( + get_existing_symbol, + get_new_symbol, + get_new_variadic_symbol, + InputType, + Placeholder, + Program, + Symbol, +) +from .var import ListVar, Var from .builder import Builder from .ops.defs._op_reqs import register_op diff --git a/coremltools/converters/mil/mil/builder.py b/coremltools/converters/mil/mil/builder.py index b3b913ce4..5954f9b4f 100644 --- a/coremltools/converters/mil/mil/builder.py +++ b/coremltools/converters/mil/mil/builder.py @@ -12,7 +12,6 @@ from coremltools.converters.mil.mil.types.symbolic import any_symbolic from .program import Program, Placeholder from .block import curr_block, Function -from .operation import is_internal_input from .input_type import ( _InputType, InternalStringInputType, @@ -26,6 +25,7 @@ ) from .var import InternalVar, Var + def is_python_value(val): return ( isinstance(val, (np.generic, np.ndarray)) @@ -252,5 +252,5 @@ def wrapper(main_block): return wrapper -"""importing ops triggers installation of all ops into Builder""" +# importing ops triggers installation of all ops into Builder from .ops import defs as _ops diff --git a/coremltools/converters/mil/mil/operation.py b/coremltools/converters/mil/mil/operation.py index da46a991c..9bcb7545c 100644 --- a/coremltools/converters/mil/mil/operation.py +++ b/coremltools/converters/mil/mil/operation.py @@ -4,6 +4,7 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np + from coremltools.converters.mil.mil import types from coremltools.converters.mil.mil.types import is_compatible_type from coremltools.converters.mil.mil.types.symbolic import is_symbolic, any_symbolic diff --git a/coremltools/converters/mil/mil/ops/__init__.py b/coremltools/converters/mil/mil/ops/__init__.py index b4e0cff86..61aafff42 100644 --- a/coremltools/converters/mil/mil/ops/__init__.py +++ b/coremltools/converters/mil/mil/ops/__init__.py @@ -2,5 +2,3 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from .defs import * diff --git a/coremltools/converters/mil/mil/ops/defs/__init__.py b/coremltools/converters/mil/mil/ops/defs/__init__.py index bbc14cba8..676738078 100644 --- a/coremltools/converters/mil/mil/ops/defs/__init__.py +++ b/coremltools/converters/mil/mil/ops/defs/__init__.py @@ -3,19 +3,198 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from .activation import * +from .activation import ( + clamped_relu, + elu, + gelu, + leaky_relu, + linear_activation, + prelu, + relu, + relu6, + scaled_tanh, + sigmoid, + sigmoid_hard, + silu, + softmax, + softplus, + softplus_parametric, + softsign, + thresholded_relu, +) + from .classify import classify -from .control_flow import * -from .conv import * -from .elementwise_binary import * -from .elementwise_unary import * -from .image_resizing import * -from .linear import * -from .normalization import * -from .pool import * -from .random import * -from .recurrent import * -from .reduction import * -from .scatter_gather import * -from .tensor_operation import * -from .tensor_transformation import * + +from .control_flow import ( + cond, + const, + list_gather, + list_length, + list_read, + list_scatter, + list_write, + make_list, + select, + while_loop, +) + +from .conv import ( + conv, + conv_quantized, + conv_transpose, +) + +from .elementwise_binary import ( + add, + elementwise_binary, + equal, + floor_div, + greater, + greater_equal, + less, + less_equal, + logical_and, + logical_or, + logical_xor, + maximum, + minimum, + mod, + mul, + not_equal, + pow, + real_div, + sub, +) + +from .elementwise_unary import ( + abs, + acos, + asin, + atan, + atanh, + cast, + ceil, + clip, + cos, + cosh, + erf, + exp, + exp2, + floor, + inverse, + log, + logical_not, + round, + rsqrt, + sign, + sin, + sinh, + sqrt, + square, + tan, + tanh, + threshold, +) + +from .image_resizing import ( + affine, + crop, + crop_resize, + resample, + resize_bilinear, + resize_nearest_neighbor, + upsample_bilinear, + upsample_nearest_neighbor, +) + +from .linear import ( + einsum, + linear, + matmul, +) + +from .normalization import ( + batch_norm, + instance_norm, + l2_norm, + layer_norm, + local_response_norm, +) + +from .pool import ( + avg_pool, + max_pool, + l2_pool +) + +from .random import ( + random_bernoulli, + random_categorical, + random_normal, + random_uniform +) + +from .recurrent import ( + gru, + lstm, + rnn +) + +from .reduction import ( + reduce_argmax, + reduce_argmin, + reduce_l1_norm, + reduce_l2_norm, + reduce_log_sum, + reduce_log_sum_exp, + reduce_max, + reduce_mean, + reduce_min, + reduce_prod, + reduce_sum, + reduce_sum_square +) + +from .scatter_gather import ( + gather, + gather_along_axis, + gather_nd, + scatter, + scatter_along_axis, + scatter_nd, +) + +from .tensor_operation import ( + argsort, + band_part, + concat, + cumsum, + fill, + flatten2d, + identity, + non_maximum_suppression, + non_zero, + one_hot, + pad, + range_1d, + shape, + split, + stack, + tile, + topk, +) + +from .tensor_transformation import ( + depth_to_space, + expand_dims, + reshape, + reverse, + reverse_sequence, + slice_by_index, + slice_by_size, + space_to_depth, + squeeze, + transpose, + pixel_shuffle, + sliding_windows, +) diff --git a/coremltools/converters/mil/mil/ops/defs/_op_reqs.py b/coremltools/converters/mil/mil/ops/defs/_op_reqs.py index b88ef04eb..3e8100f38 100644 --- a/coremltools/converters/mil/mil/ops/defs/_op_reqs.py +++ b/coremltools/converters/mil/mil/ops/defs/_op_reqs.py @@ -2,9 +2,6 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil import types -from coremltools.converters.mil.mil import Operation, precondition, VALUE -from coremltools.converters.mil.mil.input_type import * -from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry +from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry as _SSAOpRegistry -register_op = SSAOpRegistry.register_op +register_op = _SSAOpRegistry.register_op diff --git a/coremltools/converters/mil/mil/ops/defs/_utils.py b/coremltools/converters/mil/mil/ops/defs/_utils.py index ab116029f..5cca93fc7 100644 --- a/coremltools/converters/mil/mil/ops/defs/_utils.py +++ b/coremltools/converters/mil/mil/ops/defs/_utils.py @@ -5,11 +5,9 @@ import math -import coremltools.converters - -from coremltools.converters.mil.mil import get_new_symbol +from coremltools.converters.mil.mil import get_new_symbol, types from coremltools.converters.mil.mil.types.symbolic import is_symbolic -from ._op_reqs import * + def broadcast_shapes(shape_x, shape_y): """ diff --git a/coremltools/converters/mil/mil/ops/defs/activation.py b/coremltools/converters/mil/mil/ops/defs/activation.py index 1e7739bbe..a4bc14b1b 100644 --- a/coremltools/converters/mil/mil/ops/defs/activation.py +++ b/coremltools/converters/mil/mil/ops/defs/activation.py @@ -5,7 +5,7 @@ import math import numpy as np -from coremltools.converters.mil.mil import Operation, types, VALUE +from coremltools.converters.mil.mil import types from coremltools.converters.mil.mil.input_type import ( DefaultInputs, FloatInputType, @@ -14,8 +14,8 @@ ScalarOrTensorInputType, StringInputType, TensorInputType, - ) -from coremltools.converters.mil.mil.operation import precondition +) +from coremltools.converters.mil.mil.operation import Operation, precondition, VALUE from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from .elementwise_unary import elementwise_unary diff --git a/coremltools/converters/mil/mil/ops/defs/classify.py b/coremltools/converters/mil/mil/ops/defs/classify.py index 0428fefae..4a2b73fae 100644 --- a/coremltools/converters/mil/mil/ops/defs/classify.py +++ b/coremltools/converters/mil/mil/ops/defs/classify.py @@ -2,11 +2,19 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + import numpy as np -from coremltools.converters.mil.mil.ops.defs._op_reqs import * +from coremltools.converters.mil.mil import types, Operation +from coremltools.converters.mil.mil.input_type import ( + InputSpec, + ListInputType, + TensorInputType +) +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from coremltools.converters.mil.mil.types.symbolic import any_symbolic + @register_op(doc_str="") class classify(Operation): """ diff --git a/coremltools/converters/mil/mil/ops/defs/control_flow.py b/coremltools/converters/mil/mil/ops/defs/control_flow.py index 5d8479439..b95175a4c 100644 --- a/coremltools/converters/mil/mil/ops/defs/control_flow.py +++ b/coremltools/converters/mil/mil/ops/defs/control_flow.py @@ -7,18 +7,12 @@ import numpy as np import logging - from coremltools.converters.mil.mil import ( Block, - DefaultInputs, get_new_symbol, get_existing_symbol, - SYMBOL, - NONE, types, - mil_list, - Operation, - VALUE + ) from coremltools.converters.mil.mil.input_type import ( BoolInputType, @@ -35,14 +29,19 @@ TupleInputType, StringInputType, ) -from coremltools.converters.mil.mil.operation import precondition +from coremltools.converters.mil.mil.operation import ( + mil_list, + NONE, + Operation, + precondition, + SYMBOL, + VALUE +) from coremltools.converters.mil.mil.types import is_compatible_type from coremltools.converters.mil.mil.types.type_mapping import ( numpy_val_to_builtin_val, is_subtype, ) - -from coremltools.converters.mil.mil.var import Var from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op @register_op(doc_str="") @@ -127,10 +126,10 @@ def value_inference(self): return [v.val for v in self.blocks[0].outputs] return [v.val for v in self.blocks[1].outputs] -@register_op(doc_str="") -class const(Operation): + +class Const(Operation): """ - Return constant values. + A base class that returns constant values. Parameters ---------- @@ -155,7 +154,7 @@ class const(Operation): ) def __init__(self, **kwargs): - super(const, self).__init__(**kwargs) + super(Const, self).__init__(**kwargs) def type_inference(self): builtin_type, _ = self._get_type_val(self.val.val) @@ -197,6 +196,9 @@ def _get_type_val(self, value): if len(list_value) == 0: raise ValueError("'mil_list' points to an empty list") builtin_elem_type, _ = self._get_type_val(list_value[0]) + # mil_list is a special case that we want to preserve the int64 element type + if isinstance(list_value[0], np.int64): + builtin_elem_type = types.int64 from coremltools.converters.mil.mil.types.type_list import list as types_list builtin_type = types_list(builtin_elem_type, init_length=len(list_value), dynamic_length=False) return builtin_type, value @@ -208,6 +210,11 @@ def _get_type_val(self, value): _, builtin_type = numpy_val_to_builtin_val(value) return builtin_type, value +@register_op(doc_str="") +class const(Const): + def __init__(self, **kwargs): + super(const, self).__init__(**kwargs) + # Internal const can have symbolic value (for testing purpose) @register_op(doc_str="") diff --git a/coremltools/converters/mil/mil/ops/defs/conv.py b/coremltools/converters/mil/mil/ops/defs/conv.py index 3337b87f6..dd5f80450 100644 --- a/coremltools/converters/mil/mil/ops/defs/conv.py +++ b/coremltools/converters/mil/mil/ops/defs/conv.py @@ -9,7 +9,6 @@ IntInputType, IntTensorInputType, Operation, - OrderedDict, TensorInputType, types, ScalarOrTensorInputType, @@ -157,7 +156,6 @@ def type_inference(self): inshape = self.x.shape f_shape = self.weight.shape kernel_shape = f_shape[2:] - num_dims = len(inshape) - 2 C_out = f_shape[0] C_in = self.x.shape[1] groups = self.groups.val diff --git a/coremltools/converters/mil/mil/ops/defs/elementwise_binary.py b/coremltools/converters/mil/mil/ops/defs/elementwise_binary.py index 794221c89..b9ea899b5 100644 --- a/coremltools/converters/mil/mil/ops/defs/elementwise_binary.py +++ b/coremltools/converters/mil/mil/ops/defs/elementwise_binary.py @@ -3,10 +3,18 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np - import operator -from ._op_reqs import * + +from ._op_reqs import register_op from ._utils import promoted_primitive_type, broadcast_shapes +from coremltools.converters.mil.mil import ( + InputSpec, + Operation, + precondition, + ScalarOrTensorInputType, + types +) +from coremltools.converters.mil.mil.operation import VALUE class elementwise_binary(Operation): diff --git a/coremltools/converters/mil/mil/ops/defs/elementwise_unary.py b/coremltools/converters/mil/mil/ops/defs/elementwise_unary.py index b68a1a24e..0bea9e211 100644 --- a/coremltools/converters/mil/mil/ops/defs/elementwise_unary.py +++ b/coremltools/converters/mil/mil/ops/defs/elementwise_unary.py @@ -5,9 +5,9 @@ import math import numpy as np -from coremltools.converters.mil.mil import Operation, types, VALUE, SYMBOL +from coremltools.converters.mil.mil import types +from coremltools.converters.mil.mil.operation import Operation, precondition, SYMBOL, VALUE from coremltools.converters.mil.mil.types.symbolic import is_symbolic -from coremltools.converters.mil.mil.operation import precondition from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from coremltools.converters.mil.mil.input_type import ( DefaultInputs, @@ -321,7 +321,7 @@ def value_inference(self): @register_op(doc_str="") class exp(elementwise_unary): """ - Return the exponential values of the input ``x``, element-wise. + Return e^x, element-wise. Parameters ---------- @@ -348,7 +348,7 @@ def value_inference(self): @register_op(doc_str="") class exp2(elementwise_unary): """ - Return the exponential values of the input ``x``, element-wise. + Return 2^x, element-wise. Parameters ---------- @@ -518,7 +518,8 @@ def value_inference(self): @register_op(doc_str="") class round(elementwise_unary): """ - Return the round value of the input ``x``, element-wise. + Return the round value of the input ``x`` to nearest integer, element-wise. + ``0.5`` is rounded to ``0``. Parameters ---------- @@ -591,6 +592,8 @@ class sign(elementwise_unary): """ Return the sign value of the input ``x``, element-wise. + All elements in the output will be either ``-1``. or ``1``. + Parameters ---------- x: tensor<[\*d], T> (Required) @@ -697,7 +700,7 @@ def value_inference(self): @register_op(doc_str="") class square(elementwise_unary): """ - Return the square value of the input ``x``, element-wise. + Return ``x^2``, element-wise. Parameters ---------- diff --git a/coremltools/converters/mil/mil/ops/defs/image_resizing.py b/coremltools/converters/mil/mil/ops/defs/image_resizing.py index db9fb870d..49dca34d1 100644 --- a/coremltools/converters/mil/mil/ops/defs/image_resizing.py +++ b/coremltools/converters/mil/mil/ops/defs/image_resizing.py @@ -5,10 +5,23 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np -from ._op_reqs import * -from coremltools.converters.mil.mil.types.symbolic import is_symbolic -from coremltools.converters.mil.mil import get_new_symbol +from ._op_reqs import register_op +from coremltools.converters.mil.mil import ( + BoolInputType, + DefaultInputs, + FloatInputType, + get_new_symbol, + InputSpec, + IntInputType, + IntTensorInputType, + IntOrFloatInputType, + Operation, + StringInputType, + TensorInputType, + types, +) +from coremltools.converters.mil.mil.types.symbolic import is_symbolic @register_op(doc_str="") class affine(Operation): diff --git a/coremltools/converters/mil/mil/ops/defs/linear.py b/coremltools/converters/mil/mil/ops/defs/linear.py index e4dddeca6..8d37c21dc 100644 --- a/coremltools/converters/mil/mil/ops/defs/linear.py +++ b/coremltools/converters/mil/mil/ops/defs/linear.py @@ -4,8 +4,20 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np +from coremltools.converters.mil.mil import ( + BoolInputType, + DefaultInputs, + InputSpec, + Operation, + precondition, + StringInputType, + TensorInputType, + TupleInputType, + types, +) +from coremltools.converters.mil.mil.operation import VALUE from coremltools.converters.mil.mil.types.symbolic import is_symbolic -from ._op_reqs import * +from ._op_reqs import register_op from ._utils import broadcast_shapes, parse_einsum_equation @register_op(doc_str="") diff --git a/coremltools/converters/mil/mil/ops/defs/normalization.py b/coremltools/converters/mil/mil/ops/defs/normalization.py index 7b870a839..9f065a664 100644 --- a/coremltools/converters/mil/mil/ops/defs/normalization.py +++ b/coremltools/converters/mil/mil/ops/defs/normalization.py @@ -4,7 +4,19 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np -from ._op_reqs import * +from ._op_reqs import register_op +from coremltools.converters.mil.mil import ( + DefaultInputs, + FloatInputType, + InputSpec, + IntInputType, + IntTensorInputType, + Operation, + precondition, + TensorInputType, + types, +) +from coremltools.converters.mil.mil.operation import VALUE from coremltools.converters.mil.mil.types.symbolic import ( any_symbolic, ) @@ -292,7 +304,7 @@ class local_response_norm(Operation): Apply local response normalization to the n-dimensional input tensor: .. math:: - x_i \\leftarrow \\dfrac{x_i}{\\left ( k + \\dfrac{\\alpha}{\text{size}} \\sum_j x_j^2 \\right )^\\beta} + x_i \\leftarrow \\dfrac{x_i}{\\left ( k + \\dfrac{\\alpha}{\\text{size}} \\sum_j x_j^2 \\right )^\\beta} Parameters diff --git a/coremltools/converters/mil/mil/ops/defs/pool.py b/coremltools/converters/mil/mil/ops/defs/pool.py index b7efcdeec..174ca76de 100644 --- a/coremltools/converters/mil/mil/ops/defs/pool.py +++ b/coremltools/converters/mil/mil/ops/defs/pool.py @@ -3,16 +3,23 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +from coremltools.converters.mil.mil import Operation, types +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + DefaultInputs, + InputSpec, + IntTensorInputType, + TensorInputType, + StringInputType +) +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from coremltools.converters.mil.mil.ops.defs._utils import spatial_dimensions_out_shape -from ._op_reqs import * - -""" -Pooling Op Superclass -""" - class Pooling(Operation): + """ + Pooling Op Superclass + """ input_spec = InputSpec( x=TensorInputType(), kernel_sizes=IntTensorInputType(const=True), @@ -151,12 +158,11 @@ class avg_pool(Pooling): input_spec = ( InputSpec( exclude_padding_from_average=BoolInputType(const=True, - optional=True)) + optional=True)) + Pooling.input_spec ) def default_inputs(self): - num_spatial_dims = self.x.rank - 2 return super().default_inputs() + \ DefaultInputs( exclude_padding_from_average=False, diff --git a/coremltools/converters/mil/mil/ops/defs/random.py b/coremltools/converters/mil/mil/ops/defs/random.py index 3863dacde..36518b8fe 100644 --- a/coremltools/converters/mil/mil/ops/defs/random.py +++ b/coremltools/converters/mil/mil/ops/defs/random.py @@ -3,16 +3,27 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +from coremltools.converters.mil.mil import get_new_symbol, get_new_variadic_symbol, types +from coremltools.converters.mil.mil.input_type import ( + DefaultInputs, + FloatInputType, + InputSpec, + IntInputType, + IntTensorInputType, + TensorInputType, + StringInputType +) + from coremltools.converters.mil.mil.types.symbolic import any_symbolic -from coremltools.converters.mil.mil import get_new_symbol, get_new_variadic_symbol -from ._op_reqs import * +from coremltools.converters.mil.mil.operation import Operation -""" -Random Op Superclass -""" +from ._op_reqs import register_op class RandomDistribution(Operation): + """ + Random Op Superclass + """ input_spec = InputSpec(shape=IntTensorInputType(),) out_dtype = types.fp32 diff --git a/coremltools/converters/mil/mil/ops/defs/recurrent.py b/coremltools/converters/mil/mil/ops/defs/recurrent.py index 1af6bd3fd..2f3c2af4c 100644 --- a/coremltools/converters/mil/mil/ops/defs/recurrent.py +++ b/coremltools/converters/mil/mil/ops/defs/recurrent.py @@ -3,8 +3,16 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil import get_new_symbol -from ._op_reqs import * +from ._op_reqs import register_op +from coremltools.converters.mil.mil import Operation, types +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + DefaultInputs, + FloatInputType, + InputSpec, + TensorInputType, + StringInputType +) @register_op(doc_str="") @@ -26,7 +34,7 @@ class gru(Operation): Where: - * ``W_{ir}``, ``W_{io}``, and `` W_{iz}`` state input-hidden weight for reset, output + * ``W_{ir}``, ``W_{io}``, and ``W_{iz}`` state input-hidden weight for reset, output and update gate, respectively. * ``W_{h[r|o|z]}`` are recurrent weights on hidden state to reset, output, update gate. * ``h_t`` is the hidden state at time ``t``. @@ -130,7 +138,7 @@ def type_inference(self): if self.weight_hh.rank != 2: raise ValueError( "Invalid weight shape. Expecting Rank 2 input, got {}".format( - len(self.self.weight_hh.rank) + len(self.weight_hh.rank) ) ) @@ -149,8 +157,8 @@ def type_inference(self): if hidden_size != (hidden_dim // dim_factor): raise ValueError( "Incorrect weight matrix: hidden dim size mismatch. \ - Provided weight_ih {},weight_hh {}. Expecting ").format( - weight_ih.shape, weight_hh.shape + Provided weight_ih {}, weight_hh {}. Expecting ").format( + self.weight_ih.shape, self.weight_hh.shape ) out_seq_len = sequence_length if self.output_sequence.val else 1 diff --git a/coremltools/converters/mil/mil/ops/defs/reduction.py b/coremltools/converters/mil/mil/ops/defs/reduction.py index f3b310b61..c7bb5fdae 100644 --- a/coremltools/converters/mil/mil/ops/defs/reduction.py +++ b/coremltools/converters/mil/mil/ops/defs/reduction.py @@ -9,7 +9,6 @@ Operation, precondition, types, - VALUE ) from coremltools.converters.mil.mil.input_type import ( BoolInputType, @@ -19,6 +18,7 @@ IntTensorInputType, TensorInputType ) +from coremltools.converters.mil.mil.operation import VALUE class ReductionAxes(Operation): @@ -116,12 +116,7 @@ def get_operator(self): raise NotImplementedError() -""" -Reduction op implementations -""" - - -@register_op(doc_str="TODO") +@register_op(doc_str="") class reduce_arg(ReductionAxis): def __init__(self, **kwargs): super(reduce_arg, self).__init__(**kwargs) @@ -140,6 +135,10 @@ def type_inference(self): return types.tensor(types.int32, tuple(reduced_shape)) +""" +Reduction op implementations +""" + @register_op(doc_str="") class reduce_argmax(reduce_arg): """ @@ -546,7 +545,7 @@ class reduce_sum(ReductionAxes): axes: const (Optional, default="None", reduce on all axes.) * The dimensions to reduce. - + keep_dims: const (Optional, default=False) * If ``False``, the rank is reduced by ``1`` for each entry in ``axes``, otherwise retain reduced axes with length ``1``. @@ -561,7 +560,7 @@ class reduce_sum(ReductionAxes): T: f32, int32 """ - + def __init__(self, **kwargs): super(reduce_sum, self).__init__(**kwargs) diff --git a/coremltools/converters/mil/mil/ops/defs/scatter_gather.py b/coremltools/converters/mil/mil/ops/defs/scatter_gather.py index 8e6bad17e..f7337218d 100644 --- a/coremltools/converters/mil/mil/ops/defs/scatter_gather.py +++ b/coremltools/converters/mil/mil/ops/defs/scatter_gather.py @@ -5,7 +5,7 @@ import numpy as np import numbers -from coremltools.converters.mil.mil import Operation, types, SYMBOL, VALUE +from coremltools.converters.mil.mil import Operation, types from coremltools.converters.mil.mil.input_type import ( DefaultInputs, InputSpec, @@ -18,6 +18,11 @@ from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from coremltools.converters.mil.mil.types.symbolic import is_compatible_symbolic_vector, is_symbolic +from coremltools.converters.mil.mil.operation import ( + SYMBOL, + VALUE +) + @register_op(doc_str="") class gather(Operation): @@ -181,6 +186,7 @@ class scatter(Operation): mode: const string (Optional) * Can be the following modes: ``update``, ``add``, ``sub``, ``mul``, ``div``, ``max``, ``min``. + * Default value is ``update``. Returns ------- @@ -224,13 +230,9 @@ def type_inference(self): ) err = "Updates shape {} is incorrect. It should be {}.".format(self.updates.shape, expected_updates_shape) - if len(self.updates.shape) == len(expected_updates_shape): - for dim1, dim2 in zip(self.updates.shape, expected_updates_shape): - if not is_symbolic(dim1) and not is_symbolic(dim2): - if dim1 != dim2: - raise ValueError(err) - else: - raise ValueError(err) + assert is_compatible_symbolic_vector( + self.updates.shape, tuple(expected_updates_shape) + ), err return self.data.sym_type @@ -417,7 +419,9 @@ def type_inference(self): axis = self.axis.val axis = axis if axis >= 0 else axis + self.data.rank - assert self.indices.shape == self.updates.shape + assert is_compatible_symbolic_vector( + self.indices.shape, self.updates.shape + ) assert self.data.rank == self.indices.rank for i in range(self.data.rank): if i != axis: diff --git a/coremltools/converters/mil/mil/ops/defs/tensor_operation.py b/coremltools/converters/mil/mil/ops/defs/tensor_operation.py index c2d190c52..8cdfe662e 100644 --- a/coremltools/converters/mil/mil/ops/defs/tensor_operation.py +++ b/coremltools/converters/mil/mil/ops/defs/tensor_operation.py @@ -14,13 +14,32 @@ from coremltools.converters.mil.mil import ( get_new_symbol, get_new_variadic_symbol, - SYMBOL, - VALUE, + types +) +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + DefaultInputs, + FloatInputType, + InputSpec, + IntInputType, + IntOrFloatInputType, + IntOrFloatOrBoolInputType, + IntTensorInputType, + ListOrScalarOrTensorInputType, + ScalarOrTensorInputType, + StringInputType, + TensorInputType, + TupleInputType, +) +from coremltools.converters.mil.mil.operation import ( NONE, + Operation, + precondition, + SYMBOL, + VALUE ) -from coremltools.converters.mil.mil import types -from ._op_reqs import * +from ._op_reqs import register_op from ._utils import promoted_primitive_type @@ -611,7 +630,6 @@ def type_inference(self): ) raise ValueError(msg.format(len(reps), self.x.rank)) - out_shape = [] for i, rep in enumerate(reps): if not is_symbolic(rep): @@ -919,7 +937,7 @@ class concat(Operation): input_spec = InputSpec(values=TupleInputType(), axis=IntInputType(const=True), interleave=BoolInputType(const=True, - optional=True)) + optional=True)) def default_inputs(self): return DefaultInputs( @@ -1029,6 +1047,7 @@ def value_inference(self): return np.concatenate(values, axis=self.axis.val) + @register_op(doc_str="") class split(Operation): """ @@ -1181,7 +1200,7 @@ class stack(Operation): """ input_spec = InputSpec(values=TupleInputType(), - axis=IntInputType(const=True),) + axis=IntInputType(const=True),) def __init__(self, **kwargs): super(stack, self).__init__(**kwargs) diff --git a/coremltools/converters/mil/mil/ops/defs/tensor_transformation.py b/coremltools/converters/mil/mil/ops/defs/tensor_transformation.py index f4a1b5f69..2344420d4 100644 --- a/coremltools/converters/mil/mil/ops/defs/tensor_transformation.py +++ b/coremltools/converters/mil/mil/ops/defs/tensor_transformation.py @@ -18,9 +18,7 @@ get_new_variadic_symbol, Operation, precondition, - SYMBOL, types, - VALUE, ) from coremltools.converters.mil.mil.input_type import ( BoolTensorInputType, @@ -31,6 +29,11 @@ ScalarOrTensorInputType, TensorInputType ) +from coremltools.converters.mil.mil.operation import ( + SYMBOL, + VALUE +) + from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op def _solve_slice_by_index_shape(x_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask): diff --git a/coremltools/converters/mil/mil/ops/registry.py b/coremltools/converters/mil/mil/ops/registry.py index d15cbb948..e2d790956 100644 --- a/coremltools/converters/mil/mil/ops/registry.py +++ b/coremltools/converters/mil/mil/ops/registry.py @@ -15,7 +15,7 @@ class SSAOpRegistry: custom_ops = {} @staticmethod - def register_op(doc_str="", is_custom_op=False, namespace="core"): + def register_op(doc_str="", is_custom_op=False, namespace="core", allow_override=False): """ Registration routine for MIL Program operators is_custom_op: (Boolean) [Default=False] @@ -25,6 +25,8 @@ def register_op(doc_str="", is_custom_op=False, namespace="core"): Current operator is registered as `SSARegistry.custom_ops` Otherwise, current operator is registered as usual operator, i.e. registered in `SSARegistry.ops'. + allow_override: (Boolean) [Default=False] + If True, it is allowed for an operation to override the previous operation with the same op name. """ def class_wrapper(op_cls): @@ -40,7 +42,7 @@ def class_wrapper(op_cls): logging.debug("Registering {} {}".format(op_msg, op_type)) - if op_type in op_reg: + if op_type in op_reg and not allow_override: raise ValueError( "SSA {} {} already registered.".format(op_msg, op_type) ) diff --git a/coremltools/converters/mil/mil/ops/tests/test_activation.py b/coremltools/converters/mil/mil/ops/tests/test_activation.py index 01ea0d49b..d6a57a8ea 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_activation.py +++ b/coremltools/converters/mil/mil/ops/tests/test_activation.py @@ -800,7 +800,6 @@ def test_builder_eval3(self): @ssa_fn def test_builder_eval4(self): - x_val = np.array([[[-1, 3, 6]], [[-1, 2, -3]], [[4, -5, 6]]], dtype=np.float32) with pytest.raises(ValueError, match=r"x .* rank 3"): v = mb.softplus_parametric( x=[1], diff --git a/coremltools/converters/mil/mil/ops/tests/test_const.py b/coremltools/converters/mil/mil/ops/tests/test_const.py index d099cfd7f..735cadce0 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_const.py +++ b/coremltools/converters/mil/mil/ops/tests/test_const.py @@ -2,9 +2,12 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import numpy as np +import pytest from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +from coremltools.converters.mil.mil import Builder as mb, types from .testing_utils import run_compare_builder diff --git a/coremltools/converters/mil/mil/ops/tests/test_control_flow.py b/coremltools/converters/mil/mil/ops/tests/test_control_flow.py index e75e03c17..4beb14a45 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_control_flow.py +++ b/coremltools/converters/mil/mil/ops/tests/test_control_flow.py @@ -2,12 +2,14 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import pytest +import numpy as np -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * from .testing_utils import run_compare_builder, UNK_SYM - -backends = testing_reqs.backends +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import ssa_fn, random_gen class TestSelect: @@ -215,7 +217,7 @@ def cond(res, bx): return mb.less(x=bx, y=b) res, ignored = mb.while_loop(_cond=cond, _body=body, - loop_vars=([1.], [0.])) + loop_vars=([1.], [0.])) return res input_values = { @@ -276,11 +278,11 @@ def cond1(i, j): def body1(i, j): new_i = mb.while_loop(_cond=cond2, _body=body2, - loop_vars=(i,)) + loop_vars=(i,)) return mb.add(x=new_i, y=two), j return mb.while_loop(_cond=cond1, _body=body1, - loop_vars=(x, y)) + loop_vars=(x, y)) input_values = { "x": np.array([0], dtype=np.float32), @@ -307,6 +309,7 @@ def body1(i, j): backend=backend, ) + class TestList: @pytest.mark.parametrize( "use_cpu_only, backend", itertools.product([True, False], backends,) diff --git a/coremltools/converters/mil/mil/ops/tests/test_conv.py b/coremltools/converters/mil/mil/ops/tests/test_conv.py index d5f454ad2..60304b713 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_conv.py +++ b/coremltools/converters/mil/mil/ops/tests/test_conv.py @@ -3,13 +3,18 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.mil import get_new_symbol -from coremltools.converters.mil.testing_reqs import * +import itertools +import numpy as np +import pytest from .testing_utils import run_compare_builder - -backends = testing_reqs.backends +from coremltools.converters.mil import testing_reqs +from coremltools.converters.mil.mil import ( + Builder as mb, + get_new_symbol, + types +) +from coremltools.converters.mil.testing_reqs import backends class TestConvTranspose: @@ -227,8 +232,6 @@ def test_builder_to_backend_stress( groups, symbolic, ): - - D, H, W, Kd, Kh, Kw = DHWKdKhKw N, C_in, C_out = 1, 1 * groups, 2 * groups @@ -424,7 +427,6 @@ def test_builder_to_backend_stress_weights_input( input_shape = [N, C_in, H, W] paddings = [padding[0], padding[0], padding[1], padding[1]] - wts = m.state_dict() weight = wts["weight"].detach().numpy() bias = wts["bias"].detach().numpy() if has_bias else None diff --git a/coremltools/converters/mil/mil/ops/tests/test_elementwise_binary.py b/coremltools/converters/mil/mil/ops/tests/test_elementwise_binary.py index d1951cc1d..bfe4560e0 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_elementwise_binary.py +++ b/coremltools/converters/mil/mil/ops/tests/test_elementwise_binary.py @@ -2,13 +2,14 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +import itertools +import numpy as np +import pytest from .testing_utils import run_compare_builder - -backends = testing_reqs.backends +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import ssa_fn class TestElementwiseBinary: diff --git a/coremltools/converters/mil/mil/ops/tests/test_image_resizing.py b/coremltools/converters/mil/mil/ops/tests/test_image_resizing.py index 7a165dcdb..96826de50 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_image_resizing.py +++ b/coremltools/converters/mil/mil/ops/tests/test_image_resizing.py @@ -3,13 +3,23 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.mil import get_new_symbol -from coremltools.converters.mil.testing_reqs import * +import functools +import itertools +import numpy as np +import pytest from .testing_utils import run_compare_builder +from coremltools.converters.mil import testing_reqs +from coremltools.converters.mil.mil import ( + Builder as mb, + get_new_symbol, + types +) +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import random_gen -backends = testing_reqs.backends +if testing_reqs._HAS_TORCH: + import torch class TestAffine: @@ -439,7 +449,6 @@ def build_upsample_fractional(x): backend=backend, ) - @pytest.mark.skipif(not testing_reqs._HAS_TORCH, reason="PyTorch not installed.") @pytest.mark.parametrize( "use_cpu_only, backend, input_shape, scale_factor, align_corners, recompute_scale_factor", @@ -796,9 +805,8 @@ def build(x, mode=0): np.array([3.5, 5.5, 11.5, 13.5], dtype=np.float32).reshape(1, 1, 1, 2, 2), ] - import functools for mode in range(6): - ## nn-proto does not support UNALIGN_CORNERS + # nn-proto does not support UNALIGN_CORNERS if not (backend[0] == 'neuralnetwork' and mode == 5): run_compare_builder( functools.partial(build, mode=mode), diff --git a/coremltools/converters/mil/mil/ops/tests/test_linear.py b/coremltools/converters/mil/mil/ops/tests/test_linear.py index 837c4cc39..bb37ed611 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_linear.py +++ b/coremltools/converters/mil/mil/ops/tests/test_linear.py @@ -2,13 +2,14 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +import itertools +import pytest +import numpy as np from .testing_utils import run_compare_builder - -backends = testing_reqs.backends +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import ssa_fn, random_gen class TestLinear: @@ -209,6 +210,7 @@ def test_builder_y_rank_2_const(self, use_cpu_only, backend, shape_x): "x": mb.placeholder(shape=x_val.shape), } input_values = {"x": x_val} + def build(x): return [mb.matmul(x=x, y=y_val, transpose_x=False, transpose_y=False)] @@ -310,8 +312,6 @@ def build(x, y): backend=backend, ) - - @ssa_fn def test_builder_eval(self): x_val = np.arange(6).astype(np.float32).reshape((1, 3, 2)) diff --git a/coremltools/converters/mil/mil/ops/tests/test_normalization.py b/coremltools/converters/mil/mil/ops/tests/test_normalization.py index 562733601..4f598e804 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_normalization.py +++ b/coremltools/converters/mil/mil/ops/tests/test_normalization.py @@ -3,13 +3,26 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * -from coremltools.converters.mil.mil import Program, Function, get_new_symbol +import itertools +import numpy as np +from numpy import linalg as la +import pytest from .testing_utils import UNK_SYM, run_compare_builder +from coremltools.converters.mil import testing_reqs +from coremltools.converters.mil.mil import ( + Builder as mb, + Function, + get_new_symbol, + types +) +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import random_gen -backends = testing_reqs.backends +if testing_reqs._HAS_TORCH: + import torch +if testing_reqs._HAS_TF_2: + import tensorflow as tf class TestNormalizationBatchNorm: @@ -353,7 +366,7 @@ def build(x): # V2->V1 lowering (op_mappings.py): else branch mb.layer_norm(x=x, axes=[-2, -1], epsilon=1e-4), # V2->V1 lowering (op_mappings.py): if branch with scale - mb.layer_norm(x=x, axes=[2], epsilon=1e-4, gamma=gamma_val, beta=beta_val), + mb.layer_norm(x=x, axes=[2], epsilon=1e-4, gamma=gamma_val, beta=beta_val), ] expected_output_types = [(1, 3, 2, types.fp32), (1, 3, 2, types.fp32), (1, 3, 2, types.fp32)] @@ -487,11 +500,10 @@ def build(x): @pytest.mark.parametrize( "use_cpu_only, backend, rank_and_axes, epsilon, provides_gamma_beta", - itertools.product([True, False], backends, - [[3,[0,2]], [3,[-2]], [4,[0,1,3]], [5,[0,4]], [5,[-5,-4,-3,-2,-1]] - ], - [0.0001, 0.01], - [True, False]), + itertools.product([True, False], backends, + [[3,[0,2]], [3,[-2]], [4,[0,1,3]], [5,[0,4]], [5,[-5,-4,-3,-2,-1]]], + [0.0001, 0.01], + [True, False]), ) def test_builder_to_backend_stress_numpy(self, use_cpu_only, backend, rank_and_axes, epsilon, provides_gamma_beta): @@ -507,7 +519,7 @@ def test_builder_to_backend_stress_numpy(self, use_cpu_only, backend, rank_and_a gamma, beta = None, None if provides_gamma_beta: - positive_axes = [axis+rank if axis <0 else axis for axis in axes] + positive_axes = [axis+rank if axis < 0 else axis for axis in axes] normalized_shape = [shape[i] for i in range(rank) if i in positive_axes] gamma = random_gen(shape=normalized_shape, rand_min=-100, rand_max=100) beta = random_gen(shape=normalized_shape, rand_min=-100, rand_max=100) @@ -538,10 +550,9 @@ def build(x): @pytest.mark.skipif(not testing_reqs._HAS_TF_2, reason="Tensorflow not found.") @pytest.mark.parametrize( "use_cpu_only, backend, rank_and_axes, epsilon", - itertools.product([True, False], backends, - [[3,[0,2]], [3,[-2]], [4,[0,1,3]], [5,[0,4]], [5,[-5,-4,-3,-2,-1]] - ], - [0.0001, 0.01]), + itertools.product([True, False], backends, + [[3,[0,2]], [3,[-2]], [4,[0,1,3]], [5,[0,4]], [5,[-5,-4,-3,-2,-1]]], + [0.0001, 0.01]), ) def test_builder_to_backend_stress_keras(self, use_cpu_only, backend, rank_and_axes, epsilon): rank, axes = rank_and_axes @@ -581,11 +592,11 @@ def test_builder_eval_stress(self, rank_and_axes, epsilon): rank, axes = rank_and_axes shape = np.random.randint(low=2, high=6, size=rank) x_val = random_gen(shape=shape, rand_min=-100.0, rand_max=100.0) - positive_axes = [axis+rank if axis <0 else axis for axis in axes] + positive_axes = [axis+rank if axis < 0 else axis for axis in axes] normalized_shape = [shape[i] for i in range(rank) if i in positive_axes] gamma_val = random_gen(shape=normalized_shape, rand_min=-100, rand_max=100) beta_val = random_gen(shape=normalized_shape, rand_min=-100, rand_max=100) - with Function({}) as ssa_func: + with Function({}): res = mb.layer_norm(x=x_val, axes=axes, epsilon=epsilon, gamma=gamma_val, beta=beta_val) ref = TestNormalizationLayerNorm._np_layer_norm(x=x_val, axes=axes, epsilon=epsilon, gamma=gamma_val, beta=beta_val) np.testing.assert_allclose(ref, res.val, atol=1e-04, rtol=1e-05) diff --git a/coremltools/converters/mil/mil/ops/tests/test_pool.py b/coremltools/converters/mil/mil/ops/tests/test_pool.py index 99fbe2318..4747380f6 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_pool.py +++ b/coremltools/converters/mil/mil/ops/tests/test_pool.py @@ -3,12 +3,13 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +import itertools +import numpy as np +import pytest from .testing_utils import run_compare_builder - -backends = testing_reqs.backends +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends class TestAvgPool: diff --git a/coremltools/converters/mil/mil/ops/tests/test_random.py b/coremltools/converters/mil/mil/ops/tests/test_random.py index 1ed343aa8..7bc172e1b 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_random.py +++ b/coremltools/converters/mil/mil/ops/tests/test_random.py @@ -2,15 +2,17 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + +import itertools +import numpy as np +import pytest import unittest -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +from .testing_utils import UNK_SYM, run_compare_builder +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends from coremltools.converters.mil.testing_utils import get_core_ml_prediction from coremltools.models.utils import _macos_version -from .testing_utils import UNK_SYM, run_compare_builder - -backends = testing_reqs.backends class TestRandomBernoulli: diff --git a/coremltools/converters/mil/mil/ops/tests/test_recurrent.py b/coremltools/converters/mil/mil/ops/tests/test_recurrent.py index ad55991d8..7014db34c 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_recurrent.py +++ b/coremltools/converters/mil/mil/ops/tests/test_recurrent.py @@ -3,13 +3,21 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.mil import get_new_symbol -from coremltools.converters.mil.testing_reqs import * +import itertools +import numpy as np +import pytest from .testing_utils import run_compare_builder +from coremltools.converters.mil import testing_reqs +from coremltools.converters.mil.mil import ( + Builder as mb, + get_new_symbol, + types +) +from coremltools.converters.mil.testing_reqs import backends -backends = testing_reqs.backends +if testing_reqs._HAS_TORCH: + import torch class TestGRU: @@ -105,7 +113,6 @@ def get_numpy_prediction_gru(X, H, return_seq, direction, output = np.transpose(output, (1, 0, 2)) return output, output[-1, :, :] - def get_numpy_prediction_gru_single_batch(X, h, return_seq, direction, inner_activation_str='SIGMOID', activation_str='TANH'): @@ -452,7 +459,7 @@ def ifzo_to_ifoz(x): n_t = torch.flip(n_t, [0]) output, (hn, cn) = rnn(n_t, (h0, c0)) - if output_sequence == False: + if not output_sequence: output = output[-1].unsqueeze(0) output = output.detach().numpy() @@ -596,7 +603,7 @@ def ifzo_to_ifoz(x): c0 = torch.randn(2, batch_size, hidden_size) output, (hn, cn) = rnn(t, (h0, c0)) - if output_sequence == False: + if not output_sequence: output_f = output[-1].unsqueeze(0)[:, :, :hidden_size] output_r = output[0].unsqueeze(0)[:, :, hidden_size:] output = torch.cat([output_f, output_r], dim=2) @@ -728,7 +735,7 @@ def test_builder_to_backend_smoke( n_t = torch.flip(n_t, [0]) output, hn = rnn(n_t, h0) - if output_sequence == False: + if not output_sequence: output = output[-1].unsqueeze(0) output = output.detach().numpy() diff --git a/coremltools/converters/mil/mil/ops/tests/test_scatter_gather.py b/coremltools/converters/mil/mil/ops/tests/test_scatter_gather.py index b112bb209..447e815f3 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_scatter_gather.py +++ b/coremltools/converters/mil/mil/ops/tests/test_scatter_gather.py @@ -3,12 +3,19 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +import itertools +import pytest +import numpy as np from .testing_utils import run_compare_builder +from coremltools._deps import _HAS_TF_1, MSG_TF1_NOT_FOUND +from coremltools.converters.mil import testing_reqs +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import ssa_fn -backends = testing_reqs.backends +if _HAS_TF_1: + import tensorflow as tf class TestScatter: diff --git a/coremltools/converters/mil/mil/ops/tests/test_slice.py b/coremltools/converters/mil/mil/ops/tests/test_slice.py index 3de30b5ca..012062a95 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_slice.py +++ b/coremltools/converters/mil/mil/ops/tests/test_slice.py @@ -2,13 +2,14 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.testing_reqs import * +import itertools +import pytest +import numpy as np from .testing_utils import UNK_SYM, run_compare_builder - -backends = testing_reqs.backends +from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import ssa_fn class TestSliceByIndex: diff --git a/coremltools/converters/mil/mil/ops/tests/test_tensor_operation.py b/coremltools/converters/mil/mil/ops/tests/test_tensor_operation.py index 163b99b23..3ffef6a45 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_tensor_operation.py +++ b/coremltools/converters/mil/mil/ops/tests/test_tensor_operation.py @@ -2,13 +2,22 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import numpy as np +import pytest from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.mil import types, get_new_symbol -from coremltools.converters.mil.testing_reqs import * - +from coremltools.converters.mil.mil import ( + Builder as mb, + get_new_symbol, + types +) +from coremltools.converters.mil.testing_utils import random_gen, ssa_fn from .testing_utils import UNK_SYM, UNK_VARIADIC, run_compare_builder +if testing_reqs._HAS_TF_1 or testing_reqs._HAS_TF_2: + from coremltools.converters.mil.testing_reqs import tf + backends = testing_reqs.backends diff --git a/coremltools/converters/mil/mil/ops/tests/test_tensor_transformation.py b/coremltools/converters/mil/mil/ops/tests/test_tensor_transformation.py index e6af34aaa..02f3e7bff 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_tensor_transformation.py +++ b/coremltools/converters/mil/mil/ops/tests/test_tensor_transformation.py @@ -2,14 +2,23 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from coremltools.converters.mil import testing_reqs -from coremltools.converters.mil.mil import get_new_symbol, get_new_variadic_symbol -from coremltools.converters.mil.testing_reqs import * +import itertools +import pytest +import numpy as np from .testing_utils import UNK_SYM, UNK_VARIADIC, run_compare_builder +from coremltools._deps import _HAS_TORCH +from coremltools.converters.mil import testing_reqs +from coremltools.converters.mil.mil import ( + Builder as mb, + get_new_symbol, + types +) +from coremltools.converters.mil.testing_reqs import backends +from coremltools.converters.mil.testing_utils import ssa_fn -backends = testing_reqs.backends +if _HAS_TORCH: + import torch class TestDepthToSpace: @@ -286,7 +295,6 @@ def test_builder_eval(self): def test_builder_to_backend_symbolic(self, use_cpu_only, backend): s0 = get_new_symbol() s_len = get_new_symbol() - s1 = get_new_variadic_symbol() input_placeholders = { "x": mb.placeholder(shape=(2, s0)), @@ -654,6 +662,7 @@ def test_builder_eval_rank_0(self): assert type(v.val) == np.float32 assert np.isclose(np.squeeze(x), v.val) + class TestTranspose: @pytest.mark.parametrize( argnames=["use_cpu_only", "backend", "is_symbolic"], @@ -921,7 +930,6 @@ def build(x, y): ) def test_builder_to_backend_type_promotion(self, use_cpu_only, backend): t1 = np.array([[1, 2], [4, 5]], dtype=np.float32) - t2 = np.array([[7, 8]], dtype=np.float32) input_placeholders = { "x": mb.placeholder(shape=t1.shape), @@ -1065,7 +1073,7 @@ def test_builder_eval_failure(self): np.random.rand(1, 1, 3, 1), ] with pytest.raises(ValueError): - v = mb.concat(values=values, axis=2) + mb.concat(values=values, axis=2) class TestSplit: diff --git a/coremltools/converters/mil/mil/ops/tests/testing_utils.py b/coremltools/converters/mil/mil/ops/tests/testing_utils.py index b5e6379d9..a8b4eb45c 100644 --- a/coremltools/converters/mil/mil/ops/tests/testing_utils.py +++ b/coremltools/converters/mil/mil/ops/tests/testing_utils.py @@ -58,9 +58,6 @@ def run_compare_builder( That is, "ct.convert(...., useCPUOnly=use_cpu_for_conversion)" It forces the model to be loaded on the CPU context, post conversion. """ - from coremltools.converters.mil.testing_reqs import ct - - if not isinstance(expected_output_types, list): expected_output_types = [expected_output_types] diff --git a/coremltools/converters/mil/mil/passes/__init__.py b/coremltools/converters/mil/mil/passes/__init__.py index 8b45ffd20..901dc952e 100644 --- a/coremltools/converters/mil/mil/passes/__init__.py +++ b/coremltools/converters/mil/mil/passes/__init__.py @@ -3,23 +3,45 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -# Import all passes in this dir -from os.path import dirname, basename, isfile, join -import glob - -excluded_files = ["pass_registry.py", "apply_common_pass_pipeline.py", "__init__.py"] -modules = glob.glob(join(dirname(__file__), "*.py")) -pass_modules = [ - basename(f)[:-3] - for f in modules - if isfile(f) - and basename(f)[:1] != "_" # Follow python convention to hide _* files. - and basename(f)[:4] != "test" - and basename(f) not in excluded_files -] -__all__ = pass_modules - -from . import * # import everything in __all__ +from . import ( + add_conv_transpose_output_shape, + cast_optimization, + concat_to_pixel_shuffle, + const_elimination, + conv_batchnorm_fusion, + conv_bias_fusion, + conv_scale_fusion, + dead_code_elimination, + dedup_op_and_var_names, + detect_concat_interleave, + divide_to_multiply, + elementwise_batchnorm_fusion, + gelu_exact_fusion, + gelu_tanh_approximation_fusion, + graph_pass, + helper, + image_input_preprocessing, + layernorm_instancenorm_pattern_fusion, + leaky_relu_fusion, + linear_bias_fusion, + loop_invariant_elimination, + matmul_weight_bias_fusion, + merge_consecutive_paddings, + name_sanitization_utils, + noop_elimination, + onehot_matmul_to_gather, + pad_conv_connect, + quantization_passes, + rank0_expand_dims_swap, + reduce_mean_fusion, + reduce_transposes, + remove_redundant_ops, + remove_symbolic_reshape, + replace_stack_reshape, + sanitize_input_output_names, + topological_reorder, + use_reflection_padding +) from coremltools.converters.mil.experimental.passes import ( generic_gelu_tanh_approximation_fusion, @@ -28,4 +50,4 @@ generic_conv_batchnorm_fusion, generic_conv_scale_fusion, generic_conv_bias_fusion -) \ No newline at end of file +) diff --git a/coremltools/converters/mil/mil/passes/add_conv_transpose_output_shape.py b/coremltools/converters/mil/mil/passes/add_conv_transpose_output_shape.py index 2cc4ff64f..32170a290 100644 --- a/coremltools/converters/mil/mil/passes/add_conv_transpose_output_shape.py +++ b/coremltools/converters/mil/mil/passes/add_conv_transpose_output_shape.py @@ -1,18 +1,16 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.types.symbolic import any_symbolic @register_pass(namespace="common") -def add_conv_transpose_output_shape(prog): +class add_conv_transpose_output_shape(AbstractGraphPass): """ ``conv_transpose`` input ``output_shape`` is an optional input. Since we can infer the output shape from type_inference, we add @@ -26,20 +24,21 @@ def add_conv_transpose_output_shape(prog): %2: (3, i32) = const(val=[1,5,39]) %3: (1, 5, 39, fp32) = conv_transpose(..., output_shape=%2) """ - for f in prog.functions.values(): - handle_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _handle_block(f) -def match_pattern(op): +def _match_pattern(op): return op.op_type == "conv_transpose" \ and op.output_shape is None \ and not any_symbolic(op.outputs[0].shape) -def handle_block(block): +def _handle_block(block): for op in list(block.operations): for b in op.blocks: - handle_block(b) + _handle_block(b) - if not match_pattern(op): + if not _match_pattern(op): continue # matched pattern diff --git a/coremltools/converters/mil/mil/passes/apply_common_pass_pipeline.py b/coremltools/converters/mil/mil/passes/apply_common_pass_pipeline.py index d4df177a4..4b7c67718 100644 --- a/coremltools/converters/mil/mil/passes/apply_common_pass_pipeline.py +++ b/coremltools/converters/mil/mil/passes/apply_common_pass_pipeline.py @@ -25,8 +25,8 @@ def _apply(passes, name="common"): s = 'passes' if len(passes) > 1 else 'pass' for p in _tqdm(passes, desc="Running MIL {} {}".format(name, s), unit=" passes"): _logging.info('Performing pass: "{}"'.format(p)) - - PASS_REGISTRY[p](prog) if not isinstance(p, AbstractQuantizationPass) else p.apply(prog) + graph_pass = PASS_REGISTRY[p] if not isinstance(p, AbstractQuantizationPass) else p + graph_pass(prog) if isinstance(p, AbstractQuantizationPass) or not isinstance(PASS_REGISTRY[p], PassContainer): prog.validate() @@ -67,6 +67,11 @@ def _apply(passes, name="common"): "common::fuse_conv_bias", # Re-run the fuse conv bias pass after the conv and batch_norm are fused "common::detect_concat_interleave", "common::concat_to_pixel_shuffle", # should come after detect_concat_interleave and after replace_stack_reshape + # "remove_redundant_ops" pass should be applied towards the end, once other graph passes have done their optimizations. + # For instance, it should come after passes such as "reduce_transpose" that can introduce redundant transposes + # in the network (while reducing the total number of transposes), and after passes such as "fuse_layernorm_or_instancenorm" + # which detects patterns that involve redundant ops ("sub") etc. + "common::remove_redundant_ops", "common::dead_code_elimination", # always end with dce ] @@ -82,7 +87,8 @@ def _apply(passes, name="common"): "common::loop_invariant_elimination", "common::noop_elimination", "common::dedup_op_and_var_names", - "common::reduce_transposes", # fuse_layernorm_or_instancenorm can potentially adding transposes + "common::reduce_transposes", # fuse_layernorm_or_instancenorm can potentially add transposes + "common::remove_redundant_ops", "common::topological_reorder", "common::dead_code_elimination", # always end with dce ] diff --git a/coremltools/converters/mil/mil/passes/cast_optimization.py b/coremltools/converters/mil/mil/passes/cast_optimization.py index 036266e04..691363100 100644 --- a/coremltools/converters/mil/mil/passes/cast_optimization.py +++ b/coremltools/converters/mil/mil/passes/cast_optimization.py @@ -1,21 +1,14 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from __future__ import print_function as _ -from __future__ import division as _ -from __future__ import absolute_import as _ - from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb -import numpy as np - @register_pass(namespace="common") -def cast_optimization(prog): +class cast_optimization(AbstractGraphPass): """ This optimization pass, - Removes redundant cast op i.e cast where source and destination tensors have same dtypes. @@ -36,34 +29,34 @@ def cast_optimization(prog): The input graph has maximum precision of fp16 while the output graph has fp32 precision. """ - for f in prog.functions.values(): - block_changed = True - cached_vars = {} - """ - Cached vars is used when all the following conditions are met: - 1. When the output of a cast gets fed into multiple casts of same configuration - 2. And, these 2 consecutive casts can be fused into a single cast. - When above conditions are satisfied, we create a NEW fused cast op ONLY once and - the output of all these consecutive casts gets replaced with the ouptut of this fused cast. - - Input graph: - |---->cast(dtype="fp16")---->square--->out_1 - | - input---->cast(dtype="int32")---->cast(dtype="fp16")---->relu--->out_2 - | - |---->cast(dtype="fp16")---->log--->out_3 - - Output graph: - - |---->square--->out_1 - | - input---->new_fused_cast(dtype="fp16")---->relu--->out_2 - | - |---->log--->out_3 - """ - while block_changed: - block_changed = fuse_or_cancel_consecutive_casts_block(f, cached_vars) - + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + cached_vars = {} + """ + Cached vars is used when all the following conditions are met: + 1. When the output of a cast gets fed into multiple casts of same configuration + 2. And, these 2 consecutive casts can be fused into a single cast. + When above conditions are satisfied, we create a NEW fused cast op ONLY once and + the output of all these consecutive casts gets replaced with the ouptut of this fused cast. + + Input graph: + |---->cast(dtype="fp16")---->square--->out_1 + | + input---->cast(dtype="int32")---->cast(dtype="fp16")---->relu--->out_2 + | + |---->cast(dtype="fp16")---->log--->out_3 + + Output graph: + + |---->square--->out_1 + | + input---->new_fused_cast(dtype="fp16")---->relu--->out_2 + | + |---->log--->out_3 + """ + while block_changed: + block_changed = _fuse_or_cancel_consecutive_casts_block(f, cached_vars) class Node(object): def __init__(self, op_type, match_criterion=None): @@ -115,7 +108,7 @@ def _match_linear_pattern(root, pattern): return [] -def try_to_transform(root_op, cached_vars): +def _try_to_transform(root_op, cached_vars): block = root_op.enclosing_block # Scenario: Redundant cast when source and destination dtype are same. @@ -175,7 +168,7 @@ def try_to_transform(root_op, cached_vars): return True -def fuse_or_cancel_consecutive_casts_block(block, cached_vars): +def _fuse_or_cancel_consecutive_casts_block(block, cached_vars): block_changed = False for i, op in enumerate(list(block.operations)): for b in op.blocks: @@ -183,7 +176,7 @@ def fuse_or_cancel_consecutive_casts_block(block, cached_vars): nested_block_cached_vars = {} nested_block_cached_vars.update(cached_vars) while nested_block_changed: - nested_block_changed = fuse_or_cancel_consecutive_casts_block(b, nested_block_cached_vars) + nested_block_changed = _fuse_or_cancel_consecutive_casts_block(b, nested_block_cached_vars) if len(op.blocks) > 0: continue @@ -191,9 +184,8 @@ def fuse_or_cancel_consecutive_casts_block(block, cached_vars): # start pattern match if cast op is encountered if op.op_type == "cast": with block: - block_changed = try_to_transform(op, cached_vars) + block_changed = _try_to_transform(op, cached_vars) # has to break as the downstream iterator is affected. if block_changed: return block_changed return block_changed - diff --git a/coremltools/converters/mil/mil/passes/concat_to_pixel_shuffle.py b/coremltools/converters/mil/mil/passes/concat_to_pixel_shuffle.py index ced4d05f8..eebd25234 100644 --- a/coremltools/converters/mil/mil/passes/concat_to_pixel_shuffle.py +++ b/coremltools/converters/mil/mil/passes/concat_to_pixel_shuffle.py @@ -5,7 +5,7 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass - +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass def _match_pattern(op): @@ -25,7 +25,7 @@ def _match_pattern(op): inputs = list(w_concat.inputs["values"]) if len(inputs) != 2: return None - + if not inputs[0].op or not inputs[1].op: return None @@ -54,6 +54,7 @@ def _match_pattern(op): return w_concat, h_concat_0, h_concat_1 + def _replace_ops(block, w_concat, h_concat_0, h_concat_1): with block: @@ -80,14 +81,14 @@ def _concat_to_pixel_shuffle_block(block): layers = _match_pattern(op) if layers: _replace_ops(block, layers[0], layers[1], layers[2]) - + @register_pass(namespace="common") -def concat_to_pixel_shuffle(prog): +class concat_to_pixel_shuffle(AbstractGraphPass): """ - Identify nested, interleaved concats which can be replaced by a single concat and a pixel shuffle layer. + Identify nested, interleaved concats which can be replaced by a single concat and a pixel shuffle layer. - This pattern occurs with the faster up-convolution from the FCRN model (Laina et al., 2016). + This pattern occurs with the faster up-convolution from the FCRN model (Laina et al., 2016). input(N, C, H, W) ------------------- | @@ -106,12 +107,13 @@ def concat_to_pixel_shuffle(prog): | v input(N, C, H, W) -----> concat(axis=1, interleave=True) -----> pixel_shuffle(upscale_factor=2) ----> output - ^ - | + ^ + | input(N, C, H, W) --------------| | | input(N, C, H, W) --------------- """ - for f in prog.functions.values(): - _concat_to_pixel_shuffle_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _concat_to_pixel_shuffle_block(f) diff --git a/coremltools/converters/mil/mil/passes/const_elimination.py b/coremltools/converters/mil/mil/passes/const_elimination.py index 50dc07328..7aa53c436 100644 --- a/coremltools/converters/mil/mil/passes/const_elimination.py +++ b/coremltools/converters/mil/mil/passes/const_elimination.py @@ -1,46 +1,13 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass - -def const_elimination_block(block): - # shallow copy hides changes on f.operations during the loop - for op in list(block.operations): - if op.op_type == "const": - continue - - for b in op.blocks: - const_elimination_block(b) - - all_outputs_are_const = True - for i, o in enumerate(op.outputs): - if o.val is not None: - with block: - res = mb.const( - val=o.val, - before_op=op, - # same var name, but different python - # instance does not violate SSA property. - name=o.name, - ) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=o, new_var=res - ) - # rename the const output - o.set_name(o.name+'_ignored') - else: - all_outputs_are_const = False - - if all_outputs_are_const: - op.remove_from_block() - +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass @register_pass(namespace="common") -def const_elimination(prog): +class const_elimination(AbstractGraphPass): """ prog: Program @@ -57,5 +24,61 @@ def const_elimination(prog): # %4 = other_op(%2_const, %3) # """ - for f in prog.functions.values(): - const_elimination_block(f) + + def _const_elimination_block(self, block, ops_to_ignore): + # shallow copy hides changes on f.operations during the loop + for op in list(block.operations): + + if op in ops_to_ignore: + continue + + for b in op.blocks: + self._const_elimination_block(b, ops_to_ignore) + + all_outputs_are_const = True + for i, o in enumerate(op.outputs): + if o.val is not None: + with block: + res = mb.const( + val=o.val, + before_op=op, + # same var name, but different python + # instance does not violate SSA property. + name=o.name, + ) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=o, new_var=res + ) + # rename the const output + o.set_name(o.name+'_ignored') + else: + all_outputs_are_const = False + + if all_outputs_are_const: + op.remove_from_block() + + def _get_ops_to_ignore(self, prog): + """ + utility function to get the ops which cannot be removed in the const elimination pass, which is all the const ops. + """ + ops_to_ignore = set() + + def _get_ops_to_ignore_block(block): + + for op in list(block.operations): + + for b in op.blocks: + _get_ops_to_ignore_block(b) + + if op.op_type == "const": + ops_to_ignore.add(op) + + for f in prog.functions.values(): + _get_ops_to_ignore_block(f) + + return ops_to_ignore + + def apply(self, prog): + ops_to_ignore = self._get_ops_to_ignore(prog) + for f in prog.functions.values(): + self._const_elimination_block(f, ops_to_ignore) diff --git a/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py b/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py index 703ccc3c5..02986e6c3 100644 --- a/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py +++ b/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py @@ -7,10 +7,10 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb import numpy as np - def _try_to_transform(conv_op, bn_op, block): # get parameters from batch_norm layer @@ -150,7 +150,7 @@ def _match_pattern(op): @register_pass(namespace="common") -def fuse_conv_batchnorm(prog): +class fuse_conv_batchnorm(AbstractGraphPass): """ Fuse the following batch_norm layer into conv and conv_transpose That is, convert conv + batch_norm to conv, by modifying the weight and bias in the conv layer @@ -165,7 +165,8 @@ def fuse_conv_batchnorm(prog): ... """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_conv_batchnorm_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_conv_batchnorm_block(f) diff --git a/coremltools/converters/mil/mil/passes/conv_bias_fusion.py b/coremltools/converters/mil/mil/passes/conv_bias_fusion.py index 23ec16661..0a88ddec1 100644 --- a/coremltools/converters/mil/mil/passes/conv_bias_fusion.py +++ b/coremltools/converters/mil/mil/passes/conv_bias_fusion.py @@ -7,13 +7,13 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import types from .helper import _check_child_op_type import numpy as np import logging - child_op_types = ["add", "sub"] def _match_pattern(op): @@ -273,7 +273,7 @@ def _fuse_conv_bias_block(block): @register_pass(namespace="common") -def fuse_conv_bias(prog): +class fuse_conv_bias(AbstractGraphPass): """ Fold add/sub into bias of conv and conv_transpose That is, convert conv + add/sub to conv, when add/sub is adding a constant @@ -306,7 +306,8 @@ def fuse_conv_bias(prog): ... """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_conv_bias_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_conv_bias_block(f) diff --git a/coremltools/converters/mil/mil/passes/conv_scale_fusion.py b/coremltools/converters/mil/mil/passes/conv_scale_fusion.py index 850ee4f7f..1c9e5e476 100644 --- a/coremltools/converters/mil/mil/passes/conv_scale_fusion.py +++ b/coremltools/converters/mil/mil/passes/conv_scale_fusion.py @@ -7,10 +7,10 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb import numpy as np - def _try_to_transform(conv_op, scale_op, block): # get the scale @@ -167,9 +167,8 @@ def _match_pattern(op): return fusion_occurred return fusion_occurred - @register_pass(namespace="common") -def fuse_conv_scale(prog): +class fuse_conv_scale(AbstractGraphPass): """ Fold mul/div into conv/conv_transpose by updating the weight/bias of the convolution layers. @@ -188,7 +187,8 @@ def fuse_conv_scale(prog): ... """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_conv_scale_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_conv_scale_block(f) diff --git a/coremltools/converters/mil/mil/passes/dead_code_elimination.py b/coremltools/converters/mil/mil/passes/dead_code_elimination.py index 2d5f5b8d0..853214d0e 100644 --- a/coremltools/converters/mil/mil/passes/dead_code_elimination.py +++ b/coremltools/converters/mil/mil/passes/dead_code_elimination.py @@ -3,11 +3,13 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil.passes.pass_registry import register_pass import logging +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass + -def dead_code_elimination_block(block): +def _dead_code_elimination_block(block): used_vars = set() ops_to_remove = list() @@ -28,7 +30,7 @@ def dead_code_elimination_block(block): used_vars.update([input_var]) for b in op.blocks: - used_in_block = dead_code_elimination_block(b) + used_in_block = _dead_code_elimination_block(b) used_vars.update(used_in_block) for op in ops_to_remove: @@ -37,9 +39,8 @@ def dead_code_elimination_block(block): return used_vars - @register_pass(namespace="common") -def dead_code_elimination(program): +class dead_code_elimination(AbstractGraphPass): """ Eliminate unused ops in program. @@ -79,6 +80,6 @@ def dead_code_elimination(program): In this example, %matmul_0 is an op that's not used in the computation, this op and its input ops (%tx_0 and %ty_0) are eliminated in this pass. """ - - for f in program.functions.values(): - dead_code_elimination_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _dead_code_elimination_block(f) diff --git a/coremltools/converters/mil/mil/passes/dedup_op_and_var_names.py b/coremltools/converters/mil/mil/passes/dedup_op_and_var_names.py index 69eba74ff..c1fdfb10d 100644 --- a/coremltools/converters/mil/mil/passes/dedup_op_and_var_names.py +++ b/coremltools/converters/mil/mil/passes/dedup_op_and_var_names.py @@ -4,6 +4,7 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Function import re import collections @@ -19,6 +20,7 @@ def _gen_new_name(seen_names, curr_name): if new_name not in seen_names: return new_name + def _deduplicate_block(block, func_outputs, seen_var_names, seen_op_names): """ seen_var_names: set[str] @@ -43,7 +45,8 @@ def _deduplicate_block(block, func_outputs, seen_var_names, seen_op_names): v.name = _gen_new_name(seen_var_names, v.name) seen_var_names.add(v.name) -def ensure_unique_var_names(v_set): + +def _ensure_unique_var_names(v_set): """ v_set: set[Variable] @@ -58,7 +61,7 @@ def ensure_unique_var_names(v_set): 'function\'s input and output') @register_pass(namespace="common") -def dedup_op_and_var_names(prog): +class dedup_op_and_var_names(AbstractGraphPass): """ For each function, this pass renames ops and variables with the same name as any preceding ops/variables across all scopes in the given function, @@ -73,14 +76,14 @@ def dedup_op_and_var_names(prog): %input = some_op(...) return %input """ - - for _, func in prog.functions.items(): - # Handle function input/outputs as they cannot be changed (to maintain - # user interface) - inputs = list(func.function_inputs) - io_vars = set(inputs + func.outputs) - ensure_unique_var_names(io_vars) - seen_var_names = set([v.name for v in io_vars]) - seen_op_names = set() - _deduplicate_block(func, set(func.outputs), - seen_var_names, seen_op_names) + def apply(self, prog): + for func in prog.functions.values(): + # Handle function input/outputs as they cannot be changed (to maintain + # user interface) + inputs = list(func.function_inputs) + io_vars = set(inputs + func.outputs) + _ensure_unique_var_names(io_vars) + seen_var_names = set([v.name for v in io_vars]) + seen_op_names = set() + _deduplicate_block(func, set(func.outputs), + seen_var_names, seen_op_names) diff --git a/coremltools/converters/mil/mil/passes/detect_concat_interleave.py b/coremltools/converters/mil/mil/passes/detect_concat_interleave.py index 0433cef31..47ea9fd16 100644 --- a/coremltools/converters/mil/mil/passes/detect_concat_interleave.py +++ b/coremltools/converters/mil/mil/passes/detect_concat_interleave.py @@ -7,11 +7,12 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb import numpy as np from coremltools.converters.mil.mil.types.symbolic import is_symbolic, any_symbolic -def match_pattern(op): +def _match_pattern(op): if op.outputs[0] in op.enclosing_block.outputs: return None @@ -48,7 +49,8 @@ def match_pattern(op): return op return None -def try_to_transform(concat_op, add_op, block): + +def _try_to_transform(concat_op, add_op, block): all_ops = [concat_op] B, C, H, W = list(concat_op.values[0].shape) n = len(concat_op.values) @@ -117,27 +119,29 @@ def try_to_transform(concat_op, add_op, block): block.remove_ops(all_ops) return True -def fuse_concat_interleave(block): + +def _fuse_concat_interleave(block): fusion_status = False for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = fuse_concat_interleave(b) + block_changed = _fuse_concat_interleave(b) if len(op.blocks) > 0: continue - concat_op = match_pattern(op) + concat_op = _match_pattern(op) if concat_op is not None: with block: - fusion_status = try_to_transform(op, concat_op, block) + fusion_status = _try_to_transform(op, concat_op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status return fusion_status + @register_pass(namespace="common") -def detect_concat_interleave(prog): +class detect_concat_interleave(AbstractGraphPass): """ Detect the pattern "concat-->reshape--->transpose--->reshape", where concat is along the channel axis (axis=-3), and map this pattern to the concat interleave op. @@ -153,7 +157,8 @@ def detect_concat_interleave(prog): Result: %6 = concat(%1.a, %1.b, ..., axis=-3, interleave=True) """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = fuse_concat_interleave(f) \ No newline at end of file + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_concat_interleave(f) diff --git a/coremltools/converters/mil/mil/passes/divide_to_multiply.py b/coremltools/converters/mil/mil/passes/divide_to_multiply.py index 9ce5cd736..ed5a1d22b 100644 --- a/coremltools/converters/mil/mil/passes/divide_to_multiply.py +++ b/coremltools/converters/mil/mil/passes/divide_to_multiply.py @@ -8,14 +8,14 @@ import numpy as np from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import types as _types - -def divide_to_multiply_block(block): +def _divide_to_multiply_block(block): for op in list(block.operations): for b in op.blocks: - divide_to_multiply_block(b) + _divide_to_multiply_block(b) if len(op.blocks) > 0: # This op can't be divide. continue @@ -41,9 +41,10 @@ def divide_to_multiply_block(block): @register_pass(namespace="common") -def divide_to_multiply(prog): +class divide_to_multiply(AbstractGraphPass): """ Convert divide into multiply if divisor is const. """ - for f in prog.functions.values(): - divide_to_multiply_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _divide_to_multiply_block(f) diff --git a/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py b/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py index cf888aa27..bccacdc43 100644 --- a/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py +++ b/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py @@ -7,11 +7,11 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb import numpy as np - -def match_pattern(op): +def _match_pattern(op): if op.outputs[0] in op.enclosing_block.outputs: return None @@ -47,7 +47,7 @@ def _check_shape(arr): return True -def try_to_transform(mul_op, add_op, block): +def _try_to_transform(mul_op, add_op, block): non_const_input_mul = mul_op.x if mul_op.x.val is None else mul_op.y if non_const_input_mul.rank != 4: return False @@ -89,29 +89,28 @@ def try_to_transform(mul_op, add_op, block): return True -def fuse_elementwise_to_batchnorm_block(block): +def _fuse_elementwise_to_batchnorm_block(block): fusion_status = False for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = fuse_elementwise_to_batchnorm_block(b) + block_changed = _fuse_elementwise_to_batchnorm_block(b) if len(op.blocks) > 0: # This op can't be mul continue - add_op = match_pattern(op) + add_op = _match_pattern(op) if add_op is not None: with block: - fusion_status = try_to_transform(op, add_op, block) + fusion_status = _try_to_transform(op, add_op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status return fusion_status - @register_pass(namespace="common") -def fuse_elementwise_to_batchnorm(prog): +class fuse_elementwise_to_batchnorm(AbstractGraphPass): """ Fold mul + add into a batch norm, if the const feeding into the mul/add is of shape (1,C,1,1) or (C,1,1) @@ -142,7 +141,8 @@ def fuse_elementwise_to_batchnorm(prog): ... """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = fuse_elementwise_to_batchnorm_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_elementwise_to_batchnorm_block(f) diff --git a/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py b/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py index 33dd82fda..27ea7bbc7 100644 --- a/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py +++ b/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py @@ -5,6 +5,7 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from .helper import _check_child_op_type, _check_var_scalar_value import numpy as np @@ -105,9 +106,8 @@ def _fuse_gelu_exact_block(block): return fusion_occurred return fusion_occurred - @register_pass(namespace="common") -def fuse_gelu_exact(prog): +class fuse_gelu_exact(AbstractGraphPass): """ Identify the pattern that corresponds to the exact version of gelu, and replace it with a single gelu layer with mode=EXACT @@ -129,7 +129,8 @@ def fuse_gelu_exact(prog): both result in : [...] ----> gelu (mode=EXACT) ---> [...] """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_gelu_exact_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_gelu_exact_block(f) diff --git a/coremltools/converters/mil/mil/passes/gelu_tanh_approximation_fusion.py b/coremltools/converters/mil/mil/passes/gelu_tanh_approximation_fusion.py index f62ab70ae..4ef421b77 100644 --- a/coremltools/converters/mil/mil/passes/gelu_tanh_approximation_fusion.py +++ b/coremltools/converters/mil/mil/passes/gelu_tanh_approximation_fusion.py @@ -7,12 +7,12 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from .helper import _check_child_op_type, _check_var_scalar_value import numpy as np - -def try_to_transform(pow_op, block): +def _try_to_transform(pow_op, block): all_ops = [pow_op] root_var = pow_op.x @@ -125,13 +125,13 @@ def try_to_transform(pow_op, block): return True -def fuse_gelu_tanh_block(block): +def _fuse_gelu_tanh_block(block): fusion_status = False for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = fuse_gelu_tanh_block(b) + block_changed = _fuse_gelu_tanh_block(b) if len(op.blocks) > 0: # This op can't be pow continue @@ -140,15 +140,14 @@ def fuse_gelu_tanh_block(block): if op.op_type == "pow": if _check_var_scalar_value(op.y, 3): with block: - fusion_status = try_to_transform(op, block) + fusion_status = _try_to_transform(op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status return fusion_status - @register_pass(namespace="common") -def fuse_gelu_tanh_approximation(prog): +class fuse_gelu_tanh_approximation(AbstractGraphPass): """ Identify the pattern that corresponds to the tanh approximate version of gelu, and replace it with a single gelu layer with mode=TANH_APPROXIMATION @@ -162,7 +161,8 @@ def fuse_gelu_tanh_approximation(prog): """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = fuse_gelu_tanh_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_gelu_tanh_block(f) diff --git a/coremltools/converters/mil/mil/passes/graph_pass.py b/coremltools/converters/mil/mil/passes/graph_pass.py new file mode 100644 index 000000000..c272a74bf --- /dev/null +++ b/coremltools/converters/mil/mil/passes/graph_pass.py @@ -0,0 +1,15 @@ +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + +class AbstractGraphPass(): + + def __call__(self, prog): + self.apply(prog) + + def apply(self, prog): + raise NotImplementedError( + 'Graph pass transformation not implemented for "{}".'.format(self) + ) + diff --git a/coremltools/converters/mil/mil/passes/helper.py b/coremltools/converters/mil/mil/passes/helper.py index 4134c35d0..ba11cbeb2 100644 --- a/coremltools/converters/mil/mil/passes/helper.py +++ b/coremltools/converters/mil/mil/passes/helper.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2021, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be @@ -7,6 +5,8 @@ import numpy as np +from coremltools.converters.mil.mil import Var + def _check_child_op_type(op, child_op_type): """ :param op: operation @@ -69,4 +69,56 @@ def _check_var_scalar_value(x, val, tol=1e-3): if abs(x_val - val) < tol: return True - return False \ No newline at end of file + return False + +def _are_ops_identical(op1, op2): + ''' + Return True, if all inputs of op1 and op2 are identical. + non-constant inputs must refer to the same object, and constant inputs must have the same value + ''' + + def _are_values_identical(val1, val2): + np_arr1 = np.array(val1) + np_arr2 = np.array(val2) + return np.array_equal(np_arr1, np_arr2) + + def _are_vars_identical(var1, var2): + if var1.val is None and var2.val is None: + if var1 != var2: + return False + elif var1.val is not None and var2.val is not None: + if var1.dtype != var2.dtype: + return False + if not _are_values_identical(var1.val, var2.val): + return False + else: + return False + return True + + if op1 == op2: + return True + if op1.op_type != op2.op_type: + return False + if len(op1.inputs) != len(op2.inputs): + return False + + for key, value1 in op1.inputs.items(): + if key not in op2.inputs: + return False + value2 = op2.inputs[key] + if isinstance(value1, Var) and isinstance(value2, Var): + if not _are_vars_identical(value1, value2): + return False + elif isinstance(value1, (list, tuple)) and isinstance(value2, (list, tuple)): + if len(value1) != len(value2): + return False + else: + for i, v in enumerate(value1): + if not _are_vars_identical(v, value2[i]): + return False + else: + return False + + assert len(op1.blocks) == 0, "this method does not handle ops that have blocks in it" + assert len(op2.blocks) == 0, "this method does not handle ops that have blocks in it" + return True \ No newline at end of file diff --git a/coremltools/converters/mil/mil/passes/image_input_preprocessing.py b/coremltools/converters/mil/mil/passes/image_input_preprocessing.py index fcbd84025..78601c9f1 100644 --- a/coremltools/converters/mil/mil/passes/image_input_preprocessing.py +++ b/coremltools/converters/mil/mil/passes/image_input_preprocessing.py @@ -1,11 +1,15 @@ -from coremltools.converters.mil.mil.passes.pass_registry import register_pass -from coremltools.converters.mil.input_types import ImageType, Shape, EnumeratedShapes +# Copyright (c) 2020, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clausefrom +from coremltools.converters.mil.input_types import ImageType, Shape, EnumeratedShapes +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb - @register_pass(namespace='common') -def image_input_preprocess(prog): +class image_input_preprocess(AbstractGraphPass): """ Plug in transpose for image input that were NHWC to NCHW. @@ -16,13 +20,16 @@ def image_input_preprocess(prog): We do not modify this input, since channel_first is the intended behaviour for feeding images for optimal performance b) channel_first == False - We convert the input into a "channel_first" input, and plug in a + We convert the input into a "channel_first" input, and plug in a transpose for the input to maintain the remaining graph's dimensionality. """ - for f_name, f in prog.functions.items(): - if f_name == 'main': - # We need to make sure main exist and start here. - _image_input_preprocess(prog) + + def apply(self, prog): + for f_name, f in prog.functions.items(): + if f_name == 'main': + # We need to make sure main exist and start here. + _image_input_preprocess(prog) + def _transform_to_channel_first(shape): if isinstance(shape, tuple): diff --git a/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py b/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py index de2e7ad8b..fce0b8002 100644 --- a/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py +++ b/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py @@ -12,10 +12,10 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import Operation, Block, Var, Program from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass DEBUG = False # set to true to plot the block before and after the transformation - def _check_no_output_connection(block: Block, ops: List[Operation]) -> bool: """ Check that none of the op in this pattern is connected to the output @@ -822,7 +822,7 @@ def _fuse_layernorm_or_instancenorm_block(block: Block): @register_pass(namespace="common") -def fuse_layernorm_or_instancenorm(prog: Program): +class fuse_layernorm_or_instancenorm(AbstractGraphPass): """ A graph optimization pass on PyMIL to detect and fuse several different vairants of layer_norm or instance_norm. Pattern 1 is corresponding to @@ -830,26 +830,27 @@ def fuse_layernorm_or_instancenorm(prog: Program): :param prog: PyMIL Program to work on. """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - if DEBUG: - import graphviz - - graphviz.Source( - f.get_dot_string(highlight_debug_op_types=["instance_norm"],) - ).view(filename="/tmp/block_before_fuse_layernorm_or_instancenorm") - logging.debug( - "Block before fuse_layernorm_or_instancenorm transform:\n{}".format(f) - ) - - block_changed = _fuse_layernorm_or_instancenorm_block(f) - - if DEBUG: - graphviz.Source( - f.get_dot_string(highlight_debug_op_types=["instance_norm"],) - ).view(filename="/tmp/block_after_fuse_layernorm_or_instancenorm") - - logging.debug( - "Block after fuse_layernorm_or_instancenorm transform:\n{}".format(f) - ) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + if DEBUG: + import graphviz + + graphviz.Source( + f.get_dot_string(highlight_debug_op_types=["instance_norm"],) + ).view(filename="/tmp/block_before_fuse_layernorm_or_instancenorm") + logging.debug( + "Block before fuse_layernorm_or_instancenorm transform:\n{}".format(f) + ) + + block_changed = _fuse_layernorm_or_instancenorm_block(f) + + if DEBUG: + graphviz.Source( + f.get_dot_string(highlight_debug_op_types=["instance_norm"],) + ).view(filename="/tmp/block_after_fuse_layernorm_or_instancenorm") + + logging.debug( + "Block after fuse_layernorm_or_instancenorm transform:\n{}".format(f) + ) diff --git a/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py b/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py index fcae42b0c..409d4e1c4 100644 --- a/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py +++ b/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2021, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be @@ -7,10 +5,9 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from .helper import _check_var_scalar_value_in_interval, _check_child_op_type -import numpy as np - def _try_to_transform(mul_op, block): @@ -73,9 +70,8 @@ def _fuse_leaky_relu_block(block): return fusion_status return fusion_status - @register_pass(namespace="common") -def fuse_leaky_relu(prog): +class fuse_leaky_relu(AbstractGraphPass): """ Detect the "mul--->max" pattern than can be mapped to leaky relu @@ -107,7 +103,8 @@ def fuse_leaky_relu(prog): input --------> leaky_relu ---------> output """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_leaky_relu_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_leaky_relu_block(f) diff --git a/coremltools/converters/mil/mil/passes/linear_bias_fusion.py b/coremltools/converters/mil/mil/passes/linear_bias_fusion.py index 40503aa17..d356b025d 100644 --- a/coremltools/converters/mil/mil/passes/linear_bias_fusion.py +++ b/coremltools/converters/mil/mil/passes/linear_bias_fusion.py @@ -5,12 +5,11 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb import numpy as np - def _try_to_transform(linear_op, add_or_sub_op, block): if add_or_sub_op.x.val is None and add_or_sub_op.y.val is None: @@ -98,9 +97,8 @@ def _match_pattern(op): return fusion_occurred return fusion_occurred - @register_pass(namespace="common") -def fuse_linear_bias(prog): +class fuse_linear_bias(AbstractGraphPass): """ Convert linear + add/sub to a single linear by updating the weight and bias of the linear layer. @@ -130,7 +128,8 @@ def fuse_linear_bias(prog): prog: Program """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_linear_bias_block(f) \ No newline at end of file + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_linear_bias_block(f) diff --git a/coremltools/converters/mil/mil/passes/loop_invariant_elimination.py b/coremltools/converters/mil/mil/passes/loop_invariant_elimination.py index 2c25ef315..36c2c8eb1 100644 --- a/coremltools/converters/mil/mil/passes/loop_invariant_elimination.py +++ b/coremltools/converters/mil/mil/passes/loop_invariant_elimination.py @@ -1,18 +1,13 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -import numpy as np - from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass - -def detect_loop_invariants(while_op): +def _detect_loop_invariants(while_op): block = while_op.blocks[1] # body block loop_invariant_ids = [] # list of index in op.loop_vars, block.inputs for i, vx_in in enumerate(block.inputs): @@ -30,7 +25,7 @@ def detect_loop_invariants(while_op): return loop_invariant_ids -def loop_invariant_elimination_block(block): +def _loop_invariant_elimination_block(block): # Phase 1: Find vars needed to be renamed. # # while_loop outputs need to be renamed if the output will be eliminated @@ -42,12 +37,12 @@ def loop_invariant_elimination_block(block): output_rename = [] for op in list(block.operations): for b in op.blocks: - loop_invariant_elimination_block(b) + _loop_invariant_elimination_block(b) if op.op_type != "while_loop": continue - loop_invariant_ids = detect_loop_invariants(op) + loop_invariant_ids = _detect_loop_invariants(op) for i in loop_invariant_ids: output_rename.append((op.loop_vars[i], op.outputs[i], op)) if len(loop_invariant_ids) > 0: @@ -73,7 +68,7 @@ def loop_invariant_elimination_block(block): for op in list(block.operations): if op.op_type != "while_loop": continue - loop_invariant_ids = detect_loop_invariants(op) + loop_invariant_ids = _detect_loop_invariants(op) loop_variant_vars = [] @@ -124,9 +119,8 @@ def loop_invariant_elimination_block(block): # check healthy state op.enclosing_block.validate() - @register_pass(namespace="common") -def loop_invariant_elimination(prog): +class loop_invariant_elimination(AbstractGraphPass): """ prog: Program @@ -170,5 +164,6 @@ def loop_invariant_elimination(prog): # instead of 2 outputs. We also preserve the return var names with # identity. """ - for f in prog.functions.values(): - loop_invariant_elimination_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _loop_invariant_elimination_block(f) diff --git a/coremltools/converters/mil/mil/passes/matmul_weight_bias_fusion.py b/coremltools/converters/mil/mil/passes/matmul_weight_bias_fusion.py index 23feb5c86..21ffaf0de 100644 --- a/coremltools/converters/mil/mil/passes/matmul_weight_bias_fusion.py +++ b/coremltools/converters/mil/mil/passes/matmul_weight_bias_fusion.py @@ -8,12 +8,13 @@ import numpy as np from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb child_op_types = ["add", "sub"] +def _match_pattern(op): -def match_pattern(op): if op.op_type == "matmul": # find add child_ops = op.outputs[0].child_ops @@ -23,8 +24,7 @@ def match_pattern(op): return add_op_candidate return None - -def transpose(v, before_op): +def _transpose(v, before_op): """ Transpose the last 2 dims. v: Var (must be a tensor) @@ -33,8 +33,7 @@ def transpose(v, before_op): perm[-2], perm[-1] = perm[-1], perm[-2] return mb.transpose(x=v, perm=perm, before_op=before_op) - -def try_to_transform(matmul_op, add_op, block): +def _try_to_transform(matmul_op, add_op, block): if matmul_op.x.val is None and matmul_op.y.val is None: # This is a dynamic matmul. return False @@ -84,23 +83,23 @@ def try_to_transform(matmul_op, add_op, block): # If transpose_x == transpose_weight == False: # w*x = (x^T w^T)^T = linear(x^T, w)^T x_transposed = ( - transpose(linear_x, before_op=matmul_op) if not transpose_x else linear_x + _transpose(linear_x, before_op=matmul_op) if not transpose_x else linear_x ) w_no_transpose = ( - weight if not transpose_weight else transpose(weight, before_op=matmul_op) + weight if not transpose_weight else _transpose(weight, before_op=matmul_op) ) x = mb.linear( x=x_transposed, weight=w_no_transpose, bias=bias, before_op=matmul_op ) - x = transpose(x, before_op=matmul_op, name=out_name) + x = _transpose(x, before_op=matmul_op, name=out_name) else: # If transpose_x == transpose_weight == False # x*w = x*(w^T)^T = linear(x, w^T) x_no_transpose = ( - transpose(linear_x, before_op=matmul_op) if transpose_x else linear_x + _transpose(linear_x, before_op=matmul_op) if transpose_x else linear_x ) w_transposed = ( - weight if transpose_weight else transpose(weight, before_op=matmul_op) + weight if transpose_weight else _transpose(weight, before_op=matmul_op) ) x = mb.linear( x=x_no_transpose, @@ -118,29 +117,28 @@ def try_to_transform(matmul_op, add_op, block): return True -def fuse_matmul_weight_bias_block(block): +def _fuse_matmul_weight_bias_block(block): fusion_status = False for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = fuse_matmul_weight_bias_block(b) + block_changed = _fuse_matmul_weight_bias_block(b) if len(op.blocks) > 0: # This op can't be matmul continue - add_op = match_pattern(op) + add_op = _match_pattern(op) if add_op is not None: with block: - fusion_status = try_to_transform(op, add_op, block) + fusion_status = _try_to_transform(op, add_op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status return fusion_status - @register_pass(namespace="common") -def fuse_matmul_weight_bias(prog): +class fuse_matmul_weight_bias(AbstractGraphPass): """ Convert matmul + add/sub to linear whenever possible. @@ -158,7 +156,8 @@ def fuse_matmul_weight_bias(prog): prog: Program """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = fuse_matmul_weight_bias_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_matmul_weight_bias_block(f) diff --git a/coremltools/converters/mil/mil/passes/merge_consecutive_paddings.py b/coremltools/converters/mil/mil/passes/merge_consecutive_paddings.py index cc8b9b205..2314ac7a2 100644 --- a/coremltools/converters/mil/mil/passes/merge_consecutive_paddings.py +++ b/coremltools/converters/mil/mil/passes/merge_consecutive_paddings.py @@ -5,10 +5,10 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from .helper import _check_child_op_type import numpy as np - def _match_pattern(block, padding_op): if padding_op.op_type != "pad": @@ -22,7 +22,7 @@ def _match_pattern(block, padding_op): if padding_op.inputs["mode"].val != child_padding_op.inputs["mode"].val: return False - # Ensure the paddings have the same length by prepending zeros to the shorter one + # Ensure the paddings have the same length by prepending zeros to the shorter one first_pad = padding_op.inputs["pad"].val child_pad = child_padding_op.inputs["pad"].val if len(first_pad) > len(child_pad): @@ -30,12 +30,12 @@ def _match_pattern(block, padding_op): elif len(child_pad) > len(first_pad): first_pad = np.insert(first_pad, 0, [0] * (len(child_pad) - len(first_pad))) final_pad = child_pad + first_pad - + if padding_op.inputs["mode"].val == "constant": # if the padding is constant, then the values need to be equal if padding_op.inputs["constant_val"].val != child_padding_op.inputs["constant_val"].val: return False - else: + else: # if the padding is not constant, then we can't merge if both pads affected the same side of the image if any(i != 0 and j != 0 for (i,j) in zip(first_pad, child_pad)): return False @@ -48,7 +48,7 @@ def _replace_ops(block, padding_op, child_padding_op, final_pad): with block: mode = padding_op.inputs["mode"].val - x = mb.pad(x=padding_op.inputs["x"], pad=final_pad, mode=mode, constant_val=padding_op.inputs["constant_val"].val, + x = mb.pad(x=padding_op.inputs["x"], pad=final_pad, mode=mode, constant_val=padding_op.inputs["constant_val"].val, before_op=padding_op) padding_op.enclosing_block.replace_uses_of_var_after_op( anchor_op=padding_op, old_var=child_padding_op.outputs[0], new_var=x @@ -58,27 +58,29 @@ def _replace_ops(block, padding_op, child_padding_op, final_pad): return True + def _merge_padding_block(block): for op in list(block.operations): result = _match_pattern(block, op) if result: return True - return False + return False @register_pass(namespace="common") -def merge_consecutive_paddings(prog): +class merge_consecutive_paddings(AbstractGraphPass): """ - Identify two consecutive 'pad' layers which could be merged into a single 'pad' layer. This is possible when + Identify two consecutive 'pad' layers which could be merged into a single 'pad' layer. This is possible when - the paddings are "constant" and have the same "constant_val" - - OR, the paddings act along different axes. + - OR, the paddings act along different axes. - Input graph: + Input graph: input(1, 2, 6, 8) ------> pad([1, 1], mode='reflect) -----> pad([1, 1, 0, 0], mode='reflect') ---> out(1, 2, 8, 10) - Output graph: + Output graph: input(1, 2, 6, 8) ------> pad([1, 1, 1, 1], mode='reflect) ---> out(1, 2, 8, 10) """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _merge_padding_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _merge_padding_block(f) diff --git a/coremltools/converters/mil/mil/passes/name_sanitization_utils.py b/coremltools/converters/mil/mil/passes/name_sanitization_utils.py index fb4c1d17a..12498feb6 100644 --- a/coremltools/converters/mil/mil/passes/name_sanitization_utils.py +++ b/coremltools/converters/mil/mil/passes/name_sanitization_utils.py @@ -1,15 +1,15 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2021, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from collections import OrderedDict -from coremltools.converters.mil.mil import Function import re import warnings +from coremltools.converters.mil.mil import Function + + class NameSanitizer(object): def __init__(self, prefix=None): diff --git a/coremltools/converters/mil/mil/passes/noop_elimination.py b/coremltools/converters/mil/mil/passes/noop_elimination.py index eb8a76ba0..7637a0af1 100644 --- a/coremltools/converters/mil/mil/passes/noop_elimination.py +++ b/coremltools/converters/mil/mil/passes/noop_elimination.py @@ -1,9 +1,12 @@ -# -*- coding: utf-8 -*- +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from coremltools.converters.mil.mil.passes.pass_registry import register_pass -from coremltools.converters.mil.mil import Builder as mb import numpy as np +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass def _remove_elementwise_binary(op, block, x, y): # We remove the ops that has op.x == x or op.y == y @@ -30,7 +33,6 @@ def _remove_elementwise_binary(op, block, x, y): return True - def remove_elementwise(op, block): if op.op_type in {"add"}: @@ -44,7 +46,6 @@ def remove_elementwise(op, block): else: return False - def remove_same_shape(op, block): input_shape = op.x.sym_type output_shape = op.outputs[0].sym_type @@ -63,7 +64,6 @@ def remove_same_shape(op, block): block.remove_ops([op]) return True - def remove_linear(op, block): if op.alpha.val != 1 or op.beta.val != 0: return False @@ -95,7 +95,6 @@ def remove_transpose(op, block): # Remove all the ops at once block.remove_ops([op]) return True - _SUPPORTED_OPS = { "add", "mul", @@ -116,6 +115,7 @@ def remove_transpose(op, block): "crop", "linear_activation" } + op_to_removal_fn = { "add": remove_elementwise, "mul": remove_elementwise, @@ -137,8 +137,7 @@ def remove_transpose(op, block): "linear_activation": remove_linear, } - -def match_pattern(op): +def _match_pattern(op): # abort if op output is a block output if op.outputs[0] in op.enclosing_block.outputs: return None @@ -151,17 +150,16 @@ def match_pattern(op): return None - -def noop_elimination_block(block): +def _noop_elimination_block(block): for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = noop_elimination_block(b) + block_changed = _noop_elimination_block(b) if len(op.blocks) > 0: continue - remove_fn = match_pattern(op) + remove_fn = _match_pattern(op) if remove_fn is not None: with block: status = remove_fn(op, block) @@ -172,7 +170,7 @@ def noop_elimination_block(block): @register_pass(namespace="common") -def noop_elimination(prog): +class noop_elimination(AbstractGraphPass): """ We remove ops that has no effect. @@ -189,7 +187,9 @@ def noop_elimination(prog): ... """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = noop_elimination_block(f) + + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _noop_elimination_block(f) diff --git a/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.py b/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.py index 53053fda9..bead0e468 100644 --- a/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.py +++ b/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.py @@ -1,18 +1,15 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -from coremltools.converters.mil.mil.passes.pass_registry import register_pass -from coremltools.converters.mil.mil import Builder as mb from .helper import _check_child_op_type, _check_var_scalar_value -import numpy as np +from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass -def try_to_transform(onehot_op, block): +def _try_to_transform(onehot_op, block): root_var = onehot_op.indices # check that the output of the onehot op is not a block output @@ -64,13 +61,13 @@ def try_to_transform(onehot_op, block): return True -def fuse_onehot_matmul_to_gather_block(block): +def _fuse_onehot_matmul_to_gather_block(block): fusion_status = False for i, op in enumerate(list(block.operations)): for b in op.blocks: block_changed = True while block_changed: - block_changed = fuse_onehot_matmul_to_gather_block(b) + block_changed = _fuse_onehot_matmul_to_gather_block(b) if len(op.blocks) > 0: # This op can't be pow continue @@ -78,7 +75,7 @@ def fuse_onehot_matmul_to_gather_block(block): # start pattern match if one_hot op is encountered if op.op_type == "one_hot": with block: - fusion_status = try_to_transform(op, block) + fusion_status = _try_to_transform(op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status @@ -86,7 +83,7 @@ def fuse_onehot_matmul_to_gather_block(block): @register_pass(namespace="common") -def fuse_onehot_matmul_to_gather(prog): +class fuse_onehot_matmul_to_gather(AbstractGraphPass): """ Detect if onehot (axis=-1, on_value=1, off_value=0) is followed by a matmul op (no bias), then they can be replaced by a gather op. @@ -100,7 +97,8 @@ def fuse_onehot_matmul_to_gather(prog): %4 = gather(%3, %2, axis=0) """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = fuse_onehot_matmul_to_gather_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_onehot_matmul_to_gather_block(f) diff --git a/coremltools/converters/mil/mil/passes/pad_conv_connect.py b/coremltools/converters/mil/mil/passes/pad_conv_connect.py index dc3f9128f..83d5b731f 100644 --- a/coremltools/converters/mil/mil/passes/pad_conv_connect.py +++ b/coremltools/converters/mil/mil/passes/pad_conv_connect.py @@ -7,11 +7,12 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb import numpy as np import copy -def match_pattern(op): +def _match_pattern(op): ret = set([]) child_ops = op.outputs[0].child_ops @@ -27,7 +28,7 @@ def match_pattern(op): return ret if len(ret) != 0 else None -def try_to_transform(pad_op, transpose_ops, block): +def _try_to_transform(pad_op, transpose_ops, block): def _compute_new_pad_values(transpose_op): if pad_op.inputs["pad"].val is None: @@ -82,29 +83,29 @@ def _compute_new_pad_values(transpose_op): return True -def pad_conv_connect_block(block): + +def _pad_conv_connect_block(block): fusion_status = False for op in list(block.operations): for b in op.blocks: block_changed = True while block_changed: - block_changed = pad_conv_connect_block(b) + block_changed = _pad_conv_connect_block(b) if op.op_type != "pad": continue - transpose_ops = match_pattern(op) + transpose_ops = _match_pattern(op) if transpose_ops is not None: with block: - fusion_status = try_to_transform(op, transpose_ops, block) + fusion_status = _try_to_transform(op, transpose_ops, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status return fusion_status - @register_pass(namespace="common") -def pad_conv_connect(prog): +class pad_conv_connect(AbstractGraphPass): """ When we observe pad -> transpose -> conv, we move the pad to be next to conv. This allows us to meld pad + conv if possible. @@ -122,7 +123,8 @@ def pad_conv_connect(prog): ... """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = pad_conv_connect_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _pad_conv_connect_block(f) diff --git a/coremltools/converters/mil/mil/passes/pass_registry.py b/coremltools/converters/mil/mil/passes/pass_registry.py index ec7b100a7..2dc0c5e4d 100644 --- a/coremltools/converters/mil/mil/passes/pass_registry.py +++ b/coremltools/converters/mil/mil/passes/pass_registry.py @@ -8,8 +8,8 @@ class PassRegistry: def __init__(self): - # str -> func (func takes Program as input and - # modifies in-place) + # str -> an AbstractGraphPass instance, which has an 'apply' method that takes + # program as input and modifies it in place self.passes = {} def __getitem__(self, pass_id): @@ -23,27 +23,29 @@ def __getitem__(self, pass_id): def __contains__(self, pass_id): return pass_id in self.passes - def add(self, namespace, pass_func): - func_name = pass_func.__name__ - pass_id = namespace + "::" + func_name + def add(self, namespace, pass_cls, override, name): + cls_name = pass_cls.__name__ if name is None else name + pass_id = namespace + "::" + cls_name logging.debug("Registering pass {}".format(pass_id)) - if pass_id in self.passes: + if pass_id in self.passes and not override: msg = "Pass {} already registered." raise KeyError(msg.format(pass_id)) - self.passes[pass_id] = pass_func + self.passes[pass_id] = pass_cls() PASS_REGISTRY = PassRegistry() -def register_pass(namespace): +def register_pass(namespace, override=False, name=None): """ namespaces like {'common', 'nn_backend', , } + override: indicate the graph pass can override an existing pass with the same name + name: name of the graph pass. Default to class name if not provided """ - def func_wrapper(pass_func): - PASS_REGISTRY.add(namespace, pass_func) - return pass_func + def class_wrapper(pass_cls): + PASS_REGISTRY.add(namespace, pass_cls, override, name) + return pass_cls - return func_wrapper + return class_wrapper diff --git a/coremltools/converters/mil/mil/passes/quantization_passes.py b/coremltools/converters/mil/mil/passes/quantization_passes.py index 18827ae10..c38de18b9 100644 --- a/coremltools/converters/mil/mil/passes/quantization_passes.py +++ b/coremltools/converters/mil/mil/passes/quantization_passes.py @@ -9,6 +9,7 @@ from coremltools.converters.mil.mil.operation import Operation from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.types import builtin_to_string, is_tensor +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from enum import Enum @@ -28,7 +29,7 @@ def _close_to_zero(val, np_type): return np.isclose(val, 0, atol=type_min[np_type], rtol=type_eps[np_type]) -class AbstractQuantizationPass(object): +class AbstractQuantizationPass(AbstractGraphPass): """ Base class for Post-Training Quantization transforms. @@ -44,8 +45,8 @@ def __init__(self, op_selector=None): "accepts a MIL operation object and returns a boolean value." ) raise TypeError(msg) - else: - self.op_selector = op_selector + + self.op_selector = op_selector def apply(self, prog): """ diff --git a/coremltools/converters/mil/mil/passes/rank0_expand_dims_swap.py b/coremltools/converters/mil/mil/passes/rank0_expand_dims_swap.py index 0bf1526af..76e65fef1 100644 --- a/coremltools/converters/mil/mil/passes/rank0_expand_dims_swap.py +++ b/coremltools/converters/mil/mil/passes/rank0_expand_dims_swap.py @@ -4,6 +4,7 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.helper import _check_child_op_type @@ -79,6 +80,7 @@ def _try_to_transform(op, block): block.remove_ops(ops_to_remove) return True + def _rank0_expand_dims_swap(block): fusion_occurred = False for op in list(block.operations): @@ -98,9 +100,8 @@ def _rank0_expand_dims_swap(block): return fusion_occurred return fusion_occurred - @register_pass(namespace="common") -def rank0_expand_dims_swap(prog): +class rank0_expand_dims_swap(AbstractGraphPass): """ Identify the pattern that a rank-0 binary elementwise operation followed by an expand_dims op. In the MIL backend, the output of the elementwise op becomes rank 1. Hence an expand_dims op @@ -131,7 +132,9 @@ def rank0_expand_dims_swap(prog): | [scalar const] """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _rank0_expand_dims_swap(f) \ No newline at end of file + + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _rank0_expand_dims_swap(f) diff --git a/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py b/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py index 06a75c596..a7b8d4a76 100644 --- a/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py +++ b/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py @@ -1,17 +1,13 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2021, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.types.symbolic import is_symbolic from .helper import _check_var_scalar_value, _check_child_op_type -import numpy as np - def _try_to_transform(reduce_sum_op, block): @@ -22,7 +18,9 @@ def _try_to_transform(reduce_sum_op, block): input_shape = reduce_sum_op.x.shape if input_shape is None: return False - axes = reduce_sum_op.axes.val + axes = None + if reduce_sum_op.axes is not None: + axes = reduce_sum_op.axes.val if axes is None: return False count = 1 @@ -90,9 +88,8 @@ def _fuse_reduce_mean_block(block): return fusion_status return fusion_status - @register_pass(namespace="common") -def fuse_reduce_mean(prog): +class fuse_reduce_mean(AbstractGraphPass): """ Detect the "reduce_sum--->mul/real_div" pattern than can be mapped to reduce_mean. That is, the operation "reduce_sum/count == reduce_mean" @@ -108,8 +105,9 @@ def fuse_reduce_mean(prog): input --------> reduce_mean ---------> output """ - for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _fuse_reduce_mean_block(f) + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _fuse_reduce_mean_block(f) diff --git a/coremltools/converters/mil/mil/passes/reduce_transposes.py b/coremltools/converters/mil/mil/passes/reduce_transposes.py index 1b0cfa2fd..747b3aa09 100644 --- a/coremltools/converters/mil/mil/passes/reduce_transposes.py +++ b/coremltools/converters/mil/mil/passes/reduce_transposes.py @@ -7,6 +7,7 @@ from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.types.symbolic import any_symbolic from coremltools.converters.mil.mil.var import Var @@ -1227,7 +1228,7 @@ def apply_transform(self): op.type_value_inference(overwrite_output=True) -def reduce_transposes_block(block): +def _reduce_transposes_block(block): """ Only apply the optimization if the block is flat, i.e, it does not contain any op which contains a sub-block. @@ -1247,6 +1248,8 @@ def reduce_transposes_block(block): @register_pass(namespace="common") -def reduce_transposes(prog): - for f in prog.functions.values(): - reduce_transposes_block(f) +class reduce_transposes(AbstractGraphPass): + + def apply(self, prog): + for f in prog.functions.values(): + _reduce_transposes_block(f) diff --git a/coremltools/converters/mil/mil/passes/remove_redundant_ops.py b/coremltools/converters/mil/mil/passes/remove_redundant_ops.py new file mode 100644 index 000000000..f0b0eca6f --- /dev/null +++ b/coremltools/converters/mil/mil/passes/remove_redundant_ops.py @@ -0,0 +1,192 @@ +# Copyright (c) 2021, Apple Inc. All rights reserved. +# +# Use of this source code is governed by a BSD-3-clause license that can be +# found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + +from collections import OrderedDict +import numpy as np + +from .helper import _are_ops_identical +from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.pass_registry import register_pass + + +def _is_op_eligible_to_be_removed(op): + if len(op.blocks) != 0 or op.op_type.startswith("random"): + return False + else: + return True + + +def _get_candidate_ops_list(prospective_ops_list): + od = OrderedDict() + enclosing_block = [op.enclosing_block for op in prospective_ops_list] + if len(set(enclosing_block)) > 1: # all candidate ops must belong to the same block + return [] + for op in prospective_ops_list: + if _is_op_eligible_to_be_removed(op): + od[op] = enclosing_block[0].operations.index(op) + # sort the ops according to their index of appearing in block.operations, which is topologically sorted + return [x[0] for x in sorted(od.items(), key=lambda t: t[1])] + + +def _get_candidate_ops_lists_from_var(var): + ''' + Return a list of lists. + Each element is a list of a subset of the child ops of var, which satisifies the following conditions: + - they are of the same op_type + - ops are not repeated in it. The .child_ops property of a var may sometimes contain an op repeated more than once + - the ops are ordered based on the order in which they appear in the block.operations list (which is topologically sorted), + with ops appearing earlier in that list appearing first here. + ''' + candidate_ops_lists = [] + + op_types_to_ops = OrderedDict() + for op in var.child_ops: + if op.op_type in op_types_to_ops: + op_types_to_ops[op.op_type].append(op) + else: + op_types_to_ops[op.op_type] = [op] + + for v in op_types_to_ops.values(): + if len(v) > 1: + candidate_ops_list = _get_candidate_ops_list(v) + if len(candidate_ops_list) > 1: + candidate_ops_lists.append(candidate_ops_list) + + return candidate_ops_lists + + +def _try_to_remove_ops(candidate_ops_list): + # candidate_ops_list contains ops in topological order. + # All the ops in candidate_ops_list will be compared to the first op, and removed if identical to it. + # Removing ops later in the topological order is much easier, as their output vars + # can simply be replaced by the output var of the first_op, this doesn't require + # changing any op order in the block. + if len(candidate_ops_list) < 2: + return False + first_op = candidate_ops_list[0] + block = first_op.enclosing_block + + # currently, we only consider the cases when the op has 1 output. + # The replace var logic below only handles the single output case. + if len(first_op.outputs) > 1: + return False + + ops_to_remove = [] + for op in candidate_ops_list[1:]: + if op.outputs[0] not in block.outputs: # to make sure we don't remove an output op + if _are_ops_identical(first_op, op): + ops_to_remove.append(op) + + if len(ops_to_remove) == 0: + return False + + # remove uses of output vars of the ops to be removed. + # This can be safely done, since all the ops in ops_to_remove + # appear after first_op, hence first_op.outputs[0] variable is in + # scope before the op's output var + for op in ops_to_remove: + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=first_op.outputs[0] + ) + block.remove_ops(ops_to_remove) + return True + + + +def _try_to_transform(parent_var): + """ + scan the children ops to parent_var, to find and remove indentical ops, if any. + Returns True, if succesful in finding such redundant ops. + """ + candidate_ops_lists = _get_candidate_ops_lists_from_var(parent_var) + block_changed = False + for ops_list in candidate_ops_lists: + if _try_to_remove_ops(ops_list): + block_changed = True + return block_changed + + +def _remove_redundant_ops_in_block(block): + + if isinstance(block.inputs, dict): + block_input_var_list = list(block.inputs.values()) + elif isinstance(block.inputs, (list, tuple)): + block_input_var_list = block.inputs + else: + raise ValueError("Unrecognized type of block.inputs, its neither a list nor dict.") + + # iterate over the block inputs + for input_var in block_input_var_list: + if len(input_var.child_ops) > 1: + with block: + _try_to_transform(input_var) + + # iterate over the ops in the block + graph_updated = False + for op in block.operations: + if op.op_type == "const": + continue + + for b in op.blocks: + block_changed = True + while block_changed: + block_changed = _remove_redundant_ops_in_block(b) + + if len(op.outputs) > 0 and len(op.outputs[0].child_ops) > 1: + with block: + # currently, we only check the first output of the op + # this can be extended, if required, to check for other outputs. + graph_updated = _try_to_transform(op.outputs[0]) + # has to break as the downstream iterator is affected. + if graph_updated: + return graph_updated + return graph_updated + + +@register_pass(namespace="common") +class remove_redundant_ops(AbstractGraphPass): + """ + If there are multiple ops with "identical" inputs, then they are redundant and all but one of them can be removed. + This pass checks and removes such ops. + + Since all inputs to ops in MIL are named, two ops with same op_types, can be compared by comparing their + correspondingly named inputs. Inputs are treated as identical if: + - if the input is a constant var, then its value should have the same dtype and numerical value + - if the input is a non constant var, then it should be the same var object + + This pass iterates over the ops, takes its first output var, and then builds a candidate op list from the child + ops of this var. + These candidate ops list contains ops, of the same op_type, and arranged in the topological order. + From each of these candidate ops list, the 2nd, 3rd,....so on, ops are pairwise compared with the first op, + and if identical to it, they are removed. + e.g.: + + Input: + %0 = op0(...) + %1 = op1(...) + %2 = const(val=4.5) + %3 = const(val=4.5) + %4 = op2(%1, %0, %2) + %5 = op3(%1, %0, %3) + + op3 will be removed and all uses of %5 will be replaced by %4 + + Output: + %0 = op0(...) + %1 = op1(...) + %2 = const(val=4.5) + %3 = const(val=4.5) # this will get removed later by dead code elimination pass + %4 = op2(%1, %0, %2) + + For more examples, see "TestRemoveRedundantOpsPass" + """ + def apply(self, prog): + for f in prog.functions.values(): + block_changed = True + while block_changed: + block_changed = _remove_redundant_ops_in_block(f) + + diff --git a/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py b/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py index 3aca15278..16013648c 100644 --- a/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py +++ b/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py @@ -14,14 +14,14 @@ ) from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass import logging - -def remove_symbolic_reshape_block(block): +def _remove_symbolic_reshape_block(block): num_changes = 0 for op in list(block.operations): for b in op.blocks: - num_changes += remove_symbolic_reshape_block(b) + num_changes += _remove_symbolic_reshape_block(b) if op.op_type != "reshape": continue if op.shape.val is not None: @@ -63,7 +63,7 @@ def remove_symbolic_reshape_block(block): @register_pass(namespace="common") -def remove_symbolic_reshape(prog): +class remove_symbolic_reshape(AbstractGraphPass): """ Convert symbolic shape in `reshape` to integers. @@ -96,7 +96,8 @@ def remove_symbolic_reshape(prog): prog: Program """ - for f in prog.functions.values(): - num_changes = remove_symbolic_reshape_block(f) - msg = "remove_symbolic_reshape: changed {} reshapes." - logging.info(msg.format(num_changes)) + def apply(self, prog): + for f in prog.functions.values(): + num_changes = _remove_symbolic_reshape_block(f) + msg = "remove_symbolic_reshape: changed {} reshapes." + logging.info(msg.format(num_changes)) diff --git a/coremltools/converters/mil/mil/passes/replace_stack_reshape.py b/coremltools/converters/mil/mil/passes/replace_stack_reshape.py index 6f58ac75f..59722d8a7 100644 --- a/coremltools/converters/mil/mil/passes/replace_stack_reshape.py +++ b/coremltools/converters/mil/mil/passes/replace_stack_reshape.py @@ -5,6 +5,7 @@ from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass def _match_operation(stack_op): @@ -35,9 +36,9 @@ def _match_operation(stack_op): return None, None # Compare the input to stack to the output from reshape - # These shapes should differ in either the stack_axis_val place (by a factor of 2), + # These shapes should differ in either the stack_axis_val place (by a factor of 2), # or in the stack_axis_val-1 place by the same factor - input_shape = list(stack_op.inputs["values"][0].shape) + input_shape = list(stack_op.inputs["values"][0].shape) concat_axis = [idx for idx, (x, y) in enumerate(zip(input_shape, reshape_op.outputs[0].shape)) if x != y] if len(concat_axis) != 1: return None, None @@ -89,21 +90,21 @@ def _replace_stack_reshape_block(block): @register_pass(namespace="common") -def replace_stack_reshape(prog): +class replace_stack_reshape(AbstractGraphPass): """ - A stack followed by a reshape layer can be replaced by a concat if the reshape + A stack followed by a reshape layer can be replaced by a concat if the reshape simply removes the new axis and doubles the size of one of the axes next to it. - If the new axis is reshaped to the "right" (i.e. the axis just after it is + If the new axis is reshaped to the "right" (i.e. the axis just after it is doubled), then we can use a concat. If it is reshaped to the "left" (the axis just before it is doubled), then the concat needs to set the "interleaved" flag. - Examples: + Examples: Given: %1 = tensor(1, 5, 3, 4) %2 = tensor(1, 5, 3, 4) %3 = stack((%1,%2), axis=2) # shape = (1, 5, 2, 3, 4) - %4 = reshape(%3, shape=[1, 10, 3, 4]) + %4 = reshape(%3, shape=[1, 10, 3, 4]) Result: %1 = tensor(1, 5, 3, 4) @@ -114,12 +115,13 @@ def replace_stack_reshape(prog): %1 = tensor(1, 5, 3, 4) %2 = tensor(1, 5, 3, 4) %3 = stack((%1, %2), axis=1) # shape = (1, 2, 5, 3, 4) - %4 = reshape(%3, shape=[1, 10, 3, 4]) + %4 = reshape(%3, shape=[1, 10, 3, 4]) Result: %1 = tensor(1, 5, 3, 4) %2 = tensor(1, 5, 3, 4) %4 = concat((%1, %2), axis = 1) # shape = (1, 10, 3, 4) """ - for f in prog.functions.values(): - _replace_stack_reshape_block(f) + def apply(self, prog): + for f in prog.functions.values(): + _replace_stack_reshape_block(f) diff --git a/coremltools/converters/mil/mil/passes/sanitize_input_output_names.py b/coremltools/converters/mil/mil/passes/sanitize_input_output_names.py index 16b4088c7..0c9b343f5 100644 --- a/coremltools/converters/mil/mil/passes/sanitize_input_output_names.py +++ b/coremltools/converters/mil/mil/passes/sanitize_input_output_names.py @@ -6,22 +6,23 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from .name_sanitization_utils import NameSanitizer, sanitize_block @register_pass(namespace="common") -def sanitize_input_output_names(prog): +class sanitize_input_output_names(AbstractGraphPass): """ Sanitize the names of model input and output vars to make sure that they are of the format as described in the NameSanitizer class, i.e. of the format [a-zA-Z_][a-zA-Z0-9_]* """ + def apply(self, prog): + sanitizer_vars = NameSanitizer(prefix="var_") + sanitizer_ops = NameSanitizer(prefix="op_") - sanitizer_vars = NameSanitizer(prefix="var_") - sanitizer_ops = NameSanitizer(prefix="op_") - - # sanitize the input/output of the main block - sanitize_block(prog.functions["main"], - sanitizer_vars, - sanitizer_ops, - prog.main_input_types, - sanitize_model_inputs_outputs_only=True) \ No newline at end of file + # sanitize the input/output of the main block + sanitize_block(prog.functions["main"], + sanitizer_vars, + sanitizer_ops, + prog.main_input_types, + sanitize_model_inputs_outputs_only=True) diff --git a/coremltools/converters/mil/mil/passes/test_cast_optimization.py b/coremltools/converters/mil/mil/passes/test_cast_optimization.py index 626ebe405..439827fe1 100644 --- a/coremltools/converters/mil/mil/passes/test_cast_optimization.py +++ b/coremltools/converters/mil/mil/passes/test_cast_optimization.py @@ -3,6 +3,9 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import numpy as np +import unittest + from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.testing_utils import ( assert_op_count_match, @@ -10,10 +13,7 @@ get_op_types_in_program, apply_pass_and_basic_check, ) -import unittest -import pytest -import numpy as np np.random.seed(1984) diff --git a/coremltools/converters/mil/mil/passes/test_fp16_compute_precision.py b/coremltools/converters/mil/mil/passes/test_fp16_compute_precision.py index add77b35b..9f92b4e34 100644 --- a/coremltools/converters/mil/mil/passes/test_fp16_compute_precision.py +++ b/coremltools/converters/mil/mil/passes/test_fp16_compute_precision.py @@ -94,15 +94,15 @@ def prog(x): """ Input graph: - input1 ----->| + input1 ----->| concat -----> out input2 ----->| Output graph: - input1 -----> cast(dtype="fp16") ----->| + input1 -----> cast(dtype="fp16") ----->| concat -----> cast(dtype="fp32") ---> out - input2 -----> cast(dtype="fp16") ----->| - + input2 -----> cast(dtype="fp16") ----->| + """ def test_multiple_inputs_to_single_operation(self): @@ -149,15 +149,15 @@ def prog(x, y): """ Input graph: |-----> output_1 - input -----> split + input -----> split |-----> output_2 - + Output graph: - + |-----> cast(dtype="fp32") ---> output_1 input -----> cast(dtype="fp16") -----> split |-----> cast(dtype="fp32") ---> output_2 - + """ @@ -201,7 +201,7 @@ def prog(x): """ Input graph: - + |----> square ---> output_1 input| |----> relu ---> output_2 diff --git a/coremltools/converters/mil/mil/passes/test_noop_elimination.py b/coremltools/converters/mil/mil/passes/test_noop_elimination.py index 748540b92..5ef4fd5cb 100644 --- a/coremltools/converters/mil/mil/passes/test_noop_elimination.py +++ b/coremltools/converters/mil/mil/passes/test_noop_elimination.py @@ -3,18 +3,17 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import itertools +import numpy as np +import pytest + from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY from coremltools.converters.mil.testing_utils import ( assert_model_is_valid, get_op_types_in_program, apply_pass_and_basic_check, ) -from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY -import copy -import pytest -import itertools - -import numpy as np @pytest.mark.parametrize("op_type, pos, val", itertools.product(['add', 'mul', 'floor_div', 'pow', 'real_div', 'sub'], ['x', 'y'], [0, 1, [0, 0, 0, 0], [1, 1, 1, 1]])) diff --git a/coremltools/converters/mil/mil/passes/test_pad_conv_pass.py b/coremltools/converters/mil/mil/passes/test_pad_conv_pass.py index 05e881026..9f4a952aa 100644 --- a/coremltools/converters/mil/mil/passes/test_pad_conv_pass.py +++ b/coremltools/converters/mil/mil/passes/test_pad_conv_pass.py @@ -3,6 +3,9 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import numpy as np +import unittest + from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.testing_utils import ( assert_op_count_match, @@ -10,10 +13,7 @@ get_op_types_in_program, apply_pass_and_basic_check, ) -import unittest -import pytest -import numpy as np np.random.seed(1984) diff --git a/coremltools/converters/mil/mil/passes/test_passes.py b/coremltools/converters/mil/mil/passes/test_passes.py index 003bead26..e59359621 100644 --- a/coremltools/converters/mil/mil/passes/test_passes.py +++ b/coremltools/converters/mil/mil/passes/test_passes.py @@ -26,8 +26,6 @@ np.random.seed(1984) validate_model = True - - def test_const_elimination(): @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))]) def prog(x): @@ -919,8 +917,8 @@ def prog(x): block = prog.functions["main"] # Reorder `split` op to test op with multiple output case - from .topological_reorder import move_operations_to_the_end_block - move_operations_to_the_end_block(block, ['split']) + from .topological_reorder import _move_operations_to_the_end_block + _move_operations_to_the_end_block(block, ['split']) assert get_op_types_in_program(prog) == ['square', 'relu', 'split', 'square', 'add'] @@ -1121,3 +1119,449 @@ def program(x): block.outputs[2].name: (3,), }, ) + + +class TestRemoveRedundantOpsPass: + + def test_redundant_ops_just_after_input_valid_pattern_1(self): + """ + Input graph: + input----->transpose(perm=[0, 2, 1])--->add---> add ---> out + | ^ ^ + | | | + |---->transpose(perm=[0, 2, 1])---- | + | | + | | + |---->transpose(perm=[0, 2, 1])------------ + + Output graph: + input----->transpose(perm=[0, 2, 1])--->add---> add ----> out + | ^ ^ + | | | + |------------- | + | | + |-------------------- + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3, 5))]) + def prog(x): + x1 = mb.transpose(x=x, perm=[0, 2, 1]) + x2 = mb.transpose(x=x, perm=[0, 2, 1]) + x3 = mb.transpose(x=x, perm=[0, 2, 1]) + z = mb.add(x=x1, y=x2) + z = mb.add(x=z, y=x3) + return z + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["transpose", "transpose", "transpose", "add", "add"] + assert get_op_types_in_program(prog) == ["transpose", "add", "add"] + assert_model_is_valid( + prog, + {"x": (2, 3, 5)}, + expected_output_shapes={block.outputs[0].name: (2, 5, 3)}, + ) + + def test_redundant_ops_just_after_input_valid_pattern_2(self): + """ + Input graph: + input----->leaky_relu(alpha=0.3)--->add---> add ---> out + | ^ ^ + | | | + |----->leaky_relu(alpha=0.3)--- | + | | + | | + |---->leaky_relu(alpha=0.3)------------ + + Output graph: + input--------->leaky_relu(alpha=0.3)--->add---> add ----> out + | ^ ^ + | | | + |------------- | + | | + |--------------------- + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3, 5))]) + def prog(x): + x1 = mb.leaky_relu(x=x, alpha=0.3) + x2 = mb.leaky_relu(x=x, alpha=0.3) + x3 = mb.leaky_relu(x=x, alpha=0.3) + z = mb.add(x=x1, y=x2) + z = mb.add(x=z, y=x3) + return z + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["leaky_relu", "leaky_relu", "leaky_relu", "add", "add"] + assert get_op_types_in_program(prog) == ["leaky_relu", "add", "add"] + assert_model_is_valid( + prog, + {"x": (2, 3, 5)}, + expected_output_shapes={block.outputs[0].name: (2, 3, 5)}, + ) + + def test_redundant_ops_just_after_input_invalid_pattern_1(self): + """ + input----->transpose(perm=[0, 2, 1])---> reshape(shape=[-1]) -----> add ---> out + | ^ + | | + |---->transpose(perm=[1, 0, 2])----> reshape(shape=[-1])------ + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3, 5))]) + def prog(x): + x1 = mb.transpose(x=x, perm=[0, 2, 1]) + x2 = mb.transpose(x=x, perm=[1, 0, 2]) + x1 = mb.reshape(x=x1, shape=[-1]) + x2 = mb.reshape(x=x2, shape=[-1]) + z = mb.add(x=x1, y=x2) + return z + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["transpose", "transpose", "reshape", "reshape", "add"] + assert get_op_types_in_program(prog) == ["transpose", "transpose", "reshape", "reshape", "add"] + assert_model_is_valid( + prog, + {"x": (2, 3, 5)}, + expected_output_shapes={block.outputs[0].name: (30,)}, + ) + + def test_redundant_ops_just_after_input_invalid_pattern_2(self): + """ + input----->leaky_relu(alpha=0.3) -----> add ---> out + | ^ + | | + |---->leaky_relu(alpha=0.4)------- + + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3, 5))]) + def prog(x): + x1 = mb.leaky_relu(x=x, alpha=0.3) + x2 = mb.leaky_relu(x=x, alpha=0.4) + z = mb.add(x=x1, y=x2) + return z + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["leaky_relu", "leaky_relu", "add"] + assert get_op_types_in_program(prog) == ["leaky_relu", "leaky_relu", "add"] + assert_model_is_valid( + prog, + {"x": (2, 3, 5)}, + expected_output_shapes={block.outputs[0].name: (2, 3, 5)}, + ) + + def test_redundant_ops_just_after_input_invalid_pattern_3(self): + """ + test case, when inputs of 1 op is a subset of the inputs of the other op + + input----->layer_norm1 -----> add ---> out + | ^ + | | + |---->layer_norm2------- + + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 2))]) + def prog(x): + x1 = mb.layer_norm(x=x, axes=[2], epsilon=1e-4) + gamma_val = np.array([1.0, 1.0], dtype=np.float32) + beta_val = np.array([1.0, 0.0], dtype=np.float32) + x2 = mb.layer_norm(x=x, axes=[2], epsilon=1e-4, gamma=gamma_val, beta=beta_val) + z = mb.add(x=x1, y=x2) + return z + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["layer_norm", "layer_norm", "add"] + assert get_op_types_in_program(prog) == ["layer_norm", "layer_norm", "add"] + assert_model_is_valid( + prog, + {"x": (1, 3, 2)}, + expected_output_shapes={block.outputs[0].name: (1, 3, 2)}, + ) + + @staticmethod + def _make_repeated_conv_prog(redundant_conv=True): + prog = Program() + func_inputs = {"x": mb.placeholder(shape=[1, 4, 5, 5])} + with Function(func_inputs) as ssa_fun: + x = ssa_fun.inputs["x"] + x = mb.relu(x=x) + W = np.random.rand(8, 4, 3, 3) + if redundant_conv: + bias = np.random.rand(8) + x1 = mb.conv(x=x, weight=W, bias=bias, pad_type="same", strides=[1, 1]) + x2 = mb.conv(x=x, weight=W, bias=bias, pad_type="same", strides=[1, 1]) + else: + x1 = mb.conv(x=x, weight=W, bias=np.random.rand(8), pad_type="same", strides=[1, 1]) + x2 = mb.conv(x=x, weight=W, bias=np.random.rand(8), pad_type="same", strides=[1, 1]) + x1 = mb.relu(x=x1) + x2 = mb.relu(x=x2) + x1 = mb.avg_pool(x=x1, kernel_sizes=[2, 2], strides=[1, 1], pad_type="same") + z = mb.concat(values=(x1, x2), axis=-3) + ssa_fun.set_outputs([z]) + prog.add_function("main", ssa_fun) + return prog + + def test_redundant_ops_inside_graph_valid_pattern(self): + """ + Input graph: + input--> relu--------->conv------>relu----> pool ---> concat ---> out + | ^ + | | + |---->conv---->relu---------------------------- + + Output graph: + input-> relu--->conv------>relu----> pool ---> concat ---> out + | ^ + | | + |------------------- + """ + prog = self._make_repeated_conv_prog(redundant_conv=True) + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["relu", "conv", "conv", "relu", "relu", "avg_pool", "concat"] + assert get_op_types_in_program(prog) == ["relu", "conv", "relu", "avg_pool", "concat"] + assert_model_is_valid( + prog, + {"x": (1, 4, 5, 5)}, + expected_output_shapes={block.outputs[0].name: (1, 16, 5, 5)}, + ) + + def test_redundant_ops_inside_graph_invalid_pattern(self): + """ + input--->relu--------->conv1------>relu----> pool ---> concat ---> out + | ^ + | | + |---->conv2---->relu--------------------------- + """ + prog = self._make_repeated_conv_prog(redundant_conv=False) + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["relu", "conv", "conv", "relu", "relu", "avg_pool", "concat"] + assert get_op_types_in_program(prog) == ["relu", "conv", "conv", "relu", "relu", "avg_pool", "concat"] + assert_model_is_valid( + prog, + {"x": (1, 4, 5, 5)}, + expected_output_shapes={block.outputs[0].name: (1, 16, 5, 5)}, + ) + + def test_redundant_op_as_output_valid_pattern_1(self): + """ + Input graph: + input--------->relu------> out1 + | + | + |---->relu---->tanh---> out2 + + Output graph: + input--------->relu------> out1 + | + | + |---->tanh---> out2 + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3, 5))]) + def prog(x): + x1 = mb.relu(x=x) + x2 = mb.relu(x=x) + return x1, mb.tanh(x=x2) + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops" + ) + assert get_op_types_in_program(prev_prog) == ["relu", "relu", "tanh"] + assert get_op_types_in_program(prog) == ["relu", "tanh"] + assert_model_is_valid( + prog, + {"x": (2, 3, 5)}, + expected_output_shapes={block.outputs[0].name: (2, 3, 5), block.outputs[1].name: (2, 3, 5)}, + ) + + def test_redundant_op_as_output_invalid_pattern_1(self): + """ + Input graph: + input--------->relu------> out1 + | + | + |---->relu---> out2 + + "common::remove_redundant_ops" pass does not remove ops if their outputs + are block outputs. + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 3, 5))]) + def prog(x): + x1 = mb.relu(x=x) + x2 = mb.relu(x=x) + return x1, x2 + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops", + ) + assert get_op_types_in_program(prev_prog) == ["relu", "relu"] + assert get_op_types_in_program(prog) == ["relu", "relu"] + assert_model_is_valid( + prog, + {"x": (2, 3, 5)}, + expected_output_shapes={block.outputs[0].name: (2, 3, 5), block.outputs[1].name: (2, 3, 5)}, + ) + + def test_cond_block_program(self): + # to test: + # - identical ops within different blocks are not removed. The "relu" op inside true and false blocks + # are not removed since they are in different blocks + # - ops that have blocks inside them are not removed. There are two cond ops here, with identical inputs + # but they are not removed, since they are ops that have nested block inside them + @mb.program(input_specs=[mb.TensorSpec(shape=(1,))]) + def prog(x): + x1 = mb.cast(x=x, dtype="bool") + def true_fn(): + x = mb.relu(x=x1) + x = mb.cast(x=x, dtype="fp32") + return mb.add(x=x, y=1) + + def false_fn(): + x = mb.relu(x=x1) + x = mb.cast(x=x, dtype="fp32") + return mb.add(x=x, y=-1) + + z1 = mb.cond(pred=x1, _true_fn=true_fn, _false_fn=false_fn) + z2 = mb.cond(pred=x1, _true_fn=true_fn, _false_fn=false_fn) + z = mb.add(x=z1, y=z2) + return z + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops", + ) + assert get_op_types_in_program(prev_prog) == ["cast", "cond", "cond", "add"] + assert get_op_types_in_program(prog) == ["cast", "cond", "cond", "add"] + cond_op = prog.find_ops(op_type="cond")[0] + assert cond_op.blocks[0].operations[0].op_type == "relu" + assert cond_op.blocks[1].operations[0].op_type == "relu" + assert_model_is_valid( + prog, + {"x": (1,)}, + expected_output_shapes={block.outputs[0].name: (1,)}, + ) + + def test_concat_op_pattern(self): + ''' + Input graph: + ---------------> concat ------> log ------> out1 + | ^ + | | + input--------->relu------> concat ------> relu----> out2 + | ^ | + | | | + |---->tanh-------------------- + + Output graph: + |------>log ------> out1 + | + | + input--------->relu------> concat ------> relu----> out2 + | ^ + | | + |---->tanh--------- + ''' + + @mb.program(input_specs=[mb.TensorSpec(shape=(10, 5))]) + def prog(x): + x1 = mb.relu(x=x) + x2 = mb.tanh(x=x) + c1 = mb.concat(values=(x1, x2), axis=0) + c2 = mb.concat(values=(x1, x2), axis=0) + z1 = mb.log(x=c1) + z2 = mb.relu(x=c2) + return z1, z2 + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops", + ) + assert get_op_types_in_program(prev_prog) == ["relu", "tanh", "concat", "concat", "log", "relu"] + assert get_op_types_in_program(prog) == ["relu", "tanh", "concat", "log", "relu"] + assert_model_is_valid( + prog, + {"x": (10, 5)}, + expected_output_shapes={block.outputs[0].name: (20, 5), block.outputs[1].name: (20, 5)}, + ) + + def test_multiple_redundant_child_ops_pattern(self): + ''' + Input graph + + input -------------> reshape ----------> add ---------> out1 + | ^ + | | + |-------> reshape --------------- + | + |------> slice_by_size-----> add ----------> out2 + | ^ + | | + |------> slice_by_size ------- + + Output graph + + input -------------> reshape ----------> add ------------> out1 + | | ^ + | | | + | |--------- + | + |------> slice_by_size----------> add -----------------> out2 + | ^ + | | + |--------------------- + + ''' + + @mb.program(input_specs=[mb.TensorSpec(shape=(10, 5, 4))]) + def prog(x): + x1 = mb.reshape(x=x, shape=[5, 2, -1]) + x2 = mb.reshape(x=x, shape=[5, 2, -1]) + x3 = mb.slice_by_size(x=x, begin=[0, 0, 1], size=[2, 4, 3]) + x4 = mb.slice_by_size(x=x, begin=[0, 0, 1], size=[2, 4, 3]) + z1 = mb.add(x=x1, y=x2) + z2 = mb.add(x=x3, y=x4) + return z1, z2 + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops", + ) + assert get_op_types_in_program(prev_prog) == ["reshape", "reshape", "slice_by_size", "slice_by_size", "add", "add"] + assert get_op_types_in_program(prog) == ["reshape", "slice_by_size", "add", "add"] + assert_model_is_valid( + prog, + {"x": (10, 5, 4)}, + expected_output_shapes={block.outputs[0].name: (5, 2, 20), block.outputs[1].name: (2, 4, 3)}, + ) + + def test_random_distribution_op_invalid_pattern(self): + """ + Identical random ops are not removed + + input----->cast---->random_uniform------> add ---> out + | ^ + | | + |---->random_uniform------------ + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(3,))]) + def prog(shape): + shape = mb.cast(x=shape, dtype="int32") + x1 = mb.random_uniform(shape=shape, low=0.0, high=1.0, seed=11) + x2 = mb.random_uniform(shape=shape, low=0.0, high=1.0, seed=11) + return mb.add(x=x1, y=x2) + + prev_prog, _, block = apply_pass_and_basic_check( + prog, "common::remove_redundant_ops", + ) + assert get_op_types_in_program(prev_prog) == ["cast", "random_uniform", "random_uniform", "add"] + assert get_op_types_in_program(prog) == ["cast", "random_uniform", "random_uniform", "add"] + + diff --git a/coremltools/converters/mil/mil/passes/topological_reorder.py b/coremltools/converters/mil/mil/passes/topological_reorder.py index 096ef6664..ce9e05044 100644 --- a/coremltools/converters/mil/mil/passes/topological_reorder.py +++ b/coremltools/converters/mil/mil/passes/topological_reorder.py @@ -1,15 +1,14 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2021, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from collections import defaultdict -from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass -def move_operations_to_the_end_block(block, op_type_to_move): + +def _move_operations_to_the_end_block(block, op_type_to_move): # Moves ops with `op_type_to_move` in `block.operations` (list) to the end of the program. # Note: ops with `op_type_to_move` and is dead code are moved toward end, which can be eliminated # later with dead-code-elimination pass. @@ -59,7 +58,7 @@ def move_operations_to_the_end_block(block, op_type_to_move): # Collect input vars from sub-block if present relevant_inputs = set() for b in current_op.blocks: - relevant_inputs |= move_operations_to_the_end_block(b, op_type_to_move) # |= is set union + relevant_inputs |= _move_operations_to_the_end_block(b, op_type_to_move) # |= is set union # Collect vars from operation input for v in current_op.inputs.values(): @@ -99,8 +98,9 @@ def move_operations_to_the_end_block(block, op_type_to_move): vars_consumed_in_block.update(block.outputs) return vars_consumed_in_block + @register_pass(namespace="common") -def topological_reorder(prog): +class topological_reorder(AbstractGraphPass): """ Topologically reorders the list of operations in a program by re-ordering operations closers to their first use or at the end if it's not being consumed by any other operation. @@ -152,5 +152,6 @@ def topological_reorder(prog): x2 = mb.cast(x=x1_t, dtype="fp32") } -> x2, x4, x7, x8 """ - for f_name, f in prog.functions.items(): - move_operations_to_the_end_block(f, ['cast', 'transpose']) + def apply(self, prog): + for f_name, f in prog.functions.items(): + _move_operations_to_the_end_block(f, ['cast', 'transpose']) diff --git a/coremltools/converters/mil/mil/passes/use_reflection_padding.py b/coremltools/converters/mil/mil/passes/use_reflection_padding.py index 510f3643a..1eaef425b 100644 --- a/coremltools/converters/mil/mil/passes/use_reflection_padding.py +++ b/coremltools/converters/mil/mil/passes/use_reflection_padding.py @@ -4,11 +4,11 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _match_pattern(concat_op, block): - if concat_op.op_type != "concat": return False @@ -66,10 +66,10 @@ def _match_pattern(concat_op, block): return False input_shape = original_input.shape - # Check that the slices are in order + # Check that the slices are in order if begin[axis] != begin_index and begin[axis] != begin_index + input_shape[axis]: return False - begin_index = begin_index - 1 + begin_index = begin_index - 1 slice_ops_out.append(slice_op) @@ -83,7 +83,7 @@ def _replace_ops(block, concat_op, slice_ops, axis): pad_size = len(slice_ops) // 2 if axis == -1: pad = [pad_size, pad_size] - elif axis == -2: + elif axis == -2: pad = [pad_size, pad_size, 0, 0] else: return False @@ -96,13 +96,14 @@ def _replace_ops(block, concat_op, slice_ops, axis): block.remove_ops([concat_op] + slice_ops) return True + def _reflection_padding_block(block): for op in list(block.operations): _match_pattern(op, block) - + @register_pass(namespace="common") -def use_reflection_padding(prog): +class use_reflection_padding(AbstractGraphPass): """ Identify a reflection padding layer composed out of slices and concats. @@ -116,5 +117,7 @@ def use_reflection_padding(prog): Output graph: input(1, 2, 6, 8) -----0> pad(mode=reflect, size=[0, 0, 1, 1]) -----> out(1, 2, 6, 10) """ - for f in prog.functions.values(): - _reflection_padding_block(f) + + def apply(self, prog): + for f in prog.functions.values(): + _reflection_padding_block(f) diff --git a/coremltools/converters/mil/mil/program.py b/coremltools/converters/mil/mil/program.py index fcd547696..889cb8fd6 100644 --- a/coremltools/converters/mil/mil/program.py +++ b/coremltools/converters/mil/mil/program.py @@ -1,5 +1,3 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be @@ -8,6 +6,7 @@ import logging as _logging import numpy as _np import sympy as _sm + from . import types from .block import Function from .var import Var diff --git a/coremltools/converters/mil/mil/tests/test_programs.py b/coremltools/converters/mil/mil/tests/test_programs.py index f8c827f03..5371beb8e 100644 --- a/coremltools/converters/mil/mil/tests/test_programs.py +++ b/coremltools/converters/mil/mil/tests/test_programs.py @@ -1,15 +1,14 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -import pytest -import coremltools as ct + +import logging import numpy as np -from coremltools.converters.mil.mil import Builder as mb + import coremltools as ct -import logging +from coremltools.converters.mil.mil import Builder as mb + np.random.seed(0) diff --git a/coremltools/converters/mil/mil/types/__init__.py b/coremltools/converters/mil/mil/types/__init__.py index 5050e8804..5676a6ebf 100644 --- a/coremltools/converters/mil/mil/types/__init__.py +++ b/coremltools/converters/mil/mil/types/__init__.py @@ -1,9 +1,9 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +from math import log, exp + from .type_double import fp16, fp32, fp64, float, double, is_float from .type_int import ( int8, @@ -60,6 +60,5 @@ from .annotate import delay_type from .get_type_info import get_type_info from .global_methods import global_remap -from math import log, exp apply_delayed_types() diff --git a/coremltools/converters/mil/mil/types/annotate.py b/coremltools/converters/mil/mil/types/annotate.py index 0e47169e3..0ccd104c8 100644 --- a/coremltools/converters/mil/mil/types/annotate.py +++ b/coremltools/converters/mil/mil/types/annotate.py @@ -1,12 +1,9 @@ -# -*- coding: utf-8 -*- - # Copyright (c) 2020, Apple Inc. All rights reserved. # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - class delay_type_cls: def __getattr__(self, t): return t diff --git a/coremltools/converters/mil/testing_reqs.py b/coremltools/converters/mil/testing_reqs.py index 00d17d63d..6ab10a3ae 100644 --- a/coremltools/converters/mil/testing_reqs.py +++ b/coremltools/converters/mil/testing_reqs.py @@ -2,9 +2,8 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause - -import os import itertools +import os import numpy as np from numpy import linalg as la import pytest diff --git a/coremltools/converters/mil/testing_utils.py b/coremltools/converters/mil/testing_utils.py index d6f6d2282..eea8521e8 100644 --- a/coremltools/converters/mil/testing_utils.py +++ b/coremltools/converters/mil/testing_utils.py @@ -333,7 +333,8 @@ def apply_pass_and_basic_check(prog, pass_name, skip_output_name_check=False): Apply pass to the program """ prev_prog = copy.deepcopy(prog) - pass_name.apply(prog) if isinstance(pass_name, AbstractQuantizationPass) else PASS_REGISTRY[pass_name](prog) + graph_pass = pass_name if isinstance(pass_name, AbstractQuantizationPass) else PASS_REGISTRY[pass_name] + graph_pass(prog) block = prog.functions["main"] prev_block = prev_prog.functions["main"] if not skip_output_name_check: diff --git a/coremltools/models/nearest_neighbors/__init__.py b/coremltools/models/nearest_neighbors/__init__.py index 0746089ad..6dd839add 100644 --- a/coremltools/models/nearest_neighbors/__init__.py +++ b/coremltools/models/nearest_neighbors/__init__.py @@ -3,4 +3,4 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -from .builder import * +from .builder import KNearestNeighborsClassifierBuilder diff --git a/coremltools/models/nearest_neighbors/builder.py b/coremltools/models/nearest_neighbors/builder.py index d65c7dba6..9db54d9e3 100644 --- a/coremltools/models/nearest_neighbors/builder.py +++ b/coremltools/models/nearest_neighbors/builder.py @@ -3,13 +3,12 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause +import numpy as _np + from ...proto import FeatureTypes_pb2 from .. import datatypes - import coremltools -import numpy as _np - class KNearestNeighborsClassifierBuilder(object): """ diff --git a/coremltools/models/neural_network/__init__.py b/coremltools/models/neural_network/__init__.py index ac7868bea..2ea726803 100644 --- a/coremltools/models/neural_network/__init__.py +++ b/coremltools/models/neural_network/__init__.py @@ -10,4 +10,13 @@ from . import spec_inspection_utils from . import update_optimizer_utils from . import utils -from .builder import * + +# This import should be pruned rdar://84519338 +from .builder import ( + AdamParams, + datatypes, + set_training_features, + set_transform_interface_params, + SgdParams, + NeuralNetworkBuilder +) diff --git a/coremltools/models/neural_network/quantization_utils.py b/coremltools/models/neural_network/quantization_utils.py index 6be2aa71d..6b327b48a 100644 --- a/coremltools/models/neural_network/quantization_utils.py +++ b/coremltools/models/neural_network/quantization_utils.py @@ -727,7 +727,7 @@ def _quantize_nn_spec(nn_spec, nbits, qm, **kwargs): layer_type = layer.WhichOneof("layer") if not selector.do_quantize(layer): continue - print("Quantizing layer {}".format(layer.name)) + print("Quantizing layer {} of type {}".format(layer.name, layer_type)) # Convolution if layer_type == "convolution": @@ -1643,7 +1643,7 @@ def quantize_weights( "specification instead.") return qspec - quantized_model = _get_model(qspec) + quantized_model = _get_model(qspec, compute_units=full_precision_model.compute_unit) if not sample_data: return quantized_model diff --git a/coremltools/models/utils.py b/coremltools/models/utils.py index e868780e2..1ba75aea0 100644 --- a/coremltools/models/utils.py +++ b/coremltools/models/utils.py @@ -6,8 +6,6 @@ """ Utilities for the entire package. """ -from coremltools.converters.mil.mil.passes.name_sanitization_utils import NameSanitizer as _NameSanitizer -from coremltools.proto import Model_pb2 as _Model_pb2 import math as _math import numpy as _np import os as _os @@ -15,8 +13,11 @@ import sys as _sys import tempfile as _tempfile import warnings as _warnings + +from coremltools import ComputeUnit as _ComputeUnit +from coremltools.converters.mil.mil.passes.name_sanitization_utils import NameSanitizer as _NameSanitizer +from coremltools.proto import Model_pb2 as _Model_pb2 from .._deps import _HAS_SCIPY -from ..libmodelpackage import ModelPackage _MLMODEL_EXTENSION = ".mlmodel" @@ -24,9 +25,9 @@ try: - from ..libmodelpackage import ModelPackage -except ModuleNotFoundError: - pass + from ..libmodelpackage import ModelPackage as _ModelPackage +except: + _ModelPackage = None if _HAS_SCIPY: import scipy.sparse as _sp @@ -109,7 +110,12 @@ def save_spec(spec, filename, auto_set_specification_version=False): f.write(spec) if is_package: - package = ModelPackage(filename) + if _ModelPackage is None: + raise Exception( + "Unable to load libmodelpackage. Cannot save spec" + ) + + package = _ModelPackage(filename) model_name = _pathlib.Path(filename).with_suffix('.mlmodel').name # Root file is copied into the model package. Changes to in-memory JSON is commited to disk when package goes out of scope. @@ -140,13 +146,16 @@ def load_spec(filename): -------- save_spec """ - from ..proto import Model_pb2 + if _ModelPackage is None: + raise Exception( + "Unable to load libmodelpackage. Cannot make save spec." + ) - spec = Model_pb2.Model() + spec = _Model_pb2.Model() specfile = filename - if ModelPackage.isValid(filename): - specfile = ModelPackage(filename).getRootModel().path() + if _ModelPackage.isValid(filename): + specfile = _ModelPackage(filename).getRootModel().path() with open(specfile, "rb") as f: contents = f.read() @@ -264,7 +273,7 @@ def _convert_neural_network_weights_to_fp16(full_precision_model): return _get_model(_convert_neural_network_spec_weights_to_fp16(spec)) -def _get_model(spec): +def _get_model(spec, compute_units=_ComputeUnit.ALL): """ Utility to get the model and the data. """ @@ -273,7 +282,7 @@ def _get_model(spec): if isinstance(spec, MLModel): return spec else: - return MLModel(spec) + return MLModel(spec, compute_units=compute_units) def evaluate_regressor(model, data, target="target", verbose=False): diff --git a/coremltools/test/api/test_api_examples.py b/coremltools/test/api/test_api_examples.py index 27e4892cf..f31387bea 100644 --- a/coremltools/test/api/test_api_examples.py +++ b/coremltools/test/api/test_api_examples.py @@ -2,12 +2,24 @@ # # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause + import copy +from io import BytesIO +import numpy as np +import os +from os import getcwd, chdir +from os.path import exists +from PIL import Image +import pytest +from shutil import rmtree +from tempfile import mkdtemp +import urllib.request + import coremltools as ct -from coremltools.converters.mil.testing_utils import get_op_types_in_program from coremltools.converters.mil import Builder as mb -from coremltools.converters.mil.mil import Program, Function from coremltools.converters.mil.frontend.torch.test.testing_utils import _copy_input_data +from coremltools.converters.mil.mil import get_new_symbol, Program, Function +from coremltools.converters.mil.testing_utils import get_op_types_in_program from coremltools._deps import ( _HAS_TF_1, @@ -18,16 +30,17 @@ MSG_TORCH_NOT_FOUND, ) -import numpy as np -import os -from os.path import exists -from os import getcwd, chdir -import pytest -from shutil import rmtree -from tempfile import mkdtemp - if _HAS_TORCH: import torch + import torchvision + +if _HAS_TF_1 or _HAS_TF_2: + import tensorflow as tf + +if _HAS_TF_2: + import requests + from tensorflow import keras + from tensorflow.keras import layers ############################################################################### @@ -44,8 +57,6 @@ class TestTensorFlow1ConverterExamples: @staticmethod def test_convert_from_frozen_graph(tmpdir): - import tensorflow as tf - with tf.Graph().as_default() as graph: x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input") y = tf.nn.relu(x, name="output") @@ -61,7 +72,6 @@ def test_convert_from_frozen_graph(tmpdir): @staticmethod def test_convert_from_frozen_graph_file(tmpdir): # create the model to convert - import tensorflow as tf # write a toy frozen graph # Note that we usually needs to run freeze_graph() on tf.Graph() @@ -121,8 +131,6 @@ def test_convert_from_saved_model_dir(tmpdir): test_input = np.random.rand(1, 3, 5) - 0.5 # create the model to convert - import tensorflow as tf - with tf.compat.v1.Session() as sess: x = tf.placeholder(shape=(1, 3, 5), dtype=tf.float32) y = tf.nn.relu(x) @@ -151,9 +159,6 @@ def test_convert_from_saved_model_dir(tmpdir): @staticmethod def test_freeze_and_convert_matmul_graph(): # testing : https://coremltools.readme.io/docs/tensorflow-1#export-as-frozen-graph-and-convert - - import tensorflow as tf - graph = tf.Graph() with graph.as_default(): x = tf.placeholder(tf.float32, shape=[None, 20], name="input") @@ -162,11 +167,9 @@ def test_freeze_and_convert_matmul_graph(): y = tf.matmul(x, W) + b output_names = [y.op.name] - import tempfile - import os from tensorflow.python.tools.freeze_graph import freeze_graph - model_dir = tempfile.mkdtemp() + model_dir = mkdtemp() graph_def_file = os.path.join(model_dir, 'tf_graph.pb') checkpoint_file = os.path.join(model_dir, 'tf_model.ckpt') frozen_graph_file = os.path.join(model_dir, 'tf_frozen.pb') @@ -213,8 +216,6 @@ def setup_class(self): chdir(self._temp_dir) # create toy models for conversion examples - import tensorflow as tf - # write a toy tf.keras HDF5 model tf_keras_model = tf.keras.Sequential( [ @@ -235,26 +236,23 @@ def teardown_class(self): @staticmethod def test_convert_tf_keras_h5_file(tmpdir): - import tensorflow as tf - - x = tf.keras.Input(shape=(32,), name="input") - y = tf.keras.layers.Dense(16, activation="softmax")(x) - keras_model = tf.keras.Model(x, y) - save_dir = str(tmpdir) - h5_path = os.path.join(save_dir, "tf_keras_model.h5") - keras_model.save(h5_path) - - mlmodel = ct.convert(h5_path) - - test_input = np.random.rand(2, 32) - expected_val = keras_model(test_input) - results = mlmodel.predict({"input": test_input}) - np.testing.assert_allclose(results["Identity"], expected_val, rtol=1e-4) + for file_extension in ("h5", ".hdf5"): + x = tf.keras.Input(shape=(32,), name="input") + y = tf.keras.layers.Dense(16, activation="softmax")(x) + keras_model = tf.keras.Model(x, y) + save_dir = str(tmpdir) + path = os.path.join(save_dir, "tf_keras_model." + file_extension) + keras_model.save(path) + + mlmodel = ct.convert(path) + + test_input = np.random.rand(2, 32) + expected_val = keras_model(test_input) + results = mlmodel.predict({"input": test_input}) + np.testing.assert_allclose(results["Identity"], expected_val, rtol=1e-4) @staticmethod def test_convert_tf_keras_model(): - import tensorflow as tf - x = tf.keras.Input(shape=(32,), name="input") y = tf.keras.layers.Dense(16, activation="softmax")(x) keras_model = tf.keras.Model(x, y) @@ -270,8 +268,6 @@ def test_convert_tf_keras_model(): @pytest.mark.parametrize( "dtype", ['default', 'mil_type', 'np type']) def test_convert_tf_keras_applications_model(dtype): - import tensorflow as tf - tf_keras_model = tf.keras.applications.MobileNet( weights="imagenet", input_shape=(224, 224, 3) ) @@ -309,10 +305,6 @@ def test_convert_from_saved_model_dir(): def test_keras_custom_layer_model(): # testing : https://coremltools.readme.io/docs/tensorflow-2#conversion-from-user-defined-models - import tensorflow as tf - from tensorflow import keras - from tensorflow.keras import layers - class CustomDense(layers.Layer): def __init__(self, units=32): super(CustomDense, self).__init__() @@ -339,7 +331,6 @@ def call(self, inputs): @staticmethod def test_concrete_function_conversion(): # testing : https://coremltools.readme.io/docs/tensorflow-2#conversion-from-user-defined-models - import tensorflow as tf @tf.function(input_signature=[tf.TensorSpec(shape=(6,), dtype=tf.float32)]) def gelu_tanh_activation(x): @@ -354,7 +345,6 @@ def gelu_tanh_activation(x): @staticmethod def test_quickstart_example(): # testing: https://coremltools.readme.io/docs/introductory-quickstart#quickstart-example - import tensorflow as tf # TF 2.2.0 # Download MobileNetv2 (using tf.keras) keras_model = tf.keras.applications.MobileNetV2( @@ -364,7 +354,6 @@ def test_quickstart_example(): ) # Download class labels (from a separate file) - import urllib.request label_url = 'https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt' class_labels = urllib.request.urlopen(label_url).read().splitlines() class_labels = class_labels[1:] # remove the first class which is background @@ -403,10 +392,6 @@ def test_quickstart_example(): model.version = "2.0" # get an image - from PIL import Image - import requests - from io import BytesIO - img_url = 'https://files.readme.io/02e3586-daisy.jpg' response = requests.get(img_url) img = Image.open(BytesIO(response.content)) @@ -424,9 +409,6 @@ def test_quickstart_example(): class TestPyTorchConverterExamples: @staticmethod def test_convert_torch_vision_mobilenet_v2(tmpdir): - import torch - import torchvision - """ In this example, we'll instantiate a PyTorch classification model and convert it to Core ML. @@ -691,8 +673,6 @@ class TestInputs: @staticmethod @pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND) def test_input_noname(): - import tensorflow as tf - with tf.Graph().as_default() as graph: x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input") x1 = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input_1") @@ -710,8 +690,6 @@ def test_input_noname(): @staticmethod @pytest.mark.skipif(not _HAS_TF_1, reason=MSG_TF1_NOT_FOUND) def test_input_wrongname(): - import tensorflow as tf - with tf.Graph().as_default() as graph: x = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input") x1 = tf.placeholder(tf.float32, shape=(1, 2, 3), name="input_1") @@ -760,8 +738,6 @@ class TestFlexibleShape: @pytest.mark.skipif(not _HAS_TF_2, reason=MSG_TF2_NOT_FOUND) def test_tf2keras_shared_range_dim(use_symbol): # Test examples in https://coremltools.readme.io/docs/flexible-inputs - import tensorflow as tf - input_dim = 3 # None denotes seq_len dimension x1 = tf.keras.Input(shape=(None,input_dim), name="seq1") @@ -1245,6 +1221,58 @@ def test_tf2keras_optional_input(): results = mlmodel.predict({"required_input": test_input_x2}) np.testing.assert_allclose(results["Identity"], expected_val, rtol=1e-4) + @staticmethod + @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND) + def test_torch_classifier(): + class Net(torch.nn.Module): + def __init__(self): + super(Net, self).__init__() + self.linear1 = torch.nn.Linear(28 * 28, 100) + self.linear2 = torch.nn.Linear(100, 50) + self.final = torch.nn.Linear(50, 10) + self.relu = torch.nn.ReLU() + + def forward(self, img): # convert + flatten + x = img.view(-1, 28 * 28) + x = self.relu(self.linear1(x)) + x = self.relu(self.linear2(x)) + x = self.final(x) + return x + model = Net() + model.eval() + example_input = torch.rand(1, 28 * 28, 1) + traced_model = torch.jit.trace(model, example_input) + traced_model.eval() + + def _test_classifier(traced_model, example_input, class_type, backend): + label = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + if class_type == "str": + label = list(map(lambda x: str(x), label)) + classifier_config = ct.ClassifierConfig(label) + mlmodel = ct.convert( + traced_model, + source='pytorch', + convert_to=backend, + inputs=[ + ct.TensorType( + name="input", + shape=example_input.shape, + dtype=example_input.numpy().dtype, + ) + ], + classifier_config=classifier_config + ) + if ct.utils._is_macos(): + coreml_out = mlmodel.predict({"input": example_input.detach().numpy()}) + assert "classLabel" in coreml_out + key_type = str if class_type == "str" else int + assert isinstance(coreml_out["classLabel"], key_type) + + for class_type in ("str", "int"): + _test_classifier(traced_model, example_input, class_type, "neuralnetwork") + if ct.utils._macos_version() >= (12, 0): + _test_classifier(traced_model, example_input, class_type, "mlprogram") + @staticmethod @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND) def test_torch_optional_input(): @@ -1535,6 +1563,21 @@ def forward(self, x): ) assert isinstance(model._get_mil_internal(), ct.converters.mil.Program) + @staticmethod + @pytest.mark.skipif(not ct.utils._is_macos(), reason="Platform is not Mac OS") + def test_deepcopy_error_with_symbols_in_prog(): + prog = Program() + func_inputs = {"x": mb.placeholder(shape=[get_new_symbol(), 3]), + "y": mb.placeholder(shape=[2, 3])} + with Function(func_inputs) as ssa_fun: + x, y = ssa_fun.inputs["x"], ssa_fun.inputs["y"] + x = mb.relu(x=x) + z = mb.add(x=x, y=y) + ssa_fun.set_outputs([z]) + prog.add_function("main", ssa_fun) + mlmodel = ct.convert(prog, convert_to="mlprogram", compute_precision=ct.precision.FLOAT32) + prog2 = mlmodel._get_mil_internal() # this will invoke a deepcopy on the prog + @staticmethod @pytest.mark.skipif(not _HAS_TORCH or ct.utils._macos_version() < (12, 0), reason=MSG_TORCH_NOT_FOUND) @@ -1638,4 +1681,4 @@ def prog(x): with pytest.raises(ValueError, match=expected_err_str): mlmodel = ct.convert(prog, compute_precision=ct.precision.FLOAT16) with pytest.raises(ValueError, match=expected_err_str): - mlmodel = ct.convert(prog, compute_precision=ct.precision.FLOAT32) \ No newline at end of file + mlmodel = ct.convert(prog, compute_precision=ct.precision.FLOAT32) diff --git a/coremltools/test/api/test_api_visibilities.py b/coremltools/test/api/test_api_visibilities.py index c93c2a7d3..17181689e 100644 --- a/coremltools/test/api/test_api_visibilities.py +++ b/coremltools/test/api/test_api_visibilities.py @@ -58,7 +58,6 @@ def test_utils(self): "load_spec", "rename_feature", "save_spec", - "ModelPackage", ] _check_visible_modules(_get_visible_items(ct.utils), expected) diff --git a/coremltools/test/neural_network/test_quantization.py b/coremltools/test/neural_network/test_quantization.py index fa7c85a8e..a8939a3c5 100644 --- a/coremltools/test/neural_network/test_quantization.py +++ b/coremltools/test/neural_network/test_quantization.py @@ -3,31 +3,31 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause """Module containing unit tests for verifying various quantization.""" + +import numpy as np import os +import pytest import shutil import tempfile import unittest -import numpy as np -import pytest import coremltools +from coremltools import ComputeUnit as _ComputeUnit +from coremltools.models import ( + neural_network, + _MLMODEL_FULL_PRECISION, + _QUANTIZATION_MODE_LINEAR_QUANTIZATION, + _QUANTIZATION_MODE_LOOKUP_TABLE_KMEANS, + _QUANTIZATION_MODE_CUSTOM_LOOKUP_TABLE, +) import coremltools.models.datatypes as datatypes -from coremltools.models import neural_network import coremltools.models.neural_network.quantization_utils as quantization_utils from coremltools.models.neural_network.quantization_utils import ( activate_int8_int8_matrix_multiplications, MatrixMultiplyLayerSelector, _quantize_spec_weights, ) - from coremltools._deps import _HAS_KERAS2_TF -from coremltools.models import ( - _MLMODEL_FULL_PRECISION, - _QUANTIZATION_MODE_LINEAR_QUANTIZATION, - _QUANTIZATION_MODE_LOOKUP_TABLE_KMEANS, - _QUANTIZATION_MODE_CUSTOM_LOOKUP_TABLE, -) - @unittest.skipIf( not coremltools.utils._is_macos() or coremltools.utils._macos_version() < (10, 14), @@ -994,12 +994,15 @@ def test_batched_matmul_1bit_weight_quantized(self): self.compare() -@unittest.skipIf( +@pytest.mark.skipif( not coremltools.utils._is_macos() or coremltools.utils._macos_version() < (10, 15), - "Missing macOS 10.15+. Skipping tests.", + reason="Missing macOS 10.15+. Skipping tests.", ) -class QuantizeWeightsAPI(unittest.TestCase): - def test_embeddingND_quantize(self): +class TestQuantizeWeightsAPI: + @staticmethod + @pytest.mark.parametrize( + "compute_units", [_ComputeUnit.ALL, _ComputeUnit.CPU_AND_GPU, _ComputeUnit.CPU_ONLY]) + def test_embeddingND_quantize(compute_units): input_features = [("data", datatypes.Array(10, 1))] output_features = [("output", None)] builder = neural_network.NeuralNetworkBuilder( @@ -1016,44 +1019,28 @@ def test_embeddingND_quantize(self): ) spec = builder.spec - model_fp32 = coremltools.models.MLModel(spec) - self.assertEqual( - len(spec.neuralNetwork.layers[0].embeddingND.weights.floatValue), 6000 - ) + model_fp32 = coremltools.models.MLModel(spec, compute_units=compute_units) + assert len(spec.neuralNetwork.layers[0].embeddingND.weights.floatValue) == 6000 # quantize to FP16 model_fp16 = quantization_utils.quantize_weights(model_fp32, nbits=16) + assert model_fp16.compute_unit == compute_units spec_fp16 = model_fp16.get_spec() - self.assertEqual( - len(spec_fp16.neuralNetwork.layers[0].embeddingND.weights.floatValue), 0 - ) - self.assertEqual( - len(spec_fp16.neuralNetwork.layers[0].embeddingND.weights.float16Value), - 2 * 6000, - ) + assert len(spec_fp16.neuralNetwork.layers[0].embeddingND.weights.floatValue) == 0 + assert len(spec_fp16.neuralNetwork.layers[0].embeddingND.weights.float16Value) == 2 * 6000 # quantize to uint8 model_uint8 = quantization_utils.quantize_weights(model_fp32, nbits=8) + assert model_uint8.compute_unit == compute_units spec_uint8 = model_uint8.get_spec() - self.assertEqual( - len(spec_uint8.neuralNetwork.layers[0].embeddingND.weights.floatValue), 0 - ) - self.assertEqual( - len(spec_uint8.neuralNetwork.layers[0].embeddingND.weights.float16Value), 0 - ) - self.assertEqual( - len(spec_uint8.neuralNetwork.layers[0].embeddingND.weights.rawValue), 6000 - ) + assert len(spec_uint8.neuralNetwork.layers[0].embeddingND.weights.floatValue) == 0 + assert len(spec_uint8.neuralNetwork.layers[0].embeddingND.weights.float16Value) == 0 + assert len(spec_uint8.neuralNetwork.layers[0].embeddingND.weights.rawValue) == 6000 # quantize to uint5 model_uint5 = quantization_utils.quantize_weights(model_fp32, nbits=5) + assert model_uint5.compute_unit == compute_units spec_uint5 = model_uint5.get_spec() - self.assertEqual( - len(spec_uint5.neuralNetwork.layers[0].embeddingND.weights.floatValue), 0 - ) - self.assertEqual( - len(spec_uint5.neuralNetwork.layers[0].embeddingND.weights.float16Value), 0 - ) - self.assertEqual( - len(spec_uint5.neuralNetwork.layers[0].embeddingND.weights.rawValue), 3750 - ) # 3750 = 5*6000/8 + assert len(spec_uint5.neuralNetwork.layers[0].embeddingND.weights.floatValue) == 0 + assert len(spec_uint5.neuralNetwork.layers[0].embeddingND.weights.float16Value) == 0 + assert len(spec_uint5.neuralNetwork.layers[0].embeddingND.weights.rawValue) == 3750 # 3750 = 5*6000/8 diff --git a/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py b/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py index 486f8bf54..73fd76059 100644 --- a/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py +++ b/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py @@ -158,7 +158,7 @@ def test_conversion_bad_inputs(self): spec = skl_converter.convert(model, "data", "out") -@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/84898245") @unittest.skipIf(not _HAS_SKLEARN, "Missing sklearn. Skipping tests.") @unittest.skipIf(not _HAS_XGBOOST, "Skipping, no xgboost") class GradientBoostingBinaryClassifierXGboostTest(unittest.TestCase): @@ -224,7 +224,7 @@ def test_conversion_bad_inputs(self): spec = xgb_converter.convert(model, "data", "out", mode="classifier") -@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/84898245") @unittest.skipIf(not _HAS_SKLEARN, "Missing sklearn. Skipping tests.") @unittest.skipIf(not _HAS_XGBOOST, "Skipping, no xgboost") class GradientBoostingMulticlassClassifierXGboostTest(unittest.TestCase): diff --git a/coremltools/test/xgboost_tests/test_boosted_trees_classifier_numeric.py b/coremltools/test/xgboost_tests/test_boosted_trees_classifier_numeric.py index fd94325b9..713293a42 100644 --- a/coremltools/test/xgboost_tests/test_boosted_trees_classifier_numeric.py +++ b/coremltools/test/xgboost_tests/test_boosted_trees_classifier_numeric.py @@ -210,7 +210,7 @@ def _classifier_stress_test(self): self._train_convert_evaluate_assert(**arg) -@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/84898245") @unittest.skipIf(not _HAS_SKLEARN, "Missing sklearn. Skipping tests.") @unittest.skipIf(not _HAS_XGBOOST, "Skipping, no xgboost") class BoostedTreeBinaryClassificationBostonHousingXGboostNumericTest( @@ -241,7 +241,7 @@ def test_binary_classifier_stress_test(self): self._classifier_stress_test() -@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/84898245") @unittest.skipIf(not _HAS_SKLEARN, "Missing sklearn. Skipping tests.") @unittest.skipIf(not _HAS_XGBOOST, "Skipping, no xgboost") class BoostedTreeMultiClassClassificationBostonHousingXGboostNumericTest( diff --git a/coremltools/test/xgboost_tests/test_boosted_trees_regression.py b/coremltools/test/xgboost_tests/test_boosted_trees_regression.py index 10d587ef1..7e6c7c620 100644 --- a/coremltools/test/xgboost_tests/test_boosted_trees_regression.py +++ b/coremltools/test/xgboost_tests/test_boosted_trees_regression.py @@ -87,7 +87,7 @@ def test_conversion_bad_inputs(self): spec = skl_converter.convert(model, "data", "out") -@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (10, 16), "rdar://problem/84898245") @unittest.skipIf(not _HAS_SKLEARN, "Missing scikit-learn. Skipping tests.") @unittest.skipIf(not _HAS_XGBOOST, "Skipping, no xgboost") class BoostedTreeRegressorXGboostTest(unittest.TestCase): diff --git a/coremltools/test/xgboost_tests/test_boosted_trees_regression_numeric.py b/coremltools/test/xgboost_tests/test_boosted_trees_regression_numeric.py index c898caea6..bf1f09e37 100644 --- a/coremltools/test/xgboost_tests/test_boosted_trees_regression_numeric.py +++ b/coremltools/test/xgboost_tests/test_boosted_trees_regression_numeric.py @@ -3,11 +3,10 @@ # Use of this source code is governed by a BSD-3-clause license that can be # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -import unittest - import pandas as pd import itertools import pytest +import unittest from coremltools._deps import _HAS_SKLEARN, _HAS_XGBOOST from coremltools.models.utils import evaluate_regressor, _macos_version, _is_macos @@ -88,7 +87,7 @@ def test_boston_housing_parameter_stress_test(self): self._train_convert_evaluate_assert(**arg) -@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/84898245") @unittest.skipIf(not _HAS_XGBOOST, "Missing xgboost. Skipping") @unittest.skipIf(not _HAS_SKLEARN, "Missing scikit-learn. Skipping tests.") class XgboostBoosterBostonHousingNumericTest(unittest.TestCase): @@ -199,7 +198,7 @@ def test_boston_housing_parameter_stress_test(self): self._train_convert_evaluate_assert(arg) -@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/75172473") +@unittest.skipIf(_macos_version() >= (12, 0), "rdar://problem/84898245") @unittest.skipIf(not _HAS_XGBOOST, "Missing xgboost. Skipping") @unittest.skipIf(not _HAS_SKLEARN, "Missing sklearn. Skipping tests.") class XGboostRegressorBostonHousingNumericTest(unittest.TestCase): diff --git a/coremltools/version.py b/coremltools/version.py index 5c8bf39d5..2a886032f 100644 --- a/coremltools/version.py +++ b/coremltools/version.py @@ -4,4 +4,4 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause -__version__ = "5.0" # VERSION_STRING +__version__ = "5.1.0" # VERSION_STRING diff --git a/mlmodel/src/Format.hpp b/mlmodel/src/Format.hpp index 7d901cfe3..81d67b261 100644 --- a/mlmodel/src/Format.hpp +++ b/mlmodel/src/Format.hpp @@ -9,6 +9,7 @@ #if __apple_build_version__ < 10010028 #pragma clang diagnostic ignored "-Wextended-offsetof" #endif +#pragma clang diagnostic ignored "-Wdeprecated-declarations" #include #include diff --git a/reqs/build.pip b/reqs/build.pip index ca823355d..54a4425ea 100644 --- a/reqs/build.pip +++ b/reqs/build.pip @@ -2,4 +2,6 @@ numpy<1.20; platform_machine != "arm64" protobuf pytest six +sympy +tqdm wheel diff --git a/reqs/test.pip b/reqs/test.pip index 4ce664ea3..d1c7bf0f4 100644 --- a/reqs/test.pip +++ b/reqs/test.pip @@ -21,10 +21,10 @@ six sympy > 1.6 tensorflow==1.14.0; python_version < '3.8' torch==1.5.0; python_version == '3.5' -torch==1.9.0; python_version > '3.5' +torch==1.9.1; python_version > '3.5' torchvision==0.6.1; python_version == '3.5' -torchvision==0.10.0; python_version > '3.5' -xgboost +torchvision==0.10.1; python_version > '3.5' +xgboost==1.4.2 mock wrapt pyyaml > 5.3