From edf87ee02effb3d4223148a45cad10c7249ae40b Mon Sep 17 00:00:00 2001 From: Toby Roseman Date: Mon, 1 Aug 2022 14:36:01 -0700 Subject: [PATCH] 6.0b2 Release (#1566) --- CMakeLists.txt | 1 - coremlpython/CoreMLPython.mm | 38 +- coremlpython/CoreMLPythonArray.h | 5 + coremlpython/CoreMLPythonUtils.mm | 185 ++--- coremltools/__init__.py | 4 + coremltools/converters/_converters_entry.py | 3 + .../mil/_deployment_compatibility.py | 8 +- .../converters/mil/backend/mil/load.py | 2 +- .../passes/adjust_io_to_supported_types.py | 55 +- .../mil/passes/fuse_activation_silu.py | 5 +- .../mil/passes/homogenize_input_dtypes.py | 28 +- .../passes/insert_image_preprocessing_op.py | 45 +- .../mil/backend/mil/passes/mil_passes.py | 4 +- .../mil/backend/mil/passes/test_passes.py | 49 +- .../converters/mil/backend/nn/op_mapping.py | 15 +- .../mil/backend/nn/passes/test_passes.py | 4 +- coremltools/converters/mil/converter.py | 11 + .../passes/generic_pass_infrastructure.py | 18 +- .../mil/experimental/passes/readme.md | 1 - .../mil/frontend/milproto/helper.py | 4 - .../converters/mil/frontend/milproto/load.py | 39 +- .../mil/frontend/milproto/test_load.py | 76 ++- .../mil/frontend/tensorflow/converter.py | 6 +- .../mil/frontend/tensorflow/dialect_ops.py | 47 +- .../mil/frontend/tensorflow/load.py | 7 +- .../converters/mil/frontend/tensorflow/ops.py | 25 +- .../backfill_make_list_elem_type.py | 26 +- .../ssa_passes/tf_lstm_to_core_lstm.py | 48 +- .../tensorflow/test/test_conversion_api.py | 9 +- .../tensorflow/test/test_custom_ops.py | 8 +- .../mil/frontend/tensorflow/test/test_load.py | 2 +- .../mil/frontend/tensorflow/test/test_ops.py | 189 +++--- .../frontend/tensorflow/test/testing_utils.py | 19 +- .../tf_graph_pass/insert_get_tuple.py | 2 +- .../mil/frontend/tensorflow2/converter.py | 4 +- .../mil/frontend/tensorflow2/load.py | 12 +- .../ssa_passes/remove_vacuous_cond.py | 25 +- .../frontend/tensorflow2/test/test_v2_load.py | 33 + .../frontend/tensorflow2/test/test_v2_ops.py | 25 +- .../tensorflow2/test/test_v2_ops_tf_keras.py | 9 +- .../tensorflow2/test/testing_utils.py | 13 +- .../mil/frontend/torch/converter.py | 10 +- .../mil/frontend/torch/dialect_ops.py | 10 +- .../converters/mil/frontend/torch/load.py | 3 +- .../converters/mil/frontend/torch/ops.py | 128 +++- .../ssa_passes/torch_tensor_assign_to_core.py | 11 +- .../torch_upsample_to_core_upsample.py | 40 +- .../torch/test/test_conversion_api.py | 134 +++- .../frontend/torch/test/test_custom_ops.py | 2 +- .../mil/frontend/torch/test/test_torch_ops.py | 159 ++++- .../mil/frontend/torch/test/testing_utils.py | 19 +- coremltools/converters/mil/mil/block.py | 331 +++++---- coremltools/converters/mil/mil/builder.py | 17 +- coremltools/converters/mil/mil/input_type.py | 23 +- coremltools/converters/mil/mil/operation.py | 11 + .../converters/mil/mil/ops/defs/__init__.py | 205 +----- .../converters/mil/mil/ops/defs/_utils.py | 72 ++ .../mil/mil/ops/defs/iOS15/__init__.py | 204 ++++++ .../mil/ops/defs/{ => iOS15}/activation.py | 40 +- .../mil/mil/ops/defs/{ => iOS15}/classify.py | 2 +- .../mil/ops/defs/{ => iOS15}/control_flow.py | 36 +- .../mil/mil/ops/defs/{ => iOS15}/conv.py | 6 +- .../defs/{ => iOS15}/elementwise_binary.py | 40 +- .../ops/defs/{ => iOS15}/elementwise_unary.py | 54 +- .../ops/defs/{ => iOS15}/image_resizing.py | 632 +++++++++--------- .../mil/mil/ops/defs/{ => iOS15}/linear.py | 10 +- .../mil/ops/defs/{ => iOS15}/normalization.py | 13 +- .../mil/mil/ops/defs/{ => iOS15}/pool.py | 6 +- .../mil/mil/ops/defs/{ => iOS15}/random.py | 10 +- .../mil/mil/ops/defs/{ => iOS15}/recurrent.py | 8 +- .../mil/mil/ops/defs/{ => iOS15}/reduction.py | 28 +- .../ops/defs/{ => iOS15}/scatter_gather.py | 12 +- .../ops/defs/{ => iOS15}/tensor_operation.py | 39 +- .../defs/{ => iOS15}/tensor_transformation.py | 104 +-- .../mil/mil/ops/defs/iOS16/__init__.py | 20 + .../mil/ops/defs/{ => iOS16}/constexpr_ops.py | 37 +- .../mil/mil/ops/defs/iOS16/image_resizing.py | 42 ++ .../mil/ops/defs/iOS16/tensor_operation.py | 74 ++ .../ops/defs/iOS16/tensor_transformation.py | 59 ++ coremltools/converters/mil/mil/ops/helper.py | 24 + .../converters/mil/mil/ops/registry.py | 180 ++++- .../mil/mil/ops/tests/test_activation.py | 2 - .../mil/mil/ops/tests/test_const.py | 4 +- .../mil/mil/ops/tests/test_constexpr_ops.py | 59 +- .../mil/mil/ops/tests/test_control_flow.py | 7 +- .../converters/mil/mil/ops/tests/test_conv.py | 8 +- .../mil/ops/tests/test_elementwise_binary.py | 3 - .../mil/ops/tests/test_elementwise_unary.py | 9 - .../mil/mil/ops/tests/test_image_resizing.py | 48 +- .../mil/mil/ops/tests/test_linear.py | 5 +- .../mil/mil/ops/tests/test_normalization.py | 6 - .../mil/mil/ops/tests/test_reduction.py | 30 + .../mil/mil/ops/tests/test_scatter_gather.py | 7 +- .../mil/ops/tests/test_tensor_operation.py | 48 ++ .../ops/tests/test_tensor_transformation.py | 70 +- .../mil/mil/ops/tests/testing_utils.py | 2 +- .../passes/add_conv_transpose_output_shape.py | 21 +- .../mil/passes/apply_common_pass_pipeline.py | 2 + .../mil/mil/passes/cast_optimization.py | 99 +-- .../mil/mil/passes/compression_passes.py | 273 ++++---- .../mil/mil/passes/concat_to_pixel_shuffle.py | 31 +- .../mil/mil/passes/const_elimination.py | 83 ++- .../mil/mil/passes/conv_batchnorm_fusion.py | 93 ++- .../mil/mil/passes/conv_bias_fusion.py | 511 +++++++------- .../mil/mil/passes/conv_scale_fusion.py | 99 ++- .../mil/passes/detect_concat_interleave.py | 8 +- .../mil/mil/passes/divide_to_multiply.py | 30 +- .../passes/elementwise_batchnorm_fusion.py | 6 +- .../mil/mil/passes/gelu_exact_fusion.py | 21 +- .../converters/mil/mil/passes/graph_pass.py | 14 - .../converters/mil/mil/passes/helper.py | 36 +- .../layernorm_instancenorm_pattern_fusion.py | 23 +- .../mil/mil/passes/leaky_relu_fusion.py | 16 +- .../mil/mil/passes/linear_bias_fusion.py | 93 ++- .../mil/passes/loop_invariant_elimination.py | 20 +- .../mil/passes/matmul_weight_bias_fusion.py | 74 +- .../mil/passes/merge_consecutive_paddings.py | 25 +- .../mil/mil/passes/noop_elimination.py | 149 +++-- .../mil/mil/passes/onehot_matmul_to_gather.py | 8 +- .../mil/mil/passes/pad_conv_connect.py | 6 +- .../mil/mil/passes/quantization_passes.py | 83 +-- .../mil/mil/passes/rank0_expand_dims_swap.py | 12 +- .../mil/mil/passes/reduce_mean_fusion.py | 15 +- .../mil/mil/passes/reduce_transposes.py | 73 +- .../mil/mil/passes/remove_redundant_ops.py | 61 +- .../mil/mil/passes/remove_symbolic_reshape.py | 27 +- .../mil/mil/passes/replace_stack_reshape.py | 17 +- .../mil/mil/passes/test_compression_passes.py | 78 +++ .../mil/mil/passes/test_noop_elimination.py | 24 + .../converters/mil/mil/passes/test_passes.py | 310 ++++++++- .../mil/mil/passes/topological_reorder.py | 26 +- .../mil/mil/passes/use_reflection_padding.py | 31 +- coremltools/converters/mil/mil/program.py | 57 ++ .../converters/mil/mil/tests/test_block.py | 57 +- .../converters/mil/mil/tests/test_programs.py | 140 ++++ coremltools/converters/mil/mil/var.py | 50 ++ coremltools/converters/mil/testing_utils.py | 13 +- .../models/ml_program/compression_utils.py | 15 +- coremltools/models/model.py | 53 +- .../neural_network/quantization_utils.py | 2 +- coremltools/test/api/test_api_examples.py | 49 +- coremltools/test/api/test_api_visibilities.py | 16 +- .../test/ml_program/test_compression.py | 18 +- .../test/modelpackage/test_modelpackage.py | 4 + .../neural_network/test_numpy_nn_layers.py | 53 +- .../test/neural_network/test_quantization.py | 14 +- .../test_random_forest_regression.py | 2 +- .../test_boosted_trees_classifier.py | 9 +- coremltools/version.py | 2 +- reqs/build.pip | 1 + reqs/test.pip | 4 +- reqs/test_tf2.pip | 2 +- scripts/env_create.sh | 5 - 153 files changed, 4515 insertions(+), 2821 deletions(-) create mode 100644 coremltools/converters/mil/mil/ops/defs/iOS15/__init__.py rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/activation.py (97%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/classify.py (99%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/control_flow.py (97%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/conv.py (99%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/elementwise_binary.py (96%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/elementwise_unary.py (96%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/image_resizing.py (96%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/linear.py (98%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/normalization.py (98%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/pool.py (99%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/random.py (98%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/recurrent.py (99%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/reduction.py (97%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/scatter_gather.py (99%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/tensor_operation.py (98%) rename coremltools/converters/mil/mil/ops/defs/{ => iOS15}/tensor_transformation.py (92%) create mode 100644 coremltools/converters/mil/mil/ops/defs/iOS16/__init__.py rename coremltools/converters/mil/mil/ops/defs/{ => iOS16}/constexpr_ops.py (93%) create mode 100644 coremltools/converters/mil/mil/ops/defs/iOS16/image_resizing.py create mode 100644 coremltools/converters/mil/mil/ops/defs/iOS16/tensor_operation.py create mode 100644 coremltools/converters/mil/mil/ops/defs/iOS16/tensor_transformation.py create mode 100644 coremltools/converters/mil/mil/ops/helper.py create mode 100644 coremltools/converters/mil/mil/passes/test_compression_passes.py diff --git a/CMakeLists.txt b/CMakeLists.txt index f6d278ebd..b65719798 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,5 @@ cmake_minimum_required(VERSION 3.10.2) -set(CMAKE_OSX_ARCHITECTURES x86_64) set(CMAKE_DISABLE_IN_SOURCE_BUILD ON) project(coremltools) diff --git a/coremlpython/CoreMLPython.mm b/coremlpython/CoreMLPython.mm index 28a523106..ea9bad97b 100644 --- a/coremlpython/CoreMLPython.mm +++ b/coremlpython/CoreMLPython.mm @@ -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 #import "CoreMLPythonArray.h" #import "CoreMLPython.h" #import "CoreMLPythonUtils.h" #import "Globals.hpp" #import "Utils.hpp" +#import #import #import @@ -18,10 +20,19 @@ #error "ARC is off" #endif +#ifndef BUILT_WITH_MACOS13_SDK +#define BUILT_WITH_MACOS13_SDK (MAC_OS_X_VERSION_MAX_ALLOWED >= 130000) +#endif + namespace py = pybind11; using namespace CoreML::Python; +bool usingMacOS13OrHigher() { + // MLProgram class was introduced in macOS 13. + return (NSProtocolFromString(@"MLProgram") != nil); +} + Model::~Model() { @autoreleasepool { NSFileManager *fileManager = [NSFileManager defaultManager]; @@ -60,20 +71,27 @@ throw std::runtime_error(errmsg.str()); } - if (@available(macOS 10.14, *)) { - MLModelConfiguration *configuration = [MLModelConfiguration new]; - if (computeUnits == "CPU_ONLY") { - configuration.computeUnits = MLComputeUnitsCPUOnly; - } else if (computeUnits == "CPU_AND_GPU") { - configuration.computeUnits = MLComputeUnitsCPUAndGPU; + // Set compute unit + MLModelConfiguration *configuration = [MLModelConfiguration new]; + if (computeUnits == "CPU_ONLY") { + configuration.computeUnits = MLComputeUnitsCPUOnly; + } else if (computeUnits == "CPU_AND_GPU") { + configuration.computeUnits = MLComputeUnitsCPUAndGPU; + } else if (computeUnits == "CPU_AND_NE") { + if (usingMacOS13OrHigher()) { +#if BUILT_WITH_MACOS13_SDK + configuration.computeUnits = MLComputeUnitsCPUAndNeuralEngine; +#endif // BUILT_WITH_MACOS13_SDK } else { - assert(computeUnits == "ALL"); - configuration.computeUnits = MLComputeUnitsAll; + throw std::runtime_error("CPU_AND_NE is only available on macOS >= 13.0"); } - m_model = [MLModel modelWithContentsOfURL:compiledUrl configuration:configuration error:&error]; } else { - m_model = [MLModel modelWithContentsOfURL:compiledUrl error:&error]; + assert(computeUnits == "ALL"); + configuration.computeUnits = MLComputeUnitsAll; } + + // Create MLModel + m_model = [MLModel modelWithContentsOfURL:compiledUrl configuration:configuration error:&error]; Utils::handleError(error); } } diff --git a/coremlpython/CoreMLPythonArray.h b/coremlpython/CoreMLPythonArray.h index 5cabfc881..294120721 100644 --- a/coremlpython/CoreMLPythonArray.h +++ b/coremlpython/CoreMLPythonArray.h @@ -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 + #pragma clang diagnostic push #pragma clang diagnostic ignored "-Wexit-time-destructors" #pragma clang diagnostic ignored "-Wdocumentation" diff --git a/coremlpython/CoreMLPythonUtils.mm b/coremlpython/CoreMLPythonUtils.mm index 72f8c3441..9bed610d0 100644 --- a/coremlpython/CoreMLPythonUtils.mm +++ b/coremlpython/CoreMLPythonUtils.mm @@ -7,6 +7,9 @@ #include #include +#include // for std::setfill etc + +#import #if PY_MAJOR_VERSION < 3 @@ -219,6 +222,8 @@ static void handleCVReturn(CVReturn status) { format = kCVPixelFormatType_32BGRA; } else if (formatStr == "L") { format = kCVPixelFormatType_OneComponent8; + } else if (formatStr == "F") { + format = kCVPixelFormatType_OneComponent16Half; } else { std::stringstream msg; msg << "Unsupported image type " << formatStr << ". "; @@ -236,7 +241,6 @@ static void handleCVReturn(CVReturn status) { Py_ssize_t bytesLength = PyBytes_Size(bytesResult.ptr()); assert(bytesLength >= 0); const char *bytesPtr = PyBytes_AsString(bytesResult.ptr()); - std::string bytes(bytesPtr, static_cast(bytesLength)); // copy data into the CVPixelBuffer status = CVPixelBufferLockBaseAddress(pixelBuffer, 0); @@ -245,71 +249,63 @@ static void handleCVReturn(CVReturn status) { assert(baseAddress != nullptr); assert(!CVPixelBufferIsPlanar(pixelBuffer)); size_t bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer); - const char *srcPointer = bytes.data(); + const char *srcPointer = bytesPtr; + + vImage_Buffer srcBuffer; + memset(&srcBuffer, 0, sizeof(srcBuffer)); + srcBuffer.data = const_cast(srcPointer); + srcBuffer.width = width; + srcBuffer.height = height; + vImage_Buffer dstBuffer; + memset(&dstBuffer, 0, sizeof(dstBuffer)); + dstBuffer.data = baseAddress; + dstBuffer.width = width; + dstBuffer.height = height; + if (formatStr == "RGB") { // convert RGB to BGRA - assert(bytes.size() == width * height * 3); - for (size_t row = 0; row < height; row++) { - char *dstPointer = static_cast(baseAddress) + (row * bytesPerRow); - - for (size_t col = 0; col < width; col++) { - - char R = *srcPointer++; - char G = *srcPointer++; - char B = *srcPointer++; - - *dstPointer++ = B; - *dstPointer++ = G; - *dstPointer++ = R; - *dstPointer++ = 0; // A - - } - assert(bytesPerRow >= width * 4); - } - assert(srcPointer == bytes.data() + bytes.size()); + assert(bytesLength == width * height * 3); + + srcBuffer.rowBytes = width * 3; + dstBuffer.rowBytes = bytesPerRow; + vImageConvert_RGB888toBGRA8888(&srcBuffer, NULL, 255, &dstBuffer, false, 0); } else if (formatStr == "RGBA") { // convert RGBA to BGRA - assert(bytes.size() == width * height * 4); - for (size_t row = 0; row < height; row++) { - char *dstPointer = static_cast(baseAddress) + (row * bytesPerRow); - - for (size_t col = 0; col < width; col++) { - - char R = *srcPointer++; - char G = *srcPointer++; - char B = *srcPointer++; - char A = *srcPointer++; - - *dstPointer++ = B; - *dstPointer++ = G; - *dstPointer++ = R; - *dstPointer++ = A; - - } - assert(bytesPerRow >= width * 4); - } - assert(srcPointer == bytes.data() + bytes.size()); + assert(bytesLength == width * height * 4); + srcBuffer.rowBytes = width * 4; + dstBuffer.rowBytes = bytesPerRow; + uint8_t permuteMap[4] = { 2, 1, 0, 3 }; + vImagePermuteChannels_ARGB8888(&srcBuffer, &dstBuffer, permuteMap, 0); - } else { + } else if (formatStr == "L") { - // assume 8 bit grayscale (the only other case) - assert(formatStr == "L"); - assert(bytes.size() == width * height); + // 8 bit grayscale. + assert(bytesLength == width * height); + + srcBuffer.rowBytes = width; + dstBuffer.rowBytes = bytesPerRow; + vImageCopyBuffer(&srcBuffer, &dstBuffer, 1, 0); + + } else if (formatStr == "F") { - for (size_t row = 0; row < height; row++) { - char *dstPointer = static_cast(baseAddress) + (row * bytesPerRow); - - std::memcpy(dstPointer, srcPointer, width); - srcPointer += width; - } + // convert Float32 to Float16. + assert(bytesLength == width * height * sizeof(Float32)); + + srcBuffer.rowBytes = width * sizeof(Float32); + dstBuffer.rowBytes = bytesPerRow; + vImageConvert_PlanarFtoPlanar16F(&srcBuffer, &dstBuffer, 0); + + } else { + std::stringstream msg; + msg << "Unsupported image type " << formatStr << ". "; + msg << "Supported types are: RGB, RGBA, L."; + throw std::runtime_error(msg.str()); } - assert(srcPointer == bytes.data() + bytes.size()); - #ifdef COREML_SHOW_PIL_IMAGES if (formatStr == "RGB") { // for debugging purposes, convert back to PIL image and show it @@ -509,40 +505,70 @@ static size_t sizeOfArrayElement(MLMultiArrayDataType type) { // supports grayscale and BGRA format types auto formatType = CVPixelBufferGetPixelFormatType(value); - assert(formatType == kCVPixelFormatType_32BGRA || formatType == kCVPixelFormatType_OneComponent8); - + assert(formatType == kCVPixelFormatType_32BGRA + || formatType == kCVPixelFormatType_OneComponent8 + || formatType == kCVPixelFormatType_OneComponent16Half); + + auto height = CVPixelBufferGetHeight(value); + auto width = CVPixelBufferGetWidth(value); + + py::str mode; + size_t dstBytesPerRow = 0; + if (formatType == kCVPixelFormatType_32BGRA) { + mode = "RGBA"; + dstBytesPerRow = width * 4; + } else if (formatType == kCVPixelFormatType_OneComponent8) { + mode = "L"; + dstBytesPerRow = width * sizeof(uint8_t); + } else if (formatType == kCVPixelFormatType_OneComponent16Half) { + mode = "F"; + dstBytesPerRow = width * sizeof(Float32); + } else { + std::stringstream msg; + msg << "Unsupported pixel format type: " << std::hex << std::setfill('0') << std::setw(4) << formatType << ". "; + throw std::runtime_error(msg.str()); + } + + PyObject *dstPyBytes = PyBytes_FromStringAndSize(NULL, height * dstBytesPerRow); + if (!dstPyBytes) { + throw std::bad_alloc(); + } + auto result = CVPixelBufferLockBaseAddress(value, kCVPixelBufferLock_ReadOnly); assert(result == kCVReturnSuccess); uint8_t *src = reinterpret_cast(CVPixelBufferGetBaseAddress(value)); assert(src != nullptr); - auto height = CVPixelBufferGetHeight(value); - auto width = CVPixelBufferGetWidth(value); size_t srcBytesPerRow = CVPixelBufferGetBytesPerRow(value); - // Initializing this for Xcode warnings - size_t dstBytesPerRow = 0; - py::str mode; + + // Prepare for vImage blitting + vImage_Buffer srcBuffer; + memset(&srcBuffer, 0, sizeof(srcBuffer)); + srcBuffer.data = src; + srcBuffer.width = width; + srcBuffer.height = height; + srcBuffer.rowBytes = srcBytesPerRow; + + vImage_Buffer dstBuffer; + memset(&dstBuffer, 0, sizeof(dstBuffer)); + dstBuffer.data = PyBytes_AS_STRING(dstPyBytes); + dstBuffer.width = width; + dstBuffer.height = height; + dstBuffer.rowBytes = dstBytesPerRow; + if (formatType == kCVPixelFormatType_32BGRA) { - dstBytesPerRow = width * 4; - mode = "RGBA"; + // convert BGRA to RGBA + uint8_t permuteMap[4] = { 2, 1, 0, 3 }; + vImagePermuteChannels_ARGB8888(&srcBuffer, &dstBuffer, permuteMap, 0); } else if (formatType == kCVPixelFormatType_OneComponent8) { - dstBytesPerRow = width; - mode = "L"; - } - std::string array(height * dstBytesPerRow, 0); - for (size_t i=0; i(src[(i * srcBytesPerRow) + (j * 4) + 2]); - array[(i * dstBytesPerRow) + (j * 4) + 1] = static_cast(src[(i * srcBytesPerRow) + (j * 4) + 1]); - array[(i * dstBytesPerRow) + (j * 4) + 2] = static_cast(src[(i * srcBytesPerRow) + (j * 4) + 0]); - array[(i * dstBytesPerRow) + (j * 4) + 3] = static_cast(src[(i * srcBytesPerRow) + (j * 4) + 3]); - } else if (formatType == kCVPixelFormatType_OneComponent8) { - array[(i * dstBytesPerRow) + j] = static_cast(src[(i * srcBytesPerRow) + j]); - } - } + vImageCopyBuffer(&srcBuffer, &dstBuffer, 1, 0); + } else if (formatType == kCVPixelFormatType_OneComponent16Half) { + vImageConvert_Planar16FtoPlanarF(&srcBuffer, &dstBuffer, 0); + } else { + std::stringstream msg; + msg << "Unsupported pixel format type: " << std::hex << std::setfill('0') << std::setw(4) << formatType << ". "; + throw std::runtime_error(msg.str()); } result = CVPixelBufferUnlockBaseAddress(value, kCVPixelBufferLock_ReadOnly); @@ -552,7 +578,8 @@ static size_t sizeOfArrayElement(MLMultiArrayDataType type) { py::eval("import PIL.Image", scope); py::object pilImage = py::eval("PIL.Image", scope); py::object frombytes = pilImage.attr("frombytes"); - py::object img = frombytes(mode, py::make_tuple(width, height), py::bytes(array)); + py::bytes dstBytes = py::reinterpret_steal(dstPyBytes); // transfer ownership of `dstPyBytes` to `dstBytes` + py::object img = frombytes(mode, py::make_tuple(width, height), dstBytes); return img; } diff --git a/coremltools/__init__.py b/coremltools/__init__.py index c5b872218..8fa29c75e 100644 --- a/coremltools/__init__.py +++ b/coremltools/__init__.py @@ -70,9 +70,13 @@ class ComputeUnit(_Enum): ALL = 1 # Allows the model to use all compute units available, including the neural engine CPU_AND_GPU = 2 # Allows the model to use both the CPU and GPU, but not the neural engine CPU_ONLY = 3 # Limit the model to only use the CPU + CPU_AND_NE = 4 # Allows the model to use both the CPU and neural engine, but not the GPU. + # Only available on macOS >= 13.0 # A dictionary that maps the CoreML model specification version to the MLProgram/MIL opset string _OPSET = { + _SPECIFICATION_VERSION_IOS_13: "CoreML3", + _SPECIFICATION_VERSION_IOS_14: "CoreML4", _SPECIFICATION_VERSION_IOS_15: "CoreML5", _SPECIFICATION_VERSION_IOS_16: "CoreML6", } diff --git a/coremltools/converters/_converters_entry.py b/coremltools/converters/_converters_entry.py index 1169bc2d2..deddd72f1 100644 --- a/coremltools/converters/_converters_entry.py +++ b/coremltools/converters/_converters_entry.py @@ -92,6 +92,7 @@ def convert( - `HDF5 file path `_ (``.h5``) - `SavedModel `_ directory path - A `concrete function `_ + - A `GraphDef `_ * PyTorch @@ -344,6 +345,8 @@ def skip_real_div_ops(op): - ``coremltools.ComputeUnit.CPU_ONLY``: Limit the model to only use the CPU. - ``coremltools.ComputeUnit.CPU_AND_GPU``: Use both the CPU and GPU, but not the neural engine. + - ``coremltools.ComputeUnit.CPU_AND_NE``: Use both the CPU and neural engine, but + not the GPU. Only available on macOS >= 13.0. package_dir : str Post conversion, the model is saved at a temporary location and diff --git a/coremltools/converters/mil/_deployment_compatibility.py b/coremltools/converters/mil/_deployment_compatibility.py index 4fdc2d933..8478915c4 100644 --- a/coremltools/converters/mil/_deployment_compatibility.py +++ b/coremltools/converters/mil/_deployment_compatibility.py @@ -40,6 +40,12 @@ class AvailableTarget(IntEnum): tvOS14 = _SPECIFICATION_VERSION_IOS_14 tvOS15 = _SPECIFICATION_VERSION_IOS_15 tvOS16 = _SPECIFICATION_VERSION_IOS_16 + + # customized __str__ + def __str__(self): + original_str = super().__str__() + new_str = original_str.replace(type(self).__name__, "coremltools.target") + return new_str _get_features_associated_with = {} @@ -133,7 +139,7 @@ def check_deployment_compatibility(spec, representation, deployment_target): for any_target in AvailableTarget: - if any_target.value > deployment_target.value and any_target in _get_features_associated_with: + if any_target > deployment_target and any_target in _get_features_associated_with: missing_features = _get_features_associated_with[any_target](spec) if missing_features: diff --git a/coremltools/converters/mil/backend/mil/load.py b/coremltools/converters/mil/backend/mil/load.py index faadf2168..61a2abe46 100644 --- a/coremltools/converters/mil/backend/mil/load.py +++ b/coremltools/converters/mil/backend/mil/load.py @@ -281,7 +281,7 @@ def load(prog, weights_dir, resume_on_errors=False, specification_version=_SPECI if "main" not in prog.functions: raise ValueError("main function not found in program") - mil_passes.mil_backend_passes(prog, specification_version) + mil_passes.mil_backend_passes(prog) # if user has specified "ClassifierConfig", then add the "classify" op to the prog classifier_config = kwargs.get("classifier_config", None) 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 3a584c2f0..1f328ba6a 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,10 +5,11 @@ import logging +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target from coremltools.converters.mil.mil import Builder as mb, types as types -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._deployment_compatibility import AvailableTarget as target +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass @register_pass(namespace="mil_backend") @@ -24,7 +25,7 @@ class adjust_io_to_supported_types(AbstractGraphPass): types are numerical and can be reasonably replaced with 32 bit float types. The "main" function has additional rules since its I/O is mapped to CoreML model I/O: - * if minimum_deployment_target < coremltools.target.iOS16, then: + * if function.opset_version < coremltools.target.iOS16, then: * Fp16 I/O is replaced with fp32 I/O. Casts (fp32 input -> fp16) are inserted at the beginning of the program to preserve 16 bit inputs. Casts (fp16 -> fp32 output) are inserted at the end of the program to preserve 16 bit computations. @@ -63,7 +64,7 @@ class adjust_io_to_supported_types(AbstractGraphPass): 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, self.minimun_deployment_target) + _adjust_io_to_supported_types(func, is_main_funtion) __RUNTIME_SUPPORTED_TYPES = [types.fp16, types.fp32, types.int32, types.str, types.bool] @@ -76,7 +77,8 @@ def _adjust_var_dtype_helper(var, dtype): else: var._sym_type = types.tensor(dtype, var.sym_type.get_shape()) -def _adjust_main_inputs(func, min_deployment_target): +@block_context_manager +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)) \ @@ -97,13 +99,13 @@ def _adjust_main_inputs(func, min_deployment_target): "of fp32. No cast will be inserted; the previous dtype will be replaced.") _adjust_var_dtype_helper(input_var, types.fp32) elif input_var.dtype == types.fp16 \ - and min_deployment_target.value>= target.iOS16.value: + and func.opset_version >= target.iOS16: pass # do nothing, since fp16 is a valid input type for CoreML 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 # and do not have the same limitation on input types. - supported_dtypes = "{int32, fp32, fp64}" if min_deployment_target < target.iOS16 else \ + supported_dtypes = "{int32, fp32, fp64}" if func.opset_version < target.iOS16 else \ "{int32, fp16, fp32, fp64}" msg = "\nInput '{}' is of dtype {}. The " +\ "CoreML runtime does not support inputs with this dtype " +\ @@ -111,29 +113,29 @@ def _adjust_main_inputs(func, min_deployment_target): "fp32. A cast will be inserted at the beginning of the program to " +\ "convert the input to the originally defined dtype.\n" if input_var.dtype == types.fp16: - msg += "fp16 dtype input is supported if the minimum_deployment_target is chosen to be at least " \ + msg += "fp16 dtype input is supported if the function.opset_version is chosen to be at least " \ "iOS16/macOS13.\n" logging.warning(msg.format( input_var.name, input_dtype_str, supported_dtypes)) - with func: - 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) + 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) -def _adjust_main_outputs(func, min_deployment_target): +@block_context_manager +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 \ - and (min_deployment_target < target.iOS16 or output_var.dtype != types.fp16): + and (func.opset_version < target.iOS16 or output_var.dtype != types.fp16): # since fp16 is a valid output type for coreml from ios16 spec onwards, no need to cast output_dtype_str = types.builtin_to_string(output_var.dtype) - supported_dtypes = "{int32, fp32, fp64}" if min_deployment_target < target.iOS16 else \ + supported_dtypes = "{int32, fp32, fp64}" if func.opset_version < target.iOS16 else \ "{int32, fp16, fp32, fp64}" msg = "\nOutput '{}' is of dtype {}. The " +\ "CoreML runtime does not support outputs with this dtype " +\ @@ -141,7 +143,7 @@ def _adjust_main_outputs(func, min_deployment_target): "of fp32. A cast will be inserted at the end of the program to convert" +\ "the original output dtype to the dtype supported by the CoreML runtime.\n" if output_var.dtype == types.fp16: - msg += "fp16 dtype output is supported if the minimum_deployment_target is chosen to be at least " \ + msg += "fp16 dtype output is supported if function.opset_version is chosen to be at least " \ "iOS16/macOS13.\n" logging.warning(msg.format( output_var.name, @@ -152,9 +154,8 @@ def _adjust_main_outputs(func, min_deployment_target): output_var_name = output_var.name 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.set_name(output_var_name) + 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) @@ -195,6 +196,7 @@ def _adjust_block_inputs(block): for input_var in block.inputs: _adjust_var(input_var) +@block_context_manager def _adjust_ops(block): len_block = len(block.operations) i = 0 @@ -241,11 +243,10 @@ 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) - block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=new_cast_out - ) + 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 + ) block.remove_ops([op]) len_block = len(block.operations) i = i + 1 @@ -254,11 +255,11 @@ def _adjust_ops(block): ##### # The Pass ##### -def _adjust_io_to_supported_types(func, is_main, min_deployment_target): +def _adjust_io_to_supported_types(func, is_main): if is_main: - _adjust_main_inputs(func, min_deployment_target) + _adjust_main_inputs(func) _adjust_ops(func) - _adjust_main_outputs(func, min_deployment_target) + _adjust_main_outputs(func) else: _adjust_func_inputs(func) _adjust_ops(func) 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 2ce5826ee..84b7f99eb 100644 --- a/coremltools/converters/mil/backend/mil/passes/fuse_activation_silu.py +++ b/coremltools/converters/mil/backend/mil/passes/fuse_activation_silu.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.passes.helper import block_context_manager from coremltools.converters.mil.mil import Builder as mb def _match_pattern(op): @@ -39,6 +40,7 @@ def _try_to_transform(sigmoid_op, mul_op, block): return True +@block_context_manager def _fuse_activation_silu_block(block): fusion_status = False for op in list(block.operations): @@ -51,8 +53,7 @@ def _fuse_activation_silu_block(block): 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 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 f015d9d00..48a8ea8e1 100644 --- a/coremltools/converters/mil/backend/mil/passes/homogenize_input_dtypes.py +++ b/coremltools/converters/mil/backend/mil/passes/homogenize_input_dtypes.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 import Builder as mb -from coremltools.converters.mil.mil.ops.defs import elementwise_binary, matmul -from coremltools.converters.mil.mil.ops.defs.elementwise_unary import cast as cast_op_class +from coremltools.converters.mil.mil.ops.defs.iOS15 import elementwise_binary, matmul +from coremltools.converters.mil.mil.ops.defs.iOS15.elementwise_unary import cast as cast_op_class from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil.types import promote_dtypes, builtin_to_string @@ -37,6 +38,7 @@ def _promoted_var(op, var, promoted_dtype): x = mb.const(val=const_value_after_cast, name=var.name + "_promoted", before_op=op) return x +@block_context_manager def _homogenize_input_dtypes_block(block): for op in list(block.operations): for b in op.blocks: @@ -51,8 +53,7 @@ def _homogenize_input_dtypes_block(block): for i,var in enumerate(input_vars): if not _is_same_dtype(var.dtype, promoted_dtype): has_mixed_dtypes = True - with block: - input_vars[i] = _promoted_var(op, var, promoted_dtype) + input_vars[i] = _promoted_var(op, var, promoted_dtype) if has_mixed_dtypes: new_inputs = dict(zip(params, input_vars)) @@ -60,17 +61,16 @@ def _homogenize_input_dtypes_block(block): new_inputs.update( {k: v for k, v in op.inputs.items() if k not in new_inputs} ) - with block: - # create a new op with the promoted input vars - new_op_class = getattr(mb,op.op_type) - new_output = new_op_class(**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 - # but for output dtype assumes them to be same and chooses one of the two. - ) - block.remove_ops([op]) + 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 + # but for output dtype assumes them to be same and chooses one of the two. + ) + block.remove_ops([op]) @register_pass(namespace="mil_backend") class homogenize_input_dtypes(AbstractGraphPass): 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 53e1fd956..42556e319 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 @@ -3,9 +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 -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.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.input_types import ColorLayout, ImageType + # import mil internal ops to add it to the builder from coremltools.converters.mil.mil.ops import defs as _ops from coremltools.converters.mil.mil import Builder as mb @@ -23,7 +25,7 @@ def apply(self, prog): if f_name == 'main': _insert_image_preprocessing_ops(f, prog) - +@block_context_manager def _insert_image_preprocessing_ops(block, prog): input_types = list(prog.main_input_types) @@ -37,29 +39,28 @@ def _insert_image_preprocessing_ops(block, prog): first_op = block.operations[0] old_var = placeholder_op.outputs[0] has_bias = np.any(np.array(input_type.bias) != 0) - with block: - last_output = input_var - input_nptype = nptype_from_builtin(type(last_output.dtype())) - if input_type.scale != 1: - last_output = mb.mul(x=last_output, - y=np.array(input_type.scale, dtype=input_nptype), - before_op=first_op, name=input_var.name + "__scaled__") - if has_bias: - if input_type.color_layout in (ColorLayout.GRAYSCALE, ColorLayout.GRAYSCALE_FLOAT16): + last_output = input_var + input_nptype = nptype_from_builtin(type(last_output.dtype())) + if input_type.scale != 1: + last_output = mb.mul(x=last_output, + y=np.array(input_type.scale, dtype=input_nptype), + before_op=first_op, name=input_var.name + "__scaled__") + if has_bias: + if input_type.color_layout in (ColorLayout.GRAYSCALE, ColorLayout.GRAYSCALE_FLOAT16): + last_output = mb.add(x=last_output, + y=np.array(input_type.bias, dtype=input_nptype), + before_op=first_op, name=input_var.name + "__biased__") + else: + if len(last_output.shape) == 3: + last_output = mb.add(x=last_output, + y=np.array(input_type.bias, dtype=input_nptype).reshape([3, 1, 1]), + before_op=first_op, name=input_var.name + "__biased__") + elif len(last_output.shape) == 4: last_output = mb.add(x=last_output, - y=np.array(input_type.bias, dtype=input_nptype), + y=np.array(input_type.bias, dtype=input_nptype).reshape([1, 3, 1, 1]), before_op=first_op, name=input_var.name + "__biased__") else: - if len(last_output.shape) == 3: - last_output = mb.add(x=last_output, - y=np.array(input_type.bias, dtype=input_nptype).reshape([3, 1, 1]), - before_op=first_op, name=input_var.name + "__biased__") - elif len(last_output.shape) == 4: - last_output = mb.add(x=last_output, - y=np.array(input_type.bias, dtype=input_nptype).reshape([1, 3, 1, 1]), - before_op=first_op, name=input_var.name + "__biased__") - else: - raise TypeError("Unsupported rank for image input type.") + raise TypeError("Unsupported rank for image input type.") if last_output != input_var: block.replace_uses_of_var_after_op(anchor_op=last_output.op, diff --git a/coremltools/converters/mil/backend/mil/passes/mil_passes.py b/coremltools/converters/mil/backend/mil/passes/mil_passes.py index a7ef65a9c..10d197ef5 100644 --- a/coremltools/converters/mil/backend/mil/passes/mil_passes.py +++ b/coremltools/converters/mil/backend/mil/passes/mil_passes.py @@ -10,8 +10,7 @@ from coremltools.converters.mil._deployment_compatibility import AvailableTarget -def mil_backend_passes(prog, minimum_spec_version): - min_deployment_target = AvailableTarget(minimum_spec_version) +def mil_backend_passes(prog): passes = [ "common::const_elimination", "mil_backend::adjust_io_to_supported_types", @@ -36,7 +35,6 @@ def mil_backend_passes(prog, minimum_spec_version): prog.validate() for p in passes: _logging.info('Performing passes for mil backend: "{}"'.format(p)) - PASS_REGISTRY[p].minimun_deployment_target = min_deployment_target PASS_REGISTRY[p](prog) # No more validation from this point on as prog is not SSA anymore. diff --git a/coremltools/converters/mil/backend/mil/passes/test_passes.py b/coremltools/converters/mil/backend/mil/passes/test_passes.py index f6f34cebc..452e1697a 100644 --- a/coremltools/converters/mil/backend/mil/passes/test_passes.py +++ b/coremltools/converters/mil/backend/mil/passes/test_passes.py @@ -151,10 +151,10 @@ def prog(x): @pytest.mark.parametrize( - "use_ios16_deployment_target", - [False, True], + "opset_version", + [None, target.iOS13, target.iOS16], ) - def test_float16_input_output(self, use_ios16_deployment_target): + def test_float16_input_output(self, opset_version): """ Input graph: @@ -164,7 +164,7 @@ def test_float16_input_output(self, use_ios16_deployment_target): } -> (%relu_0) } - Output graph (if deployment_target < ios16): + Output graph (if opset_version < ios16): main(%x: (1, 1, 1, 1, fp32)(Tensor)) { block0() { @@ -174,15 +174,12 @@ def test_float16_input_output(self, use_ios16_deployment_target): } -> (%relu_0) } - Output graph (if deployment_target >= ios16): same as the input graph + Output graph (if opset_version >= ios16): same as the input graph """ - @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.fp16)]) + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 1, 1), dtype=types.fp16)], opset_version=opset_version) def prog(x): return mb.relu(x=x) - if use_ios16_deployment_target: - PASS_REGISTRY["mil_backend::adjust_io_to_supported_types"].minimun_deployment_target = target.iOS16 - prev_prog, prev_block, block = apply_pass_and_basic_check( prog, "mil_backend::adjust_io_to_supported_types" ) @@ -193,7 +190,7 @@ def prog(x): outputs = block.outputs assert prev_inputs[0][1].name == inputs[0][1].name assert outputs[0].name == prev_outputs[0].name - if not use_ios16_deployment_target: + if opset_version is None or opset_version < target.iOS16: assert get_op_types_in_program(prog) == ['cast', 'relu', 'cast'] assert inputs[0][1].dtype == types.fp32 assert outputs[0].dtype == types.fp32 @@ -201,7 +198,37 @@ def prog(x): assert get_op_types_in_program(prog) == ['relu'] assert inputs[0][1].dtype == types.fp16 assert block.outputs[0].dtype == types.fp16 + + def test_float16_input_output_with_opset_version_inference(self): + """ + Input graph: + main(%x: (1, 1, 4, 4, fp16)(Tensor)) { + block0() { + %pixel_unshuffle_0: (1, 4, 2, 2, fp16)(Tensor) = pixel_unshuffle(x=%x, downscale_factor=2, name="pixel_unshuffle_0") + } -> (%pixel_unshuffle_0) + } + + This function would be inferred as an iOS16 function, and the graph pass should behave properly + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4), dtype=types.fp16)]) + def prog(x): + x = mb.pixel_unshuffle(x=x, downscale_factor=np.uint32(2)) + return x + + prev_prog, prev_block, block = apply_pass_and_basic_check( + prog, "mil_backend::adjust_io_to_supported_types" + ) + + prev_inputs = list(prev_block.inputs.items()) + inputs = list(block.inputs.items()) + prev_outputs = prev_block.outputs + outputs = block.outputs + assert prev_inputs[0][1].name == inputs[0][1].name + assert outputs[0].name == prev_outputs[0].name + assert get_op_types_in_program(prog) == ['pixel_unshuffle'] + assert inputs[0][1].dtype == types.fp16 + assert block.outputs[0].dtype == types.fp16 def test_int8_input(self): """ @@ -910,9 +937,7 @@ def prog(x): return out assert get_op_types_in_program(prog) == ["mul"] - print(prog) apply_pass_and_basic_check(prog, "mil_backend::homogenize_input_dtypes") - print(prog) # verify that there is no cast op in the program, since the # const input (int32) should have been promoted to a float32 and replaced with a new const assert get_op_types_in_program(prog) == ["mul"] diff --git a/coremltools/converters/mil/backend/nn/op_mapping.py b/coremltools/converters/mil/backend/nn/op_mapping.py index d642d5c36..8cb96109d 100644 --- a/coremltools/converters/mil/backend/nn/op_mapping.py +++ b/coremltools/converters/mil/backend/nn/op_mapping.py @@ -183,14 +183,17 @@ def _try_convert_global_pool(const_context, builder, op, mode): if keep_dims is False: return False - if op.axes is not None: + axes = None + if op.axes is not None and op.axes.val is not None: axes = op.axes.val - axes = sorted([rank + axis if axis < 0 else axis for axis in axes]) + else: + axes = list(range(rank)) + + if tuple(op.outputs[0].shape[:-2]) != tuple(op.inputs["x"].shape[:-2]): + return False + if not all([s == 1 for s in op.outputs[0].shape[-2:]]): + return False - if tuple(op.outputs[0].shape[:-2]) != tuple(op.inputs["x"].shape[:-2]): - return False - if not all([s == 1 for s in op.outputs[0].shape[-2:]]): - return False builder.add_pooling( name=op.name, height=0, diff --git a/coremltools/converters/mil/backend/nn/passes/test_passes.py b/coremltools/converters/mil/backend/nn/passes/test_passes.py index ae4753e07..a8e120803 100644 --- a/coremltools/converters/mil/backend/nn/passes/test_passes.py +++ b/coremltools/converters/mil/backend/nn/passes/test_passes.py @@ -32,8 +32,8 @@ def prog(a, b): return mb.while_loop(_cond=cond, _body=body, loop_vars=(a, b)) while_op = prog.find_ops(op_type="while_loop", exactly_one=True)[0] - assert while_op.blocks[0].inputs[0].name == "a_x1" - assert while_op.blocks[0].inputs[1].name == "b_x1" + assert while_op.blocks[0].inputs[0].name == "a_x0" + assert while_op.blocks[0].inputs[1].name == "b_x0" prev_prog = copy.deepcopy(prog) PASS_REGISTRY["nn_backend::commingle_loop_vars"](prog) diff --git a/coremltools/converters/mil/converter.py b/coremltools/converters/mil/converter.py index a4ec005a9..dc12e708b 100644 --- a/coremltools/converters/mil/converter.py +++ b/coremltools/converters/mil/converter.py @@ -48,6 +48,17 @@ class MILFrontend: name = "milinternal" def __call__(self, model, *args, **kwargs): + specification_version = kwargs.get("specification_version", None) + if specification_version is not None: + max_opset_version, op = model._get_max_opset_version_and_op() + if max_opset_version > specification_version: + msg = ( + "Please update the minimum_deployment_target to {!s}," + " since op {} is only available in opset {!s} or newer." + + ).format(max_opset_version, op.op_type, max_opset_version) + raise ValueError(msg) + if "inputs" in kwargs and kwargs["inputs"] is not None: inputs = kwargs["inputs"] if not isinstance(inputs, (list, tuple)): diff --git a/coremltools/converters/mil/experimental/passes/generic_pass_infrastructure.py b/coremltools/converters/mil/experimental/passes/generic_pass_infrastructure.py index 63ed80d99..b8ea6e51a 100644 --- a/coremltools/converters/mil/experimental/passes/generic_pass_infrastructure.py +++ b/coremltools/converters/mil/experimental/passes/generic_pass_infrastructure.py @@ -7,6 +7,7 @@ import itertools from ...mil.passes import pass_registry +from coremltools.converters.mil.mil.passes.helper import block_context_manager # IMPORTANT: List of assumptions we are making about the problem # 1) The user defined pattern has exactly one root variable, and one final output operation. As such, we will be searching for a singlular @@ -163,7 +164,7 @@ def _detect_pattern(program_op, ops_arrangement_root_var, block): return False, None - +@block_context_manager def _fuse_one_block(block, ops_arrangement, var_constraints, transform_pattern): fusion_status = False for op in list(block.operations): @@ -172,16 +173,15 @@ def _fuse_one_block(block, ops_arrangement, var_constraints, transform_pattern): while block_changed: block_changed = _fuse_one_block(b, ops_arrangement, var_constraints, transform_pattern) - with block: - ops_arrangement_root_var = list(ops_arrangement.functions.values())[0].function_inputs[0] - fusion_status, pattern = _detect_pattern(op, ops_arrangement_root_var, block) + ops_arrangement_root_var = list(ops_arrangement.functions.values())[0].function_inputs[0] + fusion_status, pattern = _detect_pattern(op, ops_arrangement_root_var, block) - if fusion_status: - fusion_status &= var_constraints(pattern) + if fusion_status: + fusion_status &= var_constraints(pattern) - if fusion_status: - transform_pattern(pattern) - return fusion_status + if fusion_status: + transform_pattern(pattern) + return fusion_status return fusion_status diff --git a/coremltools/converters/mil/experimental/passes/readme.md b/coremltools/converters/mil/experimental/passes/readme.md index bd28ffb70..aff706f80 100644 --- a/coremltools/converters/mil/experimental/passes/readme.md +++ b/coremltools/converters/mil/experimental/passes/readme.md @@ -34,7 +34,6 @@ def prog(x): x = mb.reduce_mean(x=x, axes=[2, 3], keep_dims=False, name='reduce') x = mb.log(x=x, name='log') return x -print(prog) ``` * It is important that the user follows these constraints when writing their MIL program: diff --git a/coremltools/converters/mil/frontend/milproto/helper.py b/coremltools/converters/mil/frontend/milproto/helper.py index 85a05413e..c94bd1030 100644 --- a/coremltools/converters/mil/frontend/milproto/helper.py +++ b/coremltools/converters/mil/frontend/milproto/helper.py @@ -9,10 +9,6 @@ from coremltools.converters.mil.mil.program import get_new_symbol -def opstr_to_opcls(op_str): - return getattr(sys.modules["coremltools.converters.mil.mil.ops.defs"], op_str) - - def get_proto_dim(dim): if dim.WhichOneof("dimension") == "constant": return dim.constant.size diff --git a/coremltools/converters/mil/frontend/milproto/load.py b/coremltools/converters/mil/frontend/milproto/load.py index 8c7964c88..cc8d572ed 100644 --- a/coremltools/converters/mil/frontend/milproto/load.py +++ b/coremltools/converters/mil/frontend/milproto/load.py @@ -9,6 +9,7 @@ import numpy as np from coremltools import _OPSET +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as _target from coremltools.converters.mil.mil import ( Block, Builder as mb, @@ -21,10 +22,14 @@ types, Var, ) -from coremltools.converters.mil.mil.block import curr_block +from coremltools.converters.mil.mil.block import curr_block, curr_opset_version +from coremltools.converters.mil.mil.ops.registry import SSAOpRegistry as _SSAOpRegistry from coremltools.libmilstoragepython import _BlobStorageReader as BlobReader -from coremltools.proto import MIL_pb2 as pm -from .helper import proto_to_types, opstr_to_opcls +from coremltools.proto import ( + MIL_pb2 as pm, + Model_pb2 as ml +) +from .helper import proto_to_types class TranscriptionContext: @@ -304,10 +309,7 @@ def _load_operation(context, op_spec): vars.append(var) else: raise NotImplementedError("Binding {} not yet implemented".format(binding_type)) - - # TODO: rdar://92930138 (Milproto -> Pymil op translation should take account of the op version) - # we need to use the spec version of the function to pick up the correct version of op - op_cls = opstr_to_opcls(op_type) + op_cls = _SSAOpRegistry._get_core_op_cls(op_type) if len(vars) == 1 and not isinstance( op_cls.input_spec.input_types[param_name], TupleInputType ): @@ -368,12 +370,6 @@ def _load_function(context, func_spec, spec_version): if not isinstance(func_spec, pm.Function): raise TypeError("Invalid Function spec object") - opset = func_spec.opset - if opset != _OPSET[spec_version]: - raise AssertionError( - "Mismatch between provide specification version vs version implied by opset field" - ) - if func_spec.attributes: raise ValueError("Attributes on functions not supported") @@ -388,17 +384,28 @@ def _load_function(context, func_spec, spec_version): sym_shape=valuetype.get_shape(), dtype=valuetype.get_primitive(), name=name ) context.register_var_with_name(name, func_inputs[name].outputs[0]) - + + opset = func_spec.opset if opset not in func_spec.block_specializations: raise ValueError("Missing block specialization for opset {}".format(opset)) - with Function(func_inputs) as pymil_func: + with Function(func_inputs, opset_version=_target(spec_version)) as pymil_func: _load_block(context, func_spec.block_specializations[opset]) return pymil_func -def load(program_spec, specification_version, file_weights_dir="", **kwargs): +def load(model_spec, specification_version, file_weights_dir="", **kwargs): + if not isinstance(model_spec, ml.Model): + raise TypeError("Invalid Model sepc object") + + if specification_version < model_spec.specificationVersion: + raise ValueError("specification_version must be greater or equal to the input model spec version") + + if model_spec.WhichOneof("Type") != "mlProgram": + raise ValueError("Only MIL proto based mlmodels can be loaded") + + program_spec = model_spec.mlProgram if not isinstance(program_spec, pm.Program): raise TypeError("Invalid Program spec object") diff --git a/coremltools/converters/mil/frontend/milproto/test_load.py b/coremltools/converters/mil/frontend/milproto/test_load.py index 282866e43..46dbfddb7 100644 --- a/coremltools/converters/mil/frontend/milproto/test_load.py +++ b/coremltools/converters/mil/frontend/milproto/test_load.py @@ -9,6 +9,7 @@ import coremltools as ct from coremltools import ComputeUnit +from coremltools.converters.mil import Builder as mb from coremltools.converters.mil.converter import mil_convert from coremltools.converters.mil.frontend.milproto.load import load as milproto_to_pymil from coremltools.converters.mil.frontend.torch.test.test_torch_ops import TestScriptedModels as _TestScriptedModels @@ -18,40 +19,46 @@ from coremltools.converters.mil.testing_utils import get_op_types_in_program from coremltools.converters._converters_entry import _get_metadata_from_mlmodel - -def roundtrip_and_compare_mlmodel(mlmodel, input_dict): +def get_pymil_prog_from_mlmodel(mlmodel): model_spec = mlmodel.get_spec() - if model_spec.WhichOneof("Type") != "mlProgram": - raise ValueError("Only MIL proto based mlmodels can be loaded") - - program_spec = model_spec.mlProgram - model_description = model_spec.description - - pymil_prog = milproto_to_pymil( - program_spec=program_spec, + return milproto_to_pymil( + model_spec=model_spec, specification_version=model_spec.specificationVersion, file_weights_dir=mlmodel.weights_dir, - ) + ) + +def get_roundtrip_mlmodel(mlmodel): + """ + This utility function does the following roundtrip conversion: + + mlprogram proto -> pymil program -> mlprogram model + """ + pymil_prog = get_pymil_prog_from_mlmodel(mlmodel) + + # convert the pymil program to mlmodel + model_spec = mlmodel.get_spec() roundtripped_mlmodel = mil_convert( pymil_prog, convert_to="mlprogram", convert_from="milinternal", - compute_units=ComputeUnit.ALL, - model_description=model_description, + compute_units=mlmodel.compute_unit, + model_description=model_spec.description, + specification_version=model_spec.specificationVersion, ) # set MIL program attributes build_info = _get_metadata_from_mlmodel(mlmodel) roundtripped_mlmodel._set_build_info_mil_attributes(build_info) + return roundtripped_mlmodel +def roundtrip_and_compare_mlmodel(mlmodel, input_dict): + roundtripped_mlmodel = get_roundtrip_mlmodel(mlmodel) expected_outputs = mlmodel.predict(input_dict) compare_backend(roundtripped_mlmodel, input_dict, expected_outputs) class TestLoadAPIUsage: def test_mil_proto_to_pymil(self): - from coremltools.converters.mil import Builder as mb - # Define a PyMIL program @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3, 100, 100)), ]) def prog(x): @@ -65,26 +72,39 @@ def prog(x): return x # Convert it to MIL proto backed MLModel - mlmodel = mil_convert( - prog, - convert_to="mlprogram", - convert_from="milinternal", - compute_units=ComputeUnit.ALL, - ) + mlmodel = ct.convert(prog, convert_to="mlprogram") # Load MLModel back to PyMIL - model_spec = mlmodel.get_spec() - program_spec = model_spec.mlProgram - loaded_pymil_prog = milproto_to_pymil( - program_spec=program_spec, - specification_version=model_spec.specificationVersion, - file_weights_dir=mlmodel.weights_dir, - ) + loaded_pymil_prog = get_pymil_prog_from_mlmodel(mlmodel) # Assert that loaded PyMIL prog matches with defined PyMIL prog if get_op_types_in_program(loaded_pymil_prog) != get_op_types_in_program(prog): raise AssertionError("Mismatch between defined PyMIL prog and loaded PyMIL prog") + def test_mil_proto_to_pymil_with_version_handling(self): + # This test makes sure the correct version of the op is picked up during mil_proto -> pymil conversion + + # iOS15 version program with iOS13 version topk + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=ct.target.iOS15) + def prog(x): + x = mb.topk(x=x, k=1, axis=-1, ascending=True) + return x + + iOS15_mlmodel = ct.convert(prog, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS15) + iOS15_pymil_prog = get_pymil_prog_from_mlmodel(iOS15_mlmodel) + topk_op = iOS15_pymil_prog.functions["main"].find_ops(op_type="topk")[0] + assert not hasattr(topk_op, "sort") + + # iOS16 version program with iOS16 version topk + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=ct.target.iOS16) + def prog(x): + x = mb.topk(x=x, k=1, axis=-1, ascending=True) + return x + + iOS16_mlmodel = ct.convert(prog, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS16) + iOS16_pymil_prog = get_pymil_prog_from_mlmodel(iOS16_mlmodel) + topk_op = iOS16_pymil_prog.functions["main"].find_ops(op_type="topk")[0] + assert hasattr(topk_op, "sort") @pytest.mark.skipif(ct.utils._macos_version() < (12, 0), reason="mlprogram predict available only on macOS12+") class TestE2ENumericalCorrectness: diff --git a/coremltools/converters/mil/frontend/tensorflow/converter.py b/coremltools/converters/mil/frontend/tensorflow/converter.py index 448119a10..05e9bfa87 100644 --- a/coremltools/converters/mil/frontend/tensorflow/converter.py +++ b/coremltools/converters/mil/frontend/tensorflow/converter.py @@ -9,6 +9,7 @@ from .convert_utils import convert_graph from .ssa_passes.tf_passes import tensorflow_passes from .._utils import get_output_names +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as _target from coremltools.converters.mil.input_types import ( _get_shaping_class, InputType, @@ -117,7 +118,7 @@ def __contains__(self, tf_name): class TFConverter: - def __init__(self, tfssa, inputs=None, outputs=None, **kwargs): + def __init__(self, tfssa, inputs=None, outputs=None, opset_version=None): """ tfssa: TensorFlow IR. inputs: list of TensorType or ImageType, optional, defaults to None. @@ -129,6 +130,7 @@ def __init__(self, tfssa, inputs=None, outputs=None, **kwargs): self.global_type = {} self.inputs = None self.main_output_types = outputs + self.opset_version = _target(opset_version) if opset_version is not None else None output_names = get_output_names(outputs) main_func = tfssa.functions["main"] @@ -378,7 +380,7 @@ def convert_main_graph(self, prog, graph): input_type.shape.symbolic_shape, dtype=input_type.dtype) prog.set_main_input_types(self.inputs) - with Function(func_inputs) as ssa_func: + with Function(func_inputs, opset_version=self.opset_version) as ssa_func: # Get the input Var for name in func_inputs.keys(): input_var = ssa_func.inputs[name] diff --git a/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py b/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py index eb8302331..782866383 100644 --- a/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/dialect_ops.py @@ -28,7 +28,7 @@ # # tf_make_list allows elem_shape to be unspecified. core op make_list does # not allow that. -@register_op(doc_str="TODO", namespace="tf") +@register_op(namespace="tf") class tf_make_list(Operation): input_spec = InputSpec( init_length=IntInputType(optional=True), @@ -104,25 +104,24 @@ def _check_peephole_weights(self): ) -@register_op( - doc_str=""" - xh = [x, h_prev] - [i, ci, f, o] = xh * w + b - f = f + forget_bias - if not use_peephole: - wci = wcf = wco = 0 - i = sigmoid(cs_prev .* wci + i) - f = sigmoid(cs_prev .* wcf + f) - ci = tanh(ci) - cs = ci .* i + cs_prev .* f - cs = clip(cs, cell_clip) - o = sigmoid(cs * wco + o) - co = tanh(cs) - h = co .* o - """, - namespace="tf", -) +@register_op(namespace="tf") class tf_lstm_block_cell(TfLSTMBase): + """ + xh = [x, h_prev] + [i, ci, f, o] = xh * w + b + f = f + forget_bias + + if not use_peephole: + wci = wcf = wco = 0 + i = sigmoid(cs_prev .* wci + i) + f = sigmoid(cs_prev .* wcf + f) + ci = tanh(ci) + cs = ci .* i + cs_prev .* f + cs = clip(cs, cell_clip) + o = sigmoid(cs * wco + o) + co = tanh(cs) + h = co .* o + """ input_spec = ( InputSpec(x=TensorInputType(),) + TfLSTMBase.input_spec # [batch, input_dim] ) @@ -149,13 +148,11 @@ def type_inference(self): ) # h -@register_op( - doc_str=""" - Apply LSTM to an input sequence - """, - namespace="tf", -) +@register_op(namespace="tf") class tf_lstm_block(TfLSTMBase): + """ + Apply LSTM to an input sequence + """ input_spec = ( InputSpec( seq_len=IntInputType(), # int diff --git a/coremltools/converters/mil/frontend/tensorflow/load.py b/coremltools/converters/mil/frontend/tensorflow/load.py index 3dd38cbec..1b7a5826e 100644 --- a/coremltools/converters/mil/frontend/tensorflow/load.py +++ b/coremltools/converters/mil/frontend/tensorflow/load.py @@ -232,7 +232,12 @@ def _program_from_tf_ssa(self): filename="/tmp/ssa_after_tf_passes", cleanup=True ) - converter = TFConverter(self._tf_ssa, **self.kwargs) + converter = TFConverter( + tfssa=self._tf_ssa, + inputs=self.kwargs["inputs"], + outputs=self.kwargs["outputs"], + opset_version=self.kwargs["specification_version"] + ) return converter.convert() @staticmethod diff --git a/coremltools/converters/mil/frontend/tensorflow/ops.py b/coremltools/converters/mil/frontend/tensorflow/ops.py index 287436046..78ff19782 100644 --- a/coremltools/converters/mil/frontend/tensorflow/ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/ops.py @@ -9,7 +9,9 @@ from .._utils import build_einsum_mil from .convert_utils import convert_graph from .tf_op_registry import register_tf_op +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target from coremltools.converters.mil.mil import Builder as mb, types +from coremltools.converters.mil.mil.block import is_current_opset_version_compatible_with from coremltools.converters.mil.mil.ops.defs._utils import broadcast_shapes from coremltools.converters.mil.mil.types.symbolic import is_symbolic, any_symbolic @@ -2091,10 +2093,22 @@ def Tanh(context, node): @register_tf_op(tf_alias=["TopKV2"]) def TopK(context, node): x = context[node.inputs[0]] - k = context[node.inputs[1]] - x = mb.topk(x=x, k=k.val, axis=-1, name=node.name) - context.add(node.name, x) + k = context[node.inputs[1]].val + sort = node.attr["sorted"] + + kwargs = { + "x": x, + "k": k, + "axis": -1, + "name": node.name + } + + if is_current_opset_version_compatible_with(target.iOS16): + kwargs["sort"] = sort + elif not sort: + raise ValueError("For opset <= iOS16, only sorted=True supported for the topk") + context.add(node.name, mb.topk(**kwargs)) @register_tf_op(tf_alias=["InTopKV2"]) def InTopK(context, node): @@ -2165,10 +2179,7 @@ def Where(context, node): raise NotImplementedError('tf.where with x,y will be supported by ' 'MIL::select in the future') x = context[node.inputs[0]] - # rdar://78409794 (Remove cast in tf Where op lowering after rdar://77514629 - # goes into MIL build) - x_fp32 = mb.cast(x=x, dtype="fp32") - x = mb.non_zero(x=x_fp32, name=node.name) + x = mb.non_zero(x=x, name=node.name) context.add(node.name, x) 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 60a4715bc..dc7d9a63c 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 @@ -5,6 +5,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil.types.symbolic import is_symbolic from coremltools.converters.mil.mil.var import ListVar @@ -25,7 +26,7 @@ def apply(self, prog): for f in prog.functions.values(): _backfill_make_list_elem_type_block(f) - +@block_context_manager def _backfill_make_list_elem_type_block(block): # shallow copy hides changes on f.operations during the loop for op in block.operations: @@ -48,18 +49,17 @@ def _backfill_make_list_elem_type_block(block): ) raise ValueError(msg.format(op.name, op.enclosing_block)) - with block: - # elem_shape can be runtime-detemrined, which cannot be inferred here at this point, - # so we add an internal _const_symbolic node to cover both static and dynamic cases. - elem_shape = [dim.name if is_symbolic(dim) else dim for dim in elem_type.get_shape()] - new_list = mb.make_list( - init_length=op.init_length, - dynamic_length=op.dynamic_length, - elem_shape=tuple(elem_shape), - dtype=op.inputs["dtype"], - before_op=op, - name=op.name, - ) + # elem_shape can be runtime-detemrined, which cannot be inferred here at this point, + # so we add an internal _const_symbolic node to cover both static and dynamic cases. + elem_shape = [dim.name if is_symbolic(dim) else dim for dim in elem_type.get_shape()] + new_list = mb.make_list( + init_length=op.init_length, + dynamic_length=op.dynamic_length, + elem_shape=tuple(elem_shape), + dtype=op.inputs["dtype"], + before_op=op, + name=op.name, + ) block.replace_uses_of_var_after_op( anchor_op=op, old_var=op.outputs[0], new_var=new_list 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 ea3717614..1b601c9c7 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 @@ -9,6 +9,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -37,7 +38,7 @@ def apply(self, prog): for f in prog.functions.values(): _tf_lstm_to_core_lstm_block(f) - +@block_context_manager def _tf_lstm_to_core_lstm_block(block): # shallow copy hides changes on f.operations during the loop for op in block.operations: @@ -109,30 +110,29 @@ def _check_unsupported_outputs(unsupported_outputs): cell_clip = None if op.cell_clip is None else op.cell_clip.val output_sequence = op.op_type == "tf_lstm_block" - + block = op.enclosing_block - with block: - # x: [seq_len, batch, input_dim] - if op.op_type == "tf_lstm_block_cell": - x = mb.expand_dims(x=op.x, axes=[0], before_op=op) - else: # tf_lstm_block - x = op.x - new_h_all, new_h, new_cs = mb.lstm( - x=x, - initial_c=op.c_prev, - initial_h=op.h_prev, - weight_ih=w_ih, - weight_hh=w_hh, - bias=bias, - recurrent_activation="sigmoid", - cell_activation="tanh", - activation="tanh", - peephole=mb_peep, - clip=cell_clip, - output_sequence=output_sequence, - name=op.name, - before_op=op, - ) + # x: [seq_len, batch, input_dim] + if op.op_type == "tf_lstm_block_cell": + x = mb.expand_dims(x=op.x, axes=[0], before_op=op) + else: # tf_lstm_block + x = op.x + new_h_all, new_h, new_cs = mb.lstm( + x=x, + initial_c=op.c_prev, + initial_h=op.h_prev, + weight_ih=w_ih, + weight_hh=w_hh, + bias=bias, + recurrent_activation="sigmoid", + cell_activation="tanh", + activation="tanh", + peephole=mb_peep, + clip=cell_clip, + output_sequence=output_sequence, + name=op.name, + before_op=op, + ) if op.op_type == "tf_lstm_block_cell": block.replace_uses_of_var_after_op(anchor_op=op, old_var=cs, new_var=new_cs) diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_conversion_api.py b/coremltools/converters/mil/frontend/tensorflow/test/test_conversion_api.py index 3b226b856..b981aca45 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_conversion_api.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_conversion_api.py @@ -420,8 +420,7 @@ def test_grayscale_input(self, rank4_input_model, rank3_input_model, rank4_grays assert_spec_input_image_type(mlmodel._spec, expected_feature_type=ft.ImageFeatureType.GRAYSCALE_FLOAT16) assert_prog_input_type(mlmodel._mil_program, expected_dtype_str="fp16") assert_output_dtype(mlmodel, expected_type_str="fp16") - # TODO: uncomment the following when rdar://92239179 is fixed - # verify_prediction(mlmodel) + verify_prediction(mlmodel) def test_color_output(self, rank4_input_model, rank4_input_model_with_channel_first_output): # check that an error is raised if the output shape is not of form (1, 3, H, W) @@ -487,8 +486,7 @@ def test_grayscale_output(self, rank4_grayscale_input_model, rank4_grayscale_inp assert_spec_output_image_type(mlmodel._spec, expected_feature_type=ft.ImageFeatureType.GRAYSCALE_FLOAT16) assert_prog_input_type(mlmodel._mil_program, expected_dtype_str="fp16") assert_prog_output_type(mlmodel._mil_program, expected_dtype_str="fp16") - # TODO: uncomment the following when rdar://92239179 is fixed - # verify_prediction(mlmodel) + verify_prediction(mlmodel) mlmodel = ct.convert(rank4_grayscale_input_model_with_channel_first_output, inputs=[ct.ImageType(color_layout=ct.colorlayout.GRAYSCALE)], @@ -500,8 +498,7 @@ def test_grayscale_output(self, rank4_grayscale_input_model, rank4_grayscale_inp assert_spec_output_image_type(mlmodel._spec, expected_feature_type=ft.ImageFeatureType.GRAYSCALE_FLOAT16) assert_prog_input_type(mlmodel._mil_program, expected_dtype_str="fp32") assert_prog_output_type(mlmodel._mil_program, expected_dtype_str="fp16") - # TODO: uncomment the following when rdar://92239179 is fixed - # verify_prediction(mlmodel) + verify_prediction(mlmodel) def test_linear_model(self, linear_model): 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 bb7b13476..1a73c9a75 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_custom_ops.py @@ -47,7 +47,7 @@ class TestCustomMatMul: # Define SSA Custom Op for Sparse MatMul # This will map to `custom_op` in SSA with binding information # to bind input spec to the custom implementation - @register_op(doc_str="Sparse MatMul Layer", is_custom_op=True) + @register_op(is_custom_op=True) class custom_sparse_matmul(Operation): # Defining input spec for current op input_spec = InputSpec( @@ -120,7 +120,7 @@ def test_tf( self, use_cpu_only, backend, transpose_a, transpose_b, a_is_sparse, b_is_sparse, b_is_const, ): if backend[0] == 'mlprogram': - pytest.xfail("Custom layer not supported with ML Program backend") + pytest.skip("Custom layer not supported with ML Program backend") rank = 2 shape = list(np.random.randint(low=3, high=100, size=1)) * rank @@ -177,7 +177,7 @@ class TestCustomTopK: @pytest.fixture(scope="class") def create_custom_TopK(self): # Defining SSA TopK Op - @register_op(doc_str="Custom TopK Layer", is_custom_op=True) + @register_op(is_custom_op=True) class custom_topk(Operation): input_spec = InputSpec( x=TensorInputType(), @@ -243,7 +243,7 @@ def CustomTopK(context, node): @pytest.mark.usefixtures("create_custom_TopK") def test_tf(self, use_cpu_only, backend, rank, k): if backend[0] == 'mlprogram': - pytest.xfail("Custom layer not supported with ML Program backend") + pytest.skip("Custom layer not supported with ML Program backend") shape = np.random.randint(low=3, high=6, size=rank) with tf.Graph().as_default() as graph: diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_load.py b/coremltools/converters/mil/frontend/tensorflow/test/test_load.py index 445455f60..888b56aa0 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_load.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_load.py @@ -200,7 +200,7 @@ def build_model(x): inputs=[ct.TensorType(shape=(1,))]) assert mlmodel is not None - @pytest.mark.xfail(reason="Rank-0 input is not supported", run=True) + @pytest.mark.skip(reason="Rank-0 input is not supported", run=True) def test_scalar_placeholder_shape(self): x_shape = () # Scalar Placeholder Shape diff --git a/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py b/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py index b16441aaa..f4a2b47a2 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/test_ops.py @@ -12,6 +12,7 @@ import numpy as np import pytest +import coremltools as ct from coremltools import TensorType, RangeDim from coremltools.converters.mil import testing_reqs from coremltools.converters.mil.testing_utils import random_gen @@ -50,7 +51,7 @@ def test( self, use_cpu_only, backend, data_warp_shapes, ): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") data_shape, warp_shape = data_warp_shapes @@ -154,7 +155,7 @@ class TestPlaceholderAsOutput(TensorFlowBaseTest): ) def test(self, use_cpu_only, backend, rank): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=1, high=4, size=rank) @make_tf_graph([input_shape, input_shape]) @@ -181,7 +182,7 @@ class TestDuplicateOutputs(TensorFlowBaseTest): ) def test(self, use_cpu_only, backend, rank): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=1, high=4, size=rank) @make_tf_graph([input_shape]) @@ -211,7 +212,7 @@ class TestIdentity(TensorFlowBaseTest): ) def test(self, use_cpu_only, backend, rank): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=1, high=4, size=rank) @@ -265,7 +266,7 @@ class TestAddN(TensorFlowBaseTest): ) def test(self, use_cpu_only, backend, rank, num_inputs): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=1, high=4, size=rank) input_shapes = [input_shape[:] for _ in range(num_inputs)] @@ -293,9 +294,6 @@ class TestAddOrdering(TensorFlowBaseTest): itertools.product([True, False], backends), ) def test(self, use_cpu_only, backend): - if backend[0] == "mlprogram": - pytest.xfail("Not supported on ML Program backend") - @make_tf_graph([(2, 3, 4), (2, 3, 4)]) def build_model(x, y): return tf.math.add(x, y) @@ -303,13 +301,20 @@ def build_model(x, y): model, inputs, outputs = build_model input_values = [random_gen((2, 3, 4), -1, 1)] * 2 input_dict = dict(zip(inputs, input_values)) - - mlmodel, input_key_values, output_names, output_nodes = tf_graph_to_mlmodel( - model, input_dict, outputs, "tensorflow", backend + + spec, _, _, _, _, _ = TensorFlowBaseTest.run_compare_tf( + model, + input_dict, + outputs, + use_cpu_for_conversion=use_cpu_only, + frontend_only=False, + backend=backend, ) - nn_spec = mlmodel.get_spec().neuralNetwork - assert nn_spec.layers[0].input[0] == "Placeholder" - assert nn_spec.layers[0].input[1] == "Placeholder_1" + + if backend[0] == "neuralnetwork": + nn_spec = spec.neuralNetwork + assert nn_spec.layers[0].input[0] == "Placeholder" + assert nn_spec.layers[0].input[1] == "Placeholder_1" class TestActivationLeakyReLU(TensorFlowBaseTest): @@ -397,7 +402,7 @@ class TestGeluTanhApproximation(TensorFlowBaseTest): ) def test(self, use_cpu_only, backend, rank): if backend[0] == 'mlprogram': - pytest.xfail("Not supported with ML Program backend") + pytest.skip("Not supported with ML Program backend") input_shape = np.random.randint(low=1, high=4, size=rank) @@ -597,10 +602,6 @@ class TestSelect(TensorFlowBaseTest): ), ) def test_select(self, use_cpu_for_conversion, backend, rank, broadcast, dynamic): - if backend[0] == "mlprogram" and broadcast and not use_cpu_for_conversion: - # use_cpu_for_conversion == False and broadcast == False fails with MIL - pytest.xfail("rdar://77441455") - shape = np.random.randint(low=1, high=4, size=rank) cond_shape = np.array([shape[0]]) if broadcast else shape @@ -719,7 +720,7 @@ class TestCond(TensorFlowBaseTest): ) def test_cond_naive(self, use_cpu_only, backend): if (backend[0] == "mlprogram" and backend[1] == "fp16"): - pytest.xfail("rdar://83626929 (Reenable CondTests)") + pytest.xfail("rdar://96627246 (ConsTest unittest is failing)") @make_tf_graph([(1,), (1,)]) def build_model(x, y): return tf.cond(tf.constant(True), lambda: x + y, lambda: x * y) @@ -1442,7 +1443,7 @@ def build_model_dynamic_weights(x, W): ) if backend[0] == "neuralnetwork" and dynamic_weights: - pytest.xfail("dynamic conv with groups > 1 is not supported on the neuralnetwork backend") + pytest.skip("dynamic conv with groups > 1 is not supported on the neuralnetwork backend") # We do not support dynamic weight when dilations != 1. test_dynamic_W() if dynamic_weights and dilations == (1, 1) else test_static_W() @@ -1571,7 +1572,7 @@ def build_model_static_weights(x): test_static_W() if not any([True if d > 1 else False for d in dilations]): if backend[0] == "neuralnetwork": - pytest.xfail("dynamic conv with groups > 1 is not supported on the neuralnetwork backend") + pytest.skip("dynamic conv with groups > 1 is not supported on the neuralnetwork backend") test_dynamic_W() class TestConvTranspose(TensorFlowBaseTest): @@ -1840,14 +1841,8 @@ class TestElementWiseBinary(TensorFlowBaseTest): def test_binary_math(self, use_cpu_only, backend, rank, tf_op, broadcast_case): if rank == 0 or broadcast_case == 0: - pytest.xfail("Rank-0 input is not supported") - - if backend[0] == "mlprogram" and not use_cpu_only and tf_op == tf.math.floormod: - pytest.xfail("rdar://78343225 ((MIL GPU) Core ML Tools Unit Test failures [numerical error])") - - if backend[0] == "mlprogram" and not use_cpu_only and tf_op == tf.math.floordiv: - pytest.xfail("rdar://82743379") - + pytest.skip("Rank-0 input is not supported") + x_shape = y_shape = list(np.random.randint(low=2, high=4, size=rank)) # test broadcasting @@ -1915,12 +1910,9 @@ def build_model(x, y): def test_binary_compare(self, use_cpu_for_conversion, backend, rank, tf_op, broadcast_case): if rank == 0 or broadcast_case == 0: - pytest.xfail("Rank-0 input is not supported") + pytest.skip("Rank-0 input is not supported") use_cpu_only = use_cpu_for_conversion - if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("Error in building plan : MIL GPU backend failure. rdar://78218824") - x_shape = y_shape = list(np.random.randint(low=2, high=4, size=rank)) # test broadcasting @@ -1973,12 +1965,9 @@ def build_model(x, y): def test_binary_logical(self, use_cpu_for_conversion, backend, rank, tf_op, broadcast_case): if rank == 0 or broadcast_case == 0: - pytest.xfail("Rank-0 input is not supported") + pytest.skip("Rank-0 input is not supported") use_cpu_only = use_cpu_for_conversion - if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("Error in building plan : MIL GPU backend failure. rdar://78218824") - x_shape = y_shape = list(np.random.randint(low=2, high=4, size=rank)) # test broadcasting @@ -2029,9 +2018,7 @@ class TestEinsum(TensorFlowBaseTest): ) ) def test(self, use_cpu_for_conversion, backend, equation, reverse_input_order): - if backend[0] == "mlprogram" and equation == "abcd,adce->abce" and not use_cpu_for_conversion: - pytest.xfail("rdar://80397986") - + if equation == "abcd,adce->abce": input_shapes = [[3, 4, 2, 6], [3, 6, 2, 2]] elif equation == "abc,cbd->abd": @@ -2118,17 +2105,18 @@ def test_unary(self, use_cpu_for_conversion, backend, rank, mode): if not use_cpu_only and mode in self._FP16_UNSUPPORTED: return - if not use_cpu_for_conversion: - pytest.xfail("rdar://78358015 (TestElementWiseUnary failing numerically on the GPU (both NNv1 and MIL))") - atol, rtol = 1e-4, 1e-5 input_shape = np.random.randint(low=2, high=4, size=rank) - if use_cpu_only: - dtype = np.float32 - tf_dtype = tf.float32 - else: + + if backend == ("mlprogram", "fp16") and mode != "clip": + # For the clip mode with tf.float16 as input, it seems like the tf graph is producing wrong results + # It looks like a tensorflow bug, tracked by this radar: + # rdar://96850184 (Tensor clip_by_value is producing wrong numerical outputs with tf.float16 type input) dtype = np.float16 tf_dtype = tf.float16 + else: + dtype = np.float32 + tf_dtype = tf.float32 def cast_func(x): return tf.cast(x, dtype=tf.int32) @@ -2198,7 +2186,7 @@ def _get_test(test_mode): # We skip GPU here, since exp(1) already differs in backend. return None, None res = tf.exp - val = random_gen(input_shape, rand_min=-4, rand_max=20) + val = random_gen(input_shape, rand_min=-4, rand_max=4) elif test_mode == "floor": res = tf.floor eps_from_int = 0.0 @@ -2279,6 +2267,7 @@ def build_model(x): backend=backend, atol=atol, rtol=rtol, + minimum_deployment_target=ct.target.iOS16 if backend == ("mlprogram", "fp16") else None, ) @@ -2391,10 +2380,6 @@ def test_crop_and_resize( method, dynamic, ): - - if backend[0] == "mlprogram" and not use_cpu_for_conversion and dynamic: - pytest.xfail("Seg fault. rdar://77444115") - use_cpu_only = use_cpu_for_conversion if backend[0] == "mlprogram" and not use_cpu_for_conversion and crop_size == (1, 1): # in this case, there is a numerical mismatch on the GPU MIL backend. The GPU runtime tests are @@ -3248,10 +3233,6 @@ def test_reduction(self, use_cpu_for_conversion, backend, rank_and_axes, keep_di rank, axes = rank_and_axes shape = np.random.randint(low=1, high=3, size=rank) - if backend[0] == 'mlprogram' and not use_cpu_for_conversion: - if rank_and_axes == (5, None) or tf_op in {tf.reduce_logsumexp}: - pytest.xfail("Seg fault. rdar://77443572") - def parse_axes(axes): if axes is None: axes = 0 @@ -3666,9 +3647,6 @@ def build_model(x, begin, end): def test_slice_by_index_from_scratch(self, use_cpu_only, backend, testcase): input_shape = np.array([3, 4, 5]) - if backend == ("mlprogram", "fp16") and testcase == (slice(None, None, None), slice(None, None, None), slice(None, None, -1)): - pytest.xfail("rdar://80661727 (SliceByIndex FP16 TF converter unit test failing on one particular configuration)") - @make_tf_graph([input_shape]) def build_model(x): return x[testcase] @@ -4028,7 +4006,7 @@ class TestFakeQuant(TensorFlowBaseTest): ) def test_fake_quant_weight_quantization_with_conv(self, num_bits, weight_boundaries, use_cpu_only, backend): if backend[0] == 'mlprogram': - pytest.xfail("Not supported with ML Program backend") + pytest.skip("Not supported with ML Program backend") tf.reset_default_graph() filter_width = 1 @@ -4205,10 +4183,10 @@ def test_non_max_suppression( score_threshold, use_V5, ): - if backend == ("mlprogram", "fp16") and not use_cpu_only: + if backend == ("mlprogram", "fp16"): + # rdar://80661262 ([GPU failures ] NonMaximumSuppression FP16 coremltools unit tests) + # rdar://86581713 ([MIL / FP16 / CPU only] NonMaximumSuppression appears to be swapping output values pytest.xfail("rdar://80661262 ([GPU failures ] NonMaximumSuppression FP16 coremltools unit tests)") - if backend == ("mlprogram", "fp16") and use_cpu_only: - pytest.xfail("rdar://86581713 ([MIL / FP16 / CPU only] NonMaximumSuppression appears to be swapping output values") boxes_val = random_gen(shape=(num_boxes, 4), rand_min=0, rand_max=32) scores_val = random_gen(shape=(num_boxes,), rand_min=-100, rand_max=100) @@ -4357,7 +4335,7 @@ class TestIdentityN(TensorFlowBaseTest): @pytest.mark.parametrize("use_cpu_only, backend", itertools.product([True, False],backends,) ) - def test(self, use_cpu_only, backend): + def test_identity_n(self, use_cpu_only, backend): shape_1 = [1,] shape_2 = [3, 4] shape_3 = [5, 6, 7] @@ -4376,6 +4354,24 @@ def build_model(x, y ,z): TensorFlowBaseTest.run_compare_tf(model, input_dict, outputs, use_cpu_for_conversion=use_cpu_only, frontend_only=False, backend=backend) + + @pytest.mark.parametrize("use_cpu_only, backend", + itertools.product([True, False], backends,) + ) + def test_identity_n_with_downstream_op(self, use_cpu_only, backend): + shape = [3, 4] + + @make_tf_graph([shape]) + def build_model(x): + x = tf.identity_n(input=[x, x]) + return tf.reduce_max(x, 1) + + model, inputs, outputs = build_model + input_values = [np.random.rand(*shape).astype(np.float32)] + input_dict = dict(zip(inputs, input_values)) + TensorFlowBaseTest.run_compare_tf(model, input_dict, outputs, + use_cpu_for_conversion=use_cpu_only, + frontend_only=False, backend=backend) class TestPad(TensorFlowBaseTest): @@ -4602,22 +4598,24 @@ def test_tile(self, use_cpu_only, backend, rank): class TestTopK(TensorFlowBaseTest): @pytest.mark.parametrize( - "use_cpu_only, backend, rank, k", + "use_cpu_only, backend, rank, k, sort", itertools.product( - [True, False], backends, [rank for rank in range(1, 6)], [1, 2, 3], + [True, False], backends, [1, 3, 5], [1, 3], [True, False], ), ) - def test_top_k(self, use_cpu_only, backend, rank, k): + def test_top_k(self, use_cpu_only, backend, rank, k, sort): # TensorFlow only supports last dimension (axis = -1). shape = np.random.randint(low=3, high=4, size=rank) - if backend == ("mlprogram", "fp16") and not use_cpu_only: - pytest.xfail("operation is ill-conditioned on FP16") + if not sort and backend[0] == "neuralnetwork": + pytest.skip("iOS16 version topk needed for sort = False") @make_tf_graph([shape]) def build_model(x): - ref = tf.math.top_k(x, k=k, sorted=True) - return (ref[1], ref[0]) + ref = tf.math.top_k(x, k=k, sorted=sort) + if not sort: + ref = (tf.sort(ref[0]), tf.sort(ref[1])) + return ref model, inputs, outputs = build_model input_values = [random_gen(shape, rand_min=-100, rand_max=100)] @@ -4628,6 +4626,7 @@ def build_model(x): outputs, use_cpu_for_conversion=use_cpu_only, backend=backend, + minimum_deployment_target=ct.target.iOS16 if not sort else None, ) @pytest.mark.parametrize( @@ -4672,9 +4671,6 @@ class TestConcat(TensorFlowBaseTest): )) def test_concat(self, use_cpu_only, backend, op_version, rank, num_inputs): - if backend == ("mlprogram", "fp16"): - pytest.xfail("rdar://80658663 (Concat FP16 tests are failing in self.fp16_overflow method)") - import random for axis in range(-rank, rank): input_shape = np.random.randint(low=1, high=4, size=rank) @@ -4714,12 +4710,9 @@ class TestSplit(TensorFlowBaseTest): itertools.product([True, False], backends, [1, 2, 3, 4], [True, False]), ) def test_split(self, use_cpu_for_conversion, backend, rank, dynamic): - if dynamic: - pytest.xfail("rdar://85318486 (Python unit tests on Split layer failing for both NNv1 and MIL backends)") - - if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("rdar://80397986") - + if backend[0] == "mlprogram" and not use_cpu_for_conversion and dynamic: + pytest.xfail("rdar://97398133 (TestSplit::test_split is failing on mlprogram + GPU + dynamic combination)") + input_shape1 = np.random.randint(low=1, high=3, size=rank) for axis in range(-rank, rank, 2): for split_num in range(2, input_shape1[axis] + 1, 2): @@ -4877,7 +4870,7 @@ class TestPack(TensorFlowBaseTest): ) def test_pack(self, use_cpu_only, backend, rank, num_inputs): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') shape = np.random.randint(low=1, high=4, size=rank) input_shapes = [shape[:] for _ in range(num_inputs)] @@ -4922,7 +4915,9 @@ def build_model(x): return tf.argsort(x, axis=axis, direction=direction.upper()) model, inputs, outputs = build_model - input_values = [random_gen(shape, rand_min=-100, rand_max=100, allow_duplicate=False, dtype=dtype)] + input_values = np.arange(np.prod(shape)) + np.random.shuffle(input_values) + input_values = [np.reshape(input_values, shape).astype(dtype)] input_dict = dict(zip(inputs, input_values)) TensorFlowBaseTest.run_compare_tf( model, @@ -5091,7 +5086,7 @@ def build_model(x, y): itertools.product([False], backends, [[1], [1, 1], [1, 1, -1], []],), ) def test_reshape_scalar(self, use_cpu_only, backend, shape): - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = () @@ -5202,9 +5197,6 @@ def test_reverse(self, use_cpu_only, backend, rank_and_axes): rank, axes = rank_and_axes shape = np.random.randint(low=1, high=4, size=rank) - if backend == ("mlprogram", "fp16") and not use_cpu_only: - pytest.xfail("rdar://80710884 (Crash in Reverse op unit test on MIL GPU context)") - @make_tf_graph([shape]) def build_model(x): return tf.reverse(x, axis=axes) @@ -5540,9 +5532,9 @@ def test_programmatic( if backend[0] == "neuralnetwork": if block_rank == 2 and block_shape[0] != block_shape[1]: - pytest.xfail("neuralnetwork backend doesn't support unequal block shape.") + pytest.skip("neuralnetwork backend doesn't support unequal block shape.") if block_shape[0] == 1: - pytest.xfail("neuralnetwork backend doesn't support unity block shape.") + pytest.skip("neuralnetwork backend doesn't support unity block shape.") paddings = [] for i in range(block_rank): @@ -5633,9 +5625,9 @@ def test_programmatic( if backend[0] == "neuralnetwork": if block_rank == 2 and block_shape[0] != block_shape[1]: - pytest.xfail("neuralnetwork backend doesn't support unequal block shape.") + pytest.skip("neuralnetwork backend doesn't support unequal block shape.") if block_shape[0] == 1: - pytest.xfail("neuralnetwork backend doesn't support unity block shape.") + pytest.skip("neuralnetwork backend doesn't support unity block shape.") input_shape[0] = input_shape[0] * np.prod(block_shape) crops = [] @@ -6110,7 +6102,7 @@ class TestZerosLike(TensorFlowBaseTest): ) def test(self, use_cpu_only, backend, rank, dynamic): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=2, high=4, size=rank) input_value = random_gen(input_shape, rand_min=-1, rand_max=1) if dynamic: @@ -6153,10 +6145,7 @@ class TestIsFinite(TensorFlowBaseTest): ) def test(self, use_cpu_for_conversion, backend, rank, dynamic): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') - - if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("rdar://78343191") + pytest.skip('Rank 0 not supported by CoreML runtime') def _generate_num_with_inf(input_shape): res = random_gen(input_shape, rand_min=-1, rand_max=1) @@ -6215,10 +6204,8 @@ class TestLogSoftMax(TensorFlowBaseTest): backends, )) def test(self, use_cpu_only, backend): - if backend == ("mlprogram", "fp16") and not use_cpu_only: - pytest.xfail("operation is ill-conditioned on FP16") input_shape = (5, 20) - input_value = random_gen(input_shape, rand_min=-10, rand_max=10) + input_value = random_gen(input_shape, rand_min=-1, rand_max=1) @make_tf_graph([input_shape]) def build_model(x): @@ -6242,7 +6229,7 @@ class TestClipByValue(TensorFlowBaseTest): )) def test(self, use_cpu_only, backend, rank, min_and_max): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=2, high=4, size=rank) min_val, max_val = min_and_max @@ -6270,7 +6257,7 @@ class TestSize(TensorFlowBaseTest): )) def test(self, use_cpu_only, backend, rank, dynamic): if rank == 0: - pytest.xfail('Rank 0 not supported by CoreML runtime') + pytest.skip('Rank 0 not supported by CoreML runtime') input_shape = np.random.randint(low=2, high=4, size=rank) input_value = random_gen(input_shape, rand_min=-1, rand_max=1) diff --git a/coremltools/converters/mil/frontend/tensorflow/test/testing_utils.py b/coremltools/converters/mil/frontend/tensorflow/test/testing_utils.py index d4a3f3a7b..86b1e39c2 100644 --- a/coremltools/converters/mil/frontend/tensorflow/test/testing_utils.py +++ b/coremltools/converters/mil/frontend/tensorflow/test/testing_utils.py @@ -116,7 +116,7 @@ def get_tf_node_names(tf_nodes, mode="inputs"): def tf_graph_to_mlmodel( graph, feed_dict, output_nodes, frontend="tensorflow", backend=("neuralnetwork", "fp32"), use_cpu_for_conversion=False, - inputs_for_conversion=None, + inputs_for_conversion=None, minimum_deployment_target=None, ): """ Parameters @@ -136,6 +136,8 @@ def tf_graph_to_mlmodel( It forces the model to be loaded on the CPU context, post conversion. 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. + minimum_deployment_target : coremltools.target enumeration + It set the minimum_deployment_target argument in the coremltools.convert functino. ----------- Returns MLModel, Input Values, Output Names """ @@ -159,6 +161,7 @@ def tf_graph_to_mlmodel( mlmodel = ct_convert( graph, inputs=inputs, outputs=output_names, source=frontend, convert_to=backend, compute_units=compute_unit, + minimum_deployment_target=minimum_deployment_target, ) return mlmodel, input_values, output_names, output_nodes @@ -196,6 +199,7 @@ def run_compare_tf( validate_shapes_only=False, freeze_graph=False, tf_outputs=None, + minimum_deployment_target=None, ): """ Utility function to convert and compare a given TensorFlow 1.x model. @@ -230,6 +234,8 @@ def run_compare_tf( all the variables in the graph have been converted to constants. tf_outputs: float or list[float] If present, use it as TensorFlow predictions + minimum_deployment_target : coremltools.target enumeration + It set the minimum_deployment_target argument in the coremltools.convert functino. Return: Proto, mlmodel, input dictionay, prediction(if possible) @@ -269,7 +275,9 @@ def run_compare_tf( mlmodel, input_key_values, output_names, output_nodes = tf_graph_to_mlmodel( graph, feed_dict, output_nodes, frontend, backend, - use_cpu_for_conversion=use_cpu_for_conversion, inputs_for_conversion=inputs_for_conversion, + use_cpu_for_conversion=use_cpu_for_conversion, + inputs_for_conversion=inputs_for_conversion, + minimum_deployment_target=minimum_deployment_target ) if frontend_only or coremltoolsutils._macos_version() < (10, 13) \ @@ -338,7 +346,7 @@ def run_compare_tf(graph, feed_dict, output_nodes, frontend_only=False, frontend="tensorflow", backend=("neuralnetwork", "fp32"), atol=1e-04, rtol=1e-05, validate_shapes_only=False, freeze_graph=False, - tf_outputs=None): + tf_outputs=None, minimum_deployment_target=None): res = run_compare_tf(graph, feed_dict, @@ -351,9 +359,10 @@ def run_compare_tf(graph, feed_dict, output_nodes, rtol=rtol, validate_shapes_only=validate_shapes_only, freeze_graph=freeze_graph, - tf_outputs=tf_outputs + tf_outputs=tf_outputs, + minimum_deployment_target=minimum_deployment_target ) - + alist = [] if res is not None: alist = list(res) diff --git a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/insert_get_tuple.py b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/insert_get_tuple.py index 676b6df3f..bca2a4b41 100644 --- a/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/insert_get_tuple.py +++ b/coremltools/converters/mil/frontend/tensorflow/tf_graph_pass/insert_get_tuple.py @@ -63,7 +63,7 @@ def make_op(input_node, index, new_node_name, gto_make_op_cache): "TensorArrayV3", "Const", ] - inclusions = ["Split", "SplitV", "LSTMBlockCell", "TopK", "TopKV2", "Unpack", "BlockLSTM", "BlockLSTMV2", "NonMaxSuppressionV5"] + inclusions = ["IdentityN", "Split", "SplitV", "LSTMBlockCell", "TopK", "TopKV2", "Unpack", "BlockLSTM", "BlockLSTMV2", "NonMaxSuppressionV5"] gto_make_op_cache = {} for name in list(gddict.keys()): new_node = ParsedTFNode() diff --git a/coremltools/converters/mil/frontend/tensorflow2/converter.py b/coremltools/converters/mil/frontend/tensorflow2/converter.py index a459102d4..dc52b5d7a 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/converter.py +++ b/coremltools/converters/mil/frontend/tensorflow2/converter.py @@ -12,8 +12,8 @@ class TF2Converter(TFConverter): - def __init__(self, tf_ssa, inputs=None, outputs=None, **kwargs): - TFConverter.__init__(self, tf_ssa, inputs, outputs, **kwargs) + def __init__(self, tf_ssa, inputs=None, outputs=None, opset_version=None): + TFConverter.__init__(self, tf_ssa, inputs, outputs, opset_version) # Overwrite tensorflow_passes # TF 2.x uses different set of graph passes diff --git a/coremltools/converters/mil/frontend/tensorflow2/load.py b/coremltools/converters/mil/frontend/tensorflow2/load.py index 546c6ee75..5c10013c0 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/load.py +++ b/coremltools/converters/mil/frontend/tensorflow2/load.py @@ -99,18 +99,21 @@ def __init__(self, model, debug=False, **kwargs): def _get_concrete_functions_and_graph_def(self): msg = ( "Expected model format: [SavedModel | [concrete_function] | " - "tf.keras.Model | .h5], got {}" + "tf.keras.Model | .h5 | GraphDef], got {}" ) if ( isinstance(self.model, list) or isinstance(self.model, _tf.keras.Model) or isinstance(self.model, str) + or isinstance(self.model, _tf.compat.v1.GraphDef) ): cfs = [] if isinstance(self.model, list): cfs = self.model if isinstance(self.model, _tf.keras.Model): cfs = self._concrete_fn_from_tf_keras_or_h5(self.model) + elif isinstance(self.model, _tf.compat.v1.GraphDef): + return None, self.model elif isinstance(self.model, str): if not _os_path.exists(self.model): raise ValueError( @@ -196,7 +199,12 @@ def _run_tf_ssa_passes(self): def _program_from_tf_ssa(self): self._run_tf_ssa_passes() - converter = TF2Converter(self._tf_ssa, **self.kwargs) + converter = TF2Converter( + tf_ssa=self._tf_ssa, + inputs=self.kwargs["inputs"], + outputs=self.kwargs["outputs"], + opset_version=self.kwargs["specification_version"] + ) return converter.convert() def _populate_sub_graph_input_shapes(self, graph, graph_fns): 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 6d16e3561..7ef459db8 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 @@ -6,10 +6,11 @@ import logging 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.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass - +@block_context_manager def _remove_vacuous_cond_block(block): num_changes = 0 for op in list(block.operations): @@ -54,11 +55,10 @@ def _remove_vacuous_cond_block(block): continue new_var = pred_y.ls - with block: - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=new_var - ) - block.remove_ops([op]) # rely on DCE to remove extra cond inputs + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=new_var + ) + block.remove_ops([op]) # rely on DCE to remove extra cond inputs num_changes += 1 # Pattern 2: both than and else branch contains exactly 1 identity op @@ -68,12 +68,11 @@ def _remove_vacuous_cond_block(block): if then_ops[0].x != else_ops[0].x: continue - with block: - new_var = mb.identity(x=then_ops[0].x, before_op=op, name=op.name) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=new_var - ) - block.remove_ops([op]) # rely on DCE to remove extra cond inputs + new_var = mb.identity(x=then_ops[0].x, before_op=op, name=op.name) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=new_var + ) + block.remove_ops([op]) # rely on DCE to remove extra cond inputs num_changes += 1 return num_changes diff --git a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_load.py b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_load.py index 28b1e2b80..50a3a1168 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_load.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_load.py @@ -30,6 +30,14 @@ # ----------------------------------------------------------------------------- # Import TF 2.x-compatible TF 1.x test cases +from coremltools.converters.mil.frontend.tensorflow2.test.testing_utils import ( + TensorFlow2BaseTest +) +from coremltools.converters.mil.frontend.tensorflow.test.testing_utils import ( + TensorFlowBaseTest +) +TensorFlowBaseTest.run_compare_tf = TensorFlow2BaseTest.run_compare_tf2 + from coremltools.converters.mil.frontend.tensorflow.test.test_load import ( frontend, TestTf1ModelInputsOutputs as TestTf2ModelInputsOutputs, @@ -138,6 +146,31 @@ def __call__(self, x): [concrete_func], outputs=["Identity"], source=frontend ) assert mlmodel is not None + + def test_graphdef_from_tf_function(self): + class build_model(tf.Module): + def __init__(self): + self.dense = tf.keras.layers.Dense(256, activation="relu") + + input_signature = [ + tf.TensorSpec(name="input", shape=( + 128, 128), dtype=tf.float32), + ] + + @tf.function(input_signature=input_signature) + def call(self, x): + x = self.dense(x) + return x + + model = build_model() + + from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 + frozen_graph_func = convert_variables_to_constants_v2( + model.call.get_concrete_function()) + frozen_graph_def = frozen_graph_func.graph.as_graph_def() + + mlmodel = converter.convert(frozen_graph_def) + assert mlmodel is not None def test_model_metadata(self): keras_model = tf.keras.Sequential( 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 e27da5da4..4dfa4bfbc 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/test_v2_ops.py @@ -140,7 +140,7 @@ def test_resample( self, use_cpu_only, backend, data_warp_shapes, ): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") tfa = pytest.importorskip("tensorflow_addons") @@ -189,7 +189,7 @@ class TestImageTransform(TensorFlow2BaseTest): def test(self, use_cpu_only, backend, transforms, interpolation, shapes): x_shape, output_shape = shapes if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") tfa = pytest.importorskip("tensorflow_addons") @@ -227,7 +227,7 @@ class TestActivationSiLU(TensorFlow2BaseTest): ) def test(self, use_cpu_only, backend, rank, tf_op): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") x_shape = tuple(np.random.randint(low=1, high=4, size=rank)) @@ -275,12 +275,9 @@ def test_raw_ops( return if target_shape[-2] % input_shape[-2] != 0: return - - if not use_cpu_only and not half_pixel_centers and backend[0] == "mlprogram": - # use_cpu_only == False & half_pixel_centers == False, & backend[0] == mlprogram - # then there are numerical errors - pytest.xfail("rdar://78321005") - + + if backend[0] == "mlprogram" and not use_cpu_only and not half_pixel_centers: + pytest.xfail("rdar://97399545 (TestResizeNearestNeighbor failing on mlprogram + GPU + half_pixel_centers=False)") @make_tf_graph([input_shape]) def build_model(x): @@ -304,7 +301,7 @@ def build_model(x): ) def test_keras_layer(self, use_cpu_only, backend, size): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") x_shape = tuple(np.random.randint(low=1, high=4, size=4)) @@ -331,14 +328,11 @@ def build_model(x): ), ) def test_tf_image_resize(self, use_cpu_only, backend, size, method): - if backend[0] == "mlprogram" and not use_cpu_only: - pytest.xfail("rdar://78343225 ((MIL GPU) Core ML Tools Unit Test failures [numerical error])") - if backend[0] == "mlprogram" and size == (1, 1): pytest.xfail("rdar://79699954 (Nearest neighbor resize numerical mismatch when output size is (1,1))") if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") x_shape = tuple(np.random.randint(low=1, high=3, size=4)) @@ -482,9 +476,6 @@ def build_model(x): "use_cpu_only, backend", itertools.product([True, False], backends) ) def test_if_binary_add_if_else_mul(self, use_cpu_only, backend): - if backend[0] == "mlprogram": - pytest.xfail("rdar://81983176 (MLProgram Failure: Mismatched elements)") - @make_tf_graph([(1,), (1,)]) def build_model(x, y): if x > y: 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 0cdffc01a..b5dd033d4 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 @@ -495,10 +495,6 @@ def test_conv2d_padding_dynamic_input( backend, padding, ): - - if backend[0] == "mlprogram" and ct.utils._macos_version() > (12, 0): - pytest.xfail("rdar://88857567") - # Test same padding input_layer = Input(batch_size=1, shape=(None, None, 1)) layer = Conv2D( @@ -1609,7 +1605,10 @@ def test( assert len(layer.upsample.fractionalScalingFactor) == 0 class TestGelu(TensorFlowBaseTest): - @pytest.mark.xfail(_get_version(_tf.__version__) < _StrictVersion("2.4.0"), reason="Gelu is a new layer for tf 2.4.0 and above.") + @pytest.mark.skipif( + _get_version(_tf.__version__) < _StrictVersion("2.4.0"), + reason="Gelu is a new layer for tf 2.4.0 and above." + ) @pytest.mark.parametrize( "use_cpu_only, backend, rank, approximate", itertools.product( diff --git a/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py b/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py index 6129d9e56..06fe6b31a 100644 --- a/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py +++ b/coremltools/converters/mil/frontend/tensorflow2/test/testing_utils.py @@ -81,6 +81,7 @@ def run_compare_tf2( debug=False, atol=1e-04, rtol=1e-05, + minimum_deployment_target=None, ): """ Parameters @@ -107,6 +108,8 @@ def run_compare_tf2( The absolute tolerance parameter. rtol: float The relative tolerance parameter. + minimum_deployment_target: coremltools.target enumeration + The spec version for the mlmodel """ inputs = [] if inputs_for_conversion is None: @@ -147,6 +150,7 @@ def run_compare_tf2( convert_to=backend, debug=debug, compute_units=compute_unit, + minimum_deployment_target=minimum_deployment_target, ) for k,v in input_dict.items(): @@ -260,8 +264,8 @@ def run_compare_tf2(model, backend=("neuralnetwork", "fp32"), debug=False, atol=1e-04, - rtol=1e-05): - + rtol=1e-05, + minimum_deployment_target=None): res = run_compare_tf2(model, input_dict, output_names, @@ -272,7 +276,8 @@ def run_compare_tf2(model, backend=backend, debug=debug, atol=atol, - rtol=rtol) + rtol=rtol, + minimum_deployment_target=minimum_deployment_target,) alist = list(res) alist.append(TensorFlow2BaseTest.testclassname) alist.append(TensorFlow2BaseTest.testmodelname) @@ -292,4 +297,4 @@ def run_compare_tf_keras(model, input_values, inputs_for_conversion=None, use_cp alist = list(res) alist.append(TensorFlow2BaseTest.testclassname) alist.append(TensorFlow2BaseTest.testmodelname) - return tuple(alist) + return tuple(alist) \ No newline at end of file diff --git a/coremltools/converters/mil/frontend/torch/converter.py b/coremltools/converters/mil/frontend/torch/converter.py index 36d716795..5508ee7e5 100644 --- a/coremltools/converters/mil/frontend/torch/converter.py +++ b/coremltools/converters/mil/frontend/torch/converter.py @@ -5,9 +5,11 @@ from collections import OrderedDict import logging as _logging +import numpy as _np import torch as _torch from coremltools._deps import version_lt +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as _target from coremltools.converters.mil.input_types import ImageType from coremltools.converters.mil.mil import ( Builder as mb, @@ -134,7 +136,7 @@ class TorchConverter: """ def __init__( - self, torchscript, inputs, outputs=None, cut_at_symbols=None, + self, torchscript, inputs, outputs=None, cut_at_symbols=None, opset_version=None ): """ Arguments: @@ -144,6 +146,7 @@ def __init__( cut_at_symbols: A list of internal symbol name strings. Graph conversion will terminate once these symbols have been generated. For debugging use only. See kwarg in load.py. + opset_version: An int represents the coreml opset version """ assert isinstance(torchscript, _torch.jit.ScriptModule) @@ -154,6 +157,7 @@ def __init__( self.torchscript = torchscript self.outputs = outputs self.output_names = get_output_names(self.outputs) + self.opset_version = _target(opset_version) if opset_version is not None else None self.context = TranscriptionContext() raw_graph, params_dict = self._expand_and_optimize_ir(self.torchscript) self.params_dict = params_dict @@ -214,6 +218,8 @@ def check_ops(self): def convert_const(self): for name, val in self.graph.params.items(): + if val.dtype == _np.uint8: + val = val.astype(_np.int32) const = mb.const(val=val, name=name) self.context.add(const) @@ -236,7 +242,7 @@ def convert(self): prog.set_main_input_types(tuple(self.inputs)) # Initialize the SSA for conversion - with Function(ssa_func_inputs) as ssa_func: + with Function(ssa_func_inputs, opset_version=self.opset_version) as ssa_func: # Map internal @self.graph.inputs to user specified @ssa_func_inputs # If @self.graph.inputs == @ssa_func_inputs this just adds the inputs diff --git a/coremltools/converters/mil/frontend/torch/dialect_ops.py b/coremltools/converters/mil/frontend/torch/dialect_ops.py index 9d4f7a9be..651458a69 100644 --- a/coremltools/converters/mil/frontend/torch/dialect_ops.py +++ b/coremltools/converters/mil/frontend/torch/dialect_ops.py @@ -13,7 +13,7 @@ IntTensorInputType, TensorInputType ) -from coremltools.converters.mil.mil.ops.defs.tensor_transformation import _solve_slice_by_index_shape +from coremltools.converters.mil.mil.ops.defs._utils 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 @@ -31,7 +31,7 @@ # torch_upsample_nearest_neighbor is dealing with upsample layer which has flexible input shape, # and recompute_scale_factor is set to True in the original torch layer. -@register_op(doc_str="", namespace="torch") +@register_op(namespace="torch") class torch_upsample_nearest_neighbor(Operation): """ Upsample the spatial dimensions (last two dimensions) of the input by @@ -81,7 +81,7 @@ def type_inference(self): # torch_upsample_bilinear is dealing with upsample layer which has flexible input shape, # and recompute_scale_factor is set to True in the original torch layer. -@register_op(doc_str="", namespace="torch") +@register_op(namespace="torch") class torch_upsample_bilinear(Operation): """ Upsample the spatial dimensions (last two dimensions) of the input by @@ -138,7 +138,7 @@ def type_inference(self): return types.tensor(self.x.dtype, ret_shape) # torch_tensor_assign is dealing with the tensor assignment operation -@register_op(doc_str="", namespace="torch") +@register_op(namespace="torch") class torch_tensor_assign(Operation): """ Method for tensor value assignment via indexing and slicing. @@ -215,7 +215,7 @@ def type_inference(self): self.squeeze_mask.val if self.squeeze_mask is not None else [False] * data_rank ) data_shape = self.data.shape - expected_updates_shape = tuple(_solve_slice_by_index_shape(data_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask)) + expected_updates_shape = tuple(solve_slice_by_index_shape(data_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask)) if not is_compatible_symbolic_vector(expected_updates_shape, self.updates.shape): raise ValueError("The updates tensor should have shape {}. Got {}".format(expected_updates_shape, self.updates.shape)) return self.data.sym_type diff --git a/coremltools/converters/mil/frontend/torch/load.py b/coremltools/converters/mil/frontend/torch/load.py index a77d730d7..ac9af9a42 100644 --- a/coremltools/converters/mil/frontend/torch/load.py +++ b/coremltools/converters/mil/frontend/torch/load.py @@ -48,7 +48,8 @@ def load(model_spec, debug=False, **kwargs): 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) + opset_version = kwargs["specification_version"] + converter = TorchConverter(torchscript, inputs, outputs, cut_at_symbols, opset_version) return _perform_torch_convert(converter, debug) diff --git a/coremltools/converters/mil/frontend/torch/ops.py b/coremltools/converters/mil/frontend/torch/ops.py index 75e5fad56..6ea7c30a7 100644 --- a/coremltools/converters/mil/frontend/torch/ops.py +++ b/coremltools/converters/mil/frontend/torch/ops.py @@ -19,7 +19,9 @@ types, Symbol ) -from coremltools.converters.mil.mil.ops.defs.tensor_transformation import _solve_slice_by_index_shape +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target +from coremltools.converters.mil.mil.block import curr_opset_version, is_current_opset_version_compatible_with +from coremltools.converters.mil.mil.ops.defs._utils import solve_slice_by_index_shape from coremltools.converters.mil.mil.types import is_bool from coremltools.converters.mil.mil.types.symbolic import any_symbolic, is_symbolic, is_compatible_symbolic_vector from coremltools.converters.mil.mil.var import Var, ListVar @@ -32,6 +34,7 @@ # This is a magic number in PyTorch. It's used as a default value in many # functions. PYTORCH_MAGIC_DEFAULT = 9223372036854775807 +VALUE_CLOSE_TO_INFINITY = 1e+38 def _all_outputs_present(context, graph): @@ -423,30 +426,98 @@ def frobenius_norm(context, node): @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) + assert x is not None and keep_dims is not None and num is not None and dim is not None + temp = _vector_norm(x=x, order=num, dim=dim, keep_dims=keep_dims, name=node.name) + context.add(temp) + + +def _vector_norm(x, order, dim, keep_dims, name): + if order.val == 0: + # sum(x!=0) + temp = mb.not_equal(x=x, y=0) + temp = mb.cast(x=temp, dtype='int32') + temp = mb.reduce_sum(x=temp, axes=dim, keep_dims=keep_dims, name=name) + elif order.val > VALUE_CLOSE_TO_INFINITY: + # max(abs(x)) + temp = mb.abs(x=x) + temp = mb.reduce_max(x=temp, axes=dim, keep_dims=keep_dims, name=name) + elif order.val < -VALUE_CLOSE_TO_INFINITY: + # min(abs(x)) + temp = mb.abs(x=x) + temp = mb.reduce_min(x=temp, axes=dim, keep_dims=keep_dims, name=name) else: - # sum(abs(x)**num)**(1./num) + # sum(abs(x)^{ord})^{(1 / ord)} temp = mb.abs(x=x) - temp = mb.pow(x=temp, y=num) + temp = mb.pow(x=temp, y=order.val) temp = mb.reduce_sum(x=temp, axes=dim, keep_dims=keep_dims) - temp = mb.pow(x=temp, y=1./num.val, name=node.name) + temp = mb.pow(x=temp, y=1./order.val, name=name) + return temp + + +def _matrix_norm(x, order, dim, keep_dims, name): + if order.val == 1: + # min(sum(abs(x), dim=0)) + temp = mb.abs(x=x) + temp = mb.reduce_sum(x=temp, axes=[dim[0]], keep_dims=True) + temp = mb.reduce_max(x=temp, axes=dim, keep_dims=keep_dims, name=name) + elif order.val == -1: + # min(sum(abs(x), dim=0)) + temp = mb.abs(x=x) + temp = mb.reduce_sum(x=temp, axes=[dim[0]], keep_dims=True) + temp = mb.reduce_min(x=temp, axes=dim, keep_dims=keep_dims, name=name) + elif order.val == "fro": + # sum(x**2)**1/2 + temp = mb.reduce_l2_norm(x=x, axes=dim, keep_dims=keep_dims, name=name) + elif order.val > VALUE_CLOSE_TO_INFINITY: + # max(sum(abs(x), dim=1)) + temp = mb.abs(x=x) + temp = mb.reduce_sum(x=temp, axes=[dim[1]], keep_dims=True) + temp = mb.reduce_max(x=temp, axes=dim, keep_dims=keep_dims, name=name) + elif order.val < -VALUE_CLOSE_TO_INFINITY: + # min(sum(abs(x), dim=1)) + temp = mb.abs(x=x) + temp = mb.reduce_sum(x=temp, axes=[dim[1]], keep_dims=True) + temp = mb.reduce_min(x=temp, axes=dim, keep_dims=keep_dims, name=name) + else: + raise RuntimeError("Matrix norm is not defined for the current inputs") + return temp + + +@register_torch_op +def linalg_vector_norm(context, node): + x, order, dim, keep_dims, _ = _get_inputs(context, node, expected=5) + assert x is not None and keep_dims is not None and order is not None + temp = _vector_norm(x=x, order=order, dim=dim, keep_dims=keep_dims, name=node.name) context.add(temp) + +@register_torch_op +def linalg_matrix_norm(context, node): + x, order, dim, keep_dims, _ = _get_inputs(context, node, expected=5) + assert x is not None and keep_dims is not None and order is not None and dim is not None + assert len(dim.val) == 2 + temp = _matrix_norm(x=x, order=order, dim=dim.val, keep_dims=keep_dims, name=node.name) + context.add(temp) + + +@register_torch_op +def linalg_norm(context, node): + x, order, dim, keep_dims, _ = _get_inputs(context, node, expected=5) + assert x is not None and keep_dims is not None + if dim is None: + dim = _np.arange(x.rank) + else: + dim = dim.val + if order is None: + temp = mb.reduce_l2_norm(x=x, axes=dim, keep_dims=keep_dims, name=node.name) + elif len(dim)==2: + temp = _matrix_norm(x=x, order=order, dim=dim, keep_dims=keep_dims, name=node.name) + else: + temp = _vector_norm(x=x, order=order, dim=dim, keep_dims=keep_dims, name=node.name) + context.add(temp) + + @register_torch_op def hardswish(context, node): inputs = _get_inputs(context, node, expected=1) @@ -582,6 +653,12 @@ def pixel_shuffle(context, node): perm = mb.pixel_shuffle(x=inputs[0], upscale_factor=inputs[1], name=node.name) context.add(perm) +@register_torch_op +def pixel_unshuffle(context, node): + inputs = _get_inputs(context, node, expected=2) + downscale_factor = _np.uint32(inputs[1].val) + perm = mb.pixel_unshuffle(x=inputs[0], downscale_factor=downscale_factor, name=node.name) + context.add(perm) @register_torch_op(torch_alias=["bmm"]) def matmul(context, node): @@ -2946,7 +3023,7 @@ def _internal_op_tensor_inplace_fill(context, node): fill_scalar = context[node.inputs[1]] begin, end, stride, begin_mask, end_mask, squeeze_mask = _get_slice_params(context, data, node.inputs[2:]) - fill_shape = _solve_slice_by_index_shape(data.shape, begin, end, stride, begin_mask, end_mask, squeeze_mask) + fill_shape = solve_slice_by_index_shape(data.shape, begin, end, stride, begin_mask, end_mask, squeeze_mask) update_values = _np.full(fill_shape, fill_scalar.val) updated_x = mb.torch_tensor_assign( @@ -4158,15 +4235,20 @@ def topk(context, node): kwargs["ascending"] = not largest # last inputs to topk are optional - sorted and out. + sort = True if len(inputs) > 4: - if inputs[4].val is False: - raise Exception("Unsupported value for argument 'sorted' in topk. Supported values: True, but input " - "is {}".format(inputs[4].val)) + if inputs[4].val is False and curr_opset_version() < target.iOS16: + raise Exception("For opset <= iOS16, only sorted=True supported for the topk") + sort = inputs[4].val + if len(inputs) > 5: if inputs[5] is not None: raise Exception("Unsupported value for argument 'out' in topk. Supported values: None, but input " "is {}".format(inputs[5].val)) + if is_current_opset_version_compatible_with(target.iOS16): + kwargs["sort"] = sort + res = mb.topk(**kwargs) values_name = node.outputs[0] indices_name = node.outputs[1] 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 e1d5ec570..9eb926c5c 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 @@ -3,10 +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 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.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass @register_pass(namespace="torch") class torch_tensor_assign_to_core(AbstractGraphPass): @@ -37,15 +37,14 @@ def apply(self, prog): for f in prog.functions.values(): _torch_tensor_assign_to_core_block(f) - +@block_context_manager def _torch_tensor_assign_to_core_block(block): for op in block.operations[:]: for b in op.blocks: _torch_tensor_assign_to_core_block(b) if op.op_type in ["torch_tensor_assign"]: - with block: - _transform_tensor_assign(op, block) + _transform_tensor_assign(op, block) def _transform_tensor_assign(op, block): 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 1622ce1a9..1fd1fe294 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 @@ -7,6 +7,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass target_ops = [ @@ -30,7 +31,7 @@ def apply(self, prog): for f in prog.functions.values(): _torch_upsample_to_core_upsample_block(f) - +@block_context_manager def _torch_upsample_to_core_upsample_block(block): for op in block.operations[:]: for b in op.blocks: @@ -108,25 +109,24 @@ def _try_replace_with_core_upsample(op): old_upsample = op.outputs[0] block = op.enclosing_block - with block: - if op.op_type == "torch_upsample_nearest_neighbor": - new_upsample = mb.upsample_nearest_neighbor( - x=op.x, - scale_factor_height=scales_h, - scale_factor_width=scales_w, - name=op.name, - before_op=op, - ) - elif op.op_type == "torch_upsample_bilinear": - new_upsample = mb.upsample_bilinear( - x=op.x, - scale_factor_height=scales_h, - scale_factor_width=scales_w, - align_corners=op.align_corners, - name=op.name, - before_op=op, - ) - block.replace_uses_of_var_after_op(anchor_op=op, old_var=old_upsample, new_var=new_upsample) + if op.op_type == "torch_upsample_nearest_neighbor": + new_upsample = mb.upsample_nearest_neighbor( + x=op.x, + scale_factor_height=scales_h, + scale_factor_width=scales_w, + name=op.name, + before_op=op, + ) + elif op.op_type == "torch_upsample_bilinear": + new_upsample = mb.upsample_bilinear( + x=op.x, + scale_factor_height=scales_h, + scale_factor_width=scales_w, + align_corners=op.align_corners, + name=op.name, + before_op=op, + ) + block.replace_uses_of_var_after_op(anchor_op=op, old_var=old_upsample, new_var=new_upsample) block.remove_ops([op]) return True diff --git a/coremltools/converters/mil/frontend/torch/test/test_conversion_api.py b/coremltools/converters/mil/frontend/torch/test/test_conversion_api.py index 8762d8796..2da1ab971 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_conversion_api.py +++ b/coremltools/converters/mil/frontend/torch/test/test_conversion_api.py @@ -96,7 +96,7 @@ def forward(self, x): def rank4_grayscale_input_model(): class Model(torch.nn.Module): def forward(self, x): - return x + 5.5 + return x + 10 example_input = torch.randint(0, 100, (1, 1, 10, 20), dtype=torch.float32) return torch.jit.trace(Model().eval(), example_input) @@ -394,17 +394,17 @@ def test_color_input(self, rank4_input_model, rank3_input_model): def test_grayscale_input(self, rank4_input_model, rank3_input_model, rank4_grayscale_input_model): with pytest.raises(ValueError, match="must have rank 4"): - mlmodel = ct.convert(rank3_input_model, - inputs=[ct.ImageType(shape=(1, 10, 20), color_layout=ct.colorlayout.GRAYSCALE)], - minimum_deployment_target=ct.target.macOS13, - ) + ct.convert(rank3_input_model, + inputs=[ct.ImageType(shape=(1, 10, 20), color_layout=ct.colorlayout.GRAYSCALE)], + minimum_deployment_target=ct.target.macOS13, + ) # invalid shape with pytest.raises(ValueError): - mlmodel = ct.convert(rank4_input_model, - inputs=[ct.ImageType(shape=(1, 3, 10, 20), color_layout=ct.colorlayout.GRAYSCALE)], - minimum_deployment_target=ct.target.macOS13, - ) + ct.convert(rank4_input_model, + inputs=[ct.ImageType(shape=(1, 3, 10, 20), color_layout=ct.colorlayout.GRAYSCALE)], + minimum_deployment_target=ct.target.macOS13, + ) mlmodel = ct.convert(rank4_grayscale_input_model, inputs=[ct.ImageType(shape=(1, 1, 10, 20), color_layout=ct.colorlayout.GRAYSCALE)], @@ -417,18 +417,18 @@ def test_grayscale_input(self, rank4_input_model, rank3_input_model, rank4_grays verify_prediction(mlmodel) with pytest.raises(TypeError, match="float16 dtype for inputs is only supported for deployment target >= iOS16/macOS13"): - mlmodel = ct.convert(rank4_grayscale_input_model, - inputs=[ct.ImageType(shape=(1, 1, 10, 20), - color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], - minimum_deployment_target=ct.target.macOS12, - ) + ct.convert(rank4_grayscale_input_model, + inputs=[ct.ImageType(shape=(1, 1, 10, 20), + color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], + minimum_deployment_target=ct.target.macOS12, + ) # test that grayscale_16 raises error when used with neural network with pytest.raises(TypeError, match="float16 dtype for inputs is only supported for deployment target >= iOS16/macOS13"): - mlmodel = ct.convert(rank4_grayscale_input_model, - inputs=[ct.ImageType(shape=(1, 1, 10, 20), - color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], - ) + ct.convert(rank4_grayscale_input_model, + inputs=[ct.ImageType(shape=(1, 1, 10, 20), + color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], + ) mlmodel = ct.convert(rank4_grayscale_input_model, inputs=[ct.ImageType(shape=(1, 1, 10, 20), @@ -440,8 +440,7 @@ def test_grayscale_input(self, rank4_input_model, rank3_input_model, rank4_grays assert_spec_input_image_type(mlmodel._spec, expected_feature_type=ft.ImageFeatureType.GRAYSCALE_FLOAT16) assert_prog_input_type(mlmodel._mil_program, expected_dtype_str="fp16") assert_output_dtype(mlmodel, expected_type_str="fp16") - # TODO: uncomment the following when rdar://92239179 is fixed - # verify_prediction(mlmodel) + verify_prediction(mlmodel) def test_color_output(self, rank4_input_model, float32_input_model_add_op): # check that an error is raised if the output shape is not of form (1, 3, H, W) @@ -477,11 +476,11 @@ def test_color_output(self, rank4_input_model, float32_input_model_add_op): def test_grayscale_output(self, rank4_grayscale_input_model): with pytest.raises(TypeError, match="float16 dtype for outputs is only supported for deployment target >= iOS16/macOS13"): - mlmodel = ct.convert(rank4_grayscale_input_model, - inputs=[ct.TensorType(shape=(1, 1, 10, 20))], - outputs=[ct.ImageType(color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], - minimum_deployment_target=ct.target.macOS12, - ) + ct.convert(rank4_grayscale_input_model, + inputs=[ct.TensorType(shape=(1, 1, 10, 20))], + outputs=[ct.ImageType(color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], + minimum_deployment_target=ct.target.macOS12, + ) mlmodel = ct.convert(rank4_grayscale_input_model, inputs=[ct.ImageType(shape=(1, 1, 10, 20), @@ -504,8 +503,7 @@ def test_grayscale_output(self, rank4_grayscale_input_model): assert_spec_output_image_type(mlmodel._spec, expected_feature_type=ft.ImageFeatureType.GRAYSCALE_FLOAT16) assert_prog_input_type(mlmodel._mil_program, expected_dtype_str="fp16") assert_prog_output_type(mlmodel._mil_program, expected_dtype_str="fp16") - # TODO: uncomment the following when rdar://92239179 is fixed - # verify_prediction(mlmodel) + verify_prediction(mlmodel) mlmodel = ct.convert(rank4_grayscale_input_model, inputs=[ct.ImageType(shape=(1, 1, 10, 20), @@ -518,8 +516,7 @@ def test_grayscale_output(self, rank4_grayscale_input_model): assert_spec_output_image_type(mlmodel._spec, expected_feature_type=ft.ImageFeatureType.GRAYSCALE_FLOAT16) assert_prog_input_type(mlmodel._mil_program, expected_dtype_str="fp32") assert_prog_output_type(mlmodel._mil_program, expected_dtype_str="fp16") - # TODO: uncomment the following when rdar://92239179 is fixed - # verify_prediction(mlmodel) + verify_prediction(mlmodel) def test_linear_model(self, linear_model): # this will test the fuse_linear_bias pass, when the inputs are of type float16 @@ -571,3 +568,82 @@ def test_prediction_with_fp16_io(self): reference_output = traced_model(torch.from_numpy(sample_input)).detach().numpy() np.testing.assert_allclose(reference_output, model_output, rtol=1e-2, atol=1e-2) + +@pytest.mark.skipif(ct.utils._macos_version() < (13, 0), reason='Tests are for deployment target ios16/macos13') +class TestGrayscaleImagePredictions: + + def test_grayscale_input_image(self, rank4_grayscale_input_model): + mlmodel = ct.convert(rank4_grayscale_input_model, + inputs=[ct.ImageType(name="input_image", + shape=(1, 1, 10, 20), + color_layout=ct.colorlayout.GRAYSCALE)], + outputs=[ct.TensorType(name="output")], + minimum_deployment_target=ct.target.macOS13, + ) + sample_input = np.random.randint(low=0, high=246, size=(1, 1, 10, 20)) + img_input = Image.fromarray(sample_input[0, 0, :, :].astype(np.uint8), 'L') + model_output = mlmodel.predict({"input_image": img_input})['output'] + reference_output = rank4_grayscale_input_model(torch.from_numpy(sample_input.astype(np.float32))).detach().numpy() + np.testing.assert_allclose(reference_output, model_output, rtol=1e-2, atol=1e-2) + + def test_grayscale_fp16_input_image(self, rank4_grayscale_input_model): + mlmodel = ct.convert(rank4_grayscale_input_model, + inputs=[ct.ImageType(name="input_image", + shape=(1, 1, 10, 20), + color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], + outputs=[ct.TensorType(name="output")], + minimum_deployment_target=ct.target.macOS13, + ) + + # incorrect way to do prediction + with pytest.raises(TypeError, + match="must be of type PIL.Image.Image with mode=='F'", + ): + sample_input = np.random.randint(low=0, high=246, size=(1, 1, 10, 20)) + img_input = Image.fromarray(sample_input[0, 0, :, :].astype(np.uint8), 'L') + mlmodel.predict({"input_image": img_input}) + + # correct way to do prediction + sample_input = np.random.rand(1, 1, 10, 20) # in between [0, 1] + img_input = Image.fromarray(sample_input[0, 0, :, :].astype(np.float32), 'F') + model_output = mlmodel.predict({"input_image": img_input})['output'] + reference_output = rank4_grayscale_input_model(torch.from_numpy(sample_input.astype(np.float32))).detach().numpy() + np.testing.assert_allclose(reference_output, model_output, rtol=1e-2, atol=1e-2) + + def test_grayscale_output_image(self, rank4_grayscale_input_model): + mlmodel = ct.convert(rank4_grayscale_input_model, + inputs=[ct.TensorType(name="input", + shape=(1, 1, 10, 20))], + outputs=[ct.ImageType(name="output_image", + color_layout=ct.colorlayout.GRAYSCALE)], + minimum_deployment_target=ct.target.macOS13, + compute_precision=ct.precision.FLOAT32, + ) + sample_input = np.random.randint(low=0, high=200, size=(1, 1, 10, 20)).astype(np.float32) + model_output_pil_image = mlmodel.predict({"input": sample_input})['output_image'] + assert isinstance(model_output_pil_image, Image.Image) + assert model_output_pil_image.mode == "L" + model_output_as_numpy = np.array(model_output_pil_image) + reference_output = rank4_grayscale_input_model(torch.from_numpy(sample_input)).detach().numpy() + reference_output = np.squeeze(reference_output) + np.testing.assert_allclose(reference_output, model_output_as_numpy, rtol=1e-2, atol=1e-2) + + def test_grayscale_fp16_output_image(self, rank4_grayscale_input_model): + mlmodel = ct.convert(rank4_grayscale_input_model, + inputs=[ct.TensorType(name="input", + shape=(1, 1, 10, 20))], + outputs=[ct.ImageType(name="output_image", + color_layout=ct.colorlayout.GRAYSCALE_FLOAT16)], + minimum_deployment_target=ct.target.macOS13, + compute_precision=ct.precision.FLOAT32, + ) + sample_input = np.random.randint(low=0, high=200, size=(1, 1, 10, 20)).astype(np.float32) + model_output_pil_image = mlmodel.predict({"input": sample_input})['output_image'] + assert isinstance(model_output_pil_image, Image.Image) + assert model_output_pil_image.mode == "F" + model_output_as_numpy = np.array(model_output_pil_image) + reference_output = rank4_grayscale_input_model(torch.from_numpy(sample_input)).detach().numpy() + reference_output = np.squeeze(reference_output) + np.testing.assert_allclose(reference_output, model_output_as_numpy, rtol=1e-2, atol=1e-2) + + 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 468488abe..42650f200 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_custom_ops.py @@ -64,7 +64,7 @@ class TestCustomOp: # Define SSA Custom Op for Sparse MatMul # This will map to `custom_op` in SSA with binding information # to bind input spec to the custom implementation - @register_op(doc_str="Sparse MatMul Layer", is_custom_op=True) + @register_op(is_custom_op=True) class custom_torch_sparse_matmul(Operation): # Defining input spec for current op input_spec = InputSpec( 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 41394c1ea..ab1f7ac3e 100644 --- a/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py +++ b/coremltools/converters/mil/frontend/torch/test/test_torch_ops.py @@ -235,7 +235,7 @@ def test( align_corners, ): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") x_shape, target_size = x_shape_and_target_size theta = torch.rand((x_shape[0], 2, 3)) @@ -289,7 +289,7 @@ def test( align_corners, ): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") params = { "mode": mode, @@ -474,7 +474,7 @@ def test_frobenius_norm(self, shape, backend, keepdim): itertools.product( COMMON_SHAPES, backends, - [1, 2, -1, 3, np.inf, -np.inf], + [-1, 0, 1, 2, 3, np.inf, -np.inf], [True, False] ) ) @@ -486,6 +486,87 @@ def test_number_norm(self, shape, backend, p, keepdim): TorchBaseTest.run_compare_torch(shape, model, backend=backend, places=2) +class TestLinAlgNorms(TorchBaseTest): + def _is_valid_config(self, shape, order, dim): + if isinstance(dim, tuple): + if isinstance(order, int) and (order == 0 or order > 2): + return False + elif isinstance(dim, int): + if order == "fro": + return False + elif dim is None: + if order is not None: + if len(shape) > 2: + return False + elif len(shape) == 2 and not isinstance(order, str) and (order == 0 or order > 2): + return False + elif len(shape) == 1 and isinstance(order, str): + return False + return True + + @pytest.mark.parametrize( + "shape, backend, order, keepdim, dim", + itertools.product( + COMMON_SHAPES, + backends, + [-2, -1, 0, 1, 2, 3, np.inf, -np.inf, "fro", "nuc", None], + [True, False], + [-1, 0, 1, (0, 1), (0, -1), None] + ) + ) + def test_norm(self, shape, backend, order, keepdim, dim): + if not self._is_valid_config(shape, order, dim): + pytest.skip() + if isinstance(order, int) and abs(order) == 2 and ((dim is None and len(shape) == 2) or isinstance(dim, tuple)) : + pytest.xfail("Matrix norm for order 2 and -2 is not implemented") + if (order == "nuc"): + pytest.xfail("Nucleus norm not implemented") + model = ModuleWrapper( + function=torch.linalg.norm, kwargs={'ord': order, 'keepdim': keepdim, 'dim': dim} + ) + TorchBaseTest.run_compare_torch(shape, model, backend=backend, places=2) + + +class TestLinAlgMatrixNorms(TorchBaseTest): + @pytest.mark.parametrize( + "shape, backend, order, keepdim, dim", + itertools.product( + COMMON_SHAPES, + backends, + [-2, -1, 1, 2, np.inf, -np.inf, "fro", "nuc"], + [True, False], + [(0, 1), (0, -1), (1, 2), (0, 2), (2, 3)] + ) + ) + def test_norm(self, shape, backend, order, keepdim, dim): + if dim[-1] > len(shape) - 1: + pytest.skip() + if order == "nuc" or (type(order) != str and abs(order) == 2): + pytest.xfail("Matrix norm for order 2, -2 and nuc is not implemented") + model = ModuleWrapper( + function=torch.linalg.matrix_norm, kwargs={'ord': order, 'keepdim': keepdim, 'dim': dim} + ) + TorchBaseTest.run_compare_torch(shape, model, backend=backend, places=2) + + +class TestLinAlgVectorNorms(TorchBaseTest): + @pytest.mark.parametrize( + "shape, backend, order, keepdim, dim", + itertools.product( + COMMON_SHAPES, + backends, + [-2, -1, 0, 1, 2, np.inf, -np.inf], + [True, False], + [-1, 0, 1, (0, 1), (0, -1), None] + ) + ) + def test_norm(self, shape, backend, order, keepdim, dim): + model = ModuleWrapper( + function=torch.linalg.vector_norm, kwargs={'ord': order, 'keepdim': keepdim, 'dim': dim} + ) + TorchBaseTest.run_compare_torch(shape, model, backend=backend, places=2) + + class TestHardswish(TorchBaseTest): class HardswishModel(nn.Module): @@ -794,8 +875,6 @@ def test_convolution1d( backend, groups=1, ): - if backend[0] == 'mlprogram': - pytest.xfail("Not supported on ML Program backend") class DynamicConv(nn.Module): def __init__(self): @@ -1392,9 +1471,6 @@ def test_avg_pool1d( ): if padding > kernel_size / 2: return - if kernel_size == 1 and stride == 2 and padding == 0 and ceil_mode and input_shape[-1] % 2 == 0: - pytest.xfail(reason="rdar://73894185 (CoreML sometimes returns 'nan's " - "for avg_pool when ceil_mode is True and kernel=1,stride=2,pad=0)") model = nn.AvgPool1d( kernel_size, @@ -1426,10 +1502,7 @@ def test_avg_pool2d( ): if padding > kernel_size / 2: return - if kernel_size == 1 and stride == 2 and padding == 0 and ceil_mode and \ - (input_shape[-2] % 2 == 0 or input_shape[-1] % 2 == 0): - pytest.xfail(reason="rdar://73894185 (CoreML sometimes returns 'nan's " - "for avg_pool when ceil_mode is True and kernel=1,stride=2,pad=0)") + model = nn.AvgPool2d( kernel_size, stride, @@ -1460,14 +1533,10 @@ def test_avg_pool3d( ): if padding > kernel_size / 2: return - if kernel_size == 1 and stride == 2 and padding == 0 and ceil_mode and \ - (input_shape[-3] % 2 == 0 or input_shape[-2] % 2 == 0 or input_shape[-1] % 2 == 0): - pytest.xfail(reason="rdar://73894185 (CoreML sometimes returns 'nan's " - "for avg_pool when ceil_mode is True and kernel=1,stride=2,pad=0)") + if include_pad and ceil_mode and stride > 1: # skip: MIL/CoreML does not support this configuration - # rdar://73723194 - return + pytest.xfail("rdar://73723194 (Support 3D Avg pooling with ceil_mode=True and include_pad = True, in MIL)") model = nn.AvgPool3d( kernel_size, stride, @@ -2375,6 +2444,20 @@ def test_pixel_shuffle(self, batch_size, CHW, r, backend): model = nn.PixelShuffle(upscale_factor=r) self.run_compare_torch(input_shape, model, backend=backend) +class TestPixelUnshuffle(TorchBaseTest): + @pytest.mark.parametrize( + "batch_size, CHW, r, backend", + itertools.product([1, 3], [(1, 4, 4), (3, 2, 3)], [2, 4], backends), + ) + def test_pixel_shuffle(self, batch_size, CHW, r, backend): + if backend[0] == "neuralnetwork": + pytest.skip("pixel_unshuffle only supported in mlprogram backend.") + + C, H, W = CHW + input_shape = (batch_size, C, H * r, W * r) + model = nn.PixelUnshuffle(downscale_factor=r) + self.run_compare_torch(input_shape, model, backend=backend, minimum_deployment_target=ct.target.iOS16) + class TestExpand(TorchBaseTest): @pytest.mark.parametrize( @@ -2714,9 +2797,8 @@ def test_relu6(self, backend, shape): ), ) def test_prelu(self, backend, alpha, shape, single_alpha): - if (backend[0] == "mlprogram" and backend[1] == "fp16"): + if (backend[0] == "mlprogram" and backend[1] == "fp16" or (len(shape) == 5)): pytest.xfail("rdar://92175249 ([MIL] TestActivation::test_prelu[backend=(mlprogram, fp16)] CI failure)") - input_shape = shape num_parameters = input_shape[1] if len(input_shape) >= 2 else 1 if single_alpha: @@ -3111,15 +3193,13 @@ def test(self, backend, shape, dims): class TestTo(TorchBaseTest): @pytest.mark.parametrize( - "use_cpu_for_conversion, backend", itertools.product([True, False], backends,) + "use_cpu_for_conversion, backend", + itertools.product( + [True, False], + backends + ) ) def test_cast_bug(self, use_cpu_for_conversion, 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])") - - if backend[0] == "mlprogram" and use_cpu_for_conversion: - pytest.xfail("numerical mismatch : rdar://78952850") - class TestModel(torch.nn.Module): def forward(self, spans, embedding): spans = spans.float().relu().int() @@ -3132,8 +3212,12 @@ def forward(self, spans, embedding): return sigmoided_scores model = TestModel() - self.run_compare_torch([(1, 21, 2), (1, 6, 384)], model, backend=backend, - use_cpu_for_conversion=use_cpu_for_conversion)# [spans.shape, embedding.shape] + self.run_compare_torch( + [(1, 4, 2), (1, 6, 3)], + model, + backend=backend, + use_cpu_for_conversion=use_cpu_for_conversion, + ) class TestSlice(TorchBaseTest): @pytest.mark.skipif(_python_version() < (3, 6), reason="requires python 3.6") @@ -3316,10 +3400,11 @@ def forward(self, x): class TestTopk(TorchBaseTest): @pytest.mark.parametrize( - "backend, largest, shape_dim_k", + "backend, largest, sort, shape_dim_k", itertools.product( backends, [True, False], + [True, False], [ ((4, 6, 7, 3), -1, 2), ((10, 3, 4), 2, 2), @@ -3327,7 +3412,9 @@ class TestTopk(TorchBaseTest): ], ), ) - def test_topk(self, backend, largest, shape_dim_k): + def test_topk(self, backend, largest, sort, shape_dim_k): + if not sort and backend[0] == "neuralnetwork": + pytest.xfail("iOS16 version topk needed for sort = False") input_shape = shape_dim_k[0] dim = shape_dim_k[1] k = shape_dim_k[2] @@ -3337,18 +3424,23 @@ def __init__(self): super(TopkModel, self).__init__() def forward(self, x): - return torch.topk(x, k, dim=dim, largest=largest) + topk = torch.topk(x, k, dim=dim, largest=largest, sorted=sort) + values, indices = topk.values, topk.indices + if not sort: + values, _ = torch.sort(values, dim=dim) + indices, _ = torch.sort(indices, dim=dim) + return values, indices input_data = torch.rand(input_shape) model = TopkModel() expected_results = model(input_data) - expected_results = [expected_results.values, expected_results.indices] self.run_compare_torch( input_data, model, expected_results=expected_results, input_as_shape=False, backend=backend, + minimum_deployment_target=ct.target.iOS16 if not sort else None, ) class TestLog10(TorchBaseTest): @@ -3926,7 +4018,6 @@ def forward(self, x): 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__() diff --git a/coremltools/converters/mil/frontend/torch/test/testing_utils.py b/coremltools/converters/mil/frontend/torch/test/testing_utils.py index 88b8c5b3f..e17145569 100644 --- a/coremltools/converters/mil/frontend/torch/test/testing_utils.py +++ b/coremltools/converters/mil/frontend/torch/test/testing_utils.py @@ -71,7 +71,8 @@ def convert_to_coreml_inputs(input_description, inputs): def convert_to_mlmodel(model_spec, tensor_inputs, backend=("neuralnetwork", "fp32"), - converter_input_type=None, use_cpu_for_conversion=False): + converter_input_type=None, use_cpu_for_conversion=False, + minimum_deployment_target=None): def _convert_to_inputtype(inputs): if isinstance(inputs, list): return [_convert_to_inputtype(x) for x in inputs] @@ -97,7 +98,8 @@ def _convert_to_inputtype(inputs): compute_unit = ComputeUnit.ALL return ct_convert(model_spec, inputs=inputs, convert_to=backend, - source="pytorch", compute_units=compute_unit) + source="pytorch", compute_units=compute_unit, + minimum_deployment_target=minimum_deployment_target) def generate_input_data(input_size, rand_range=(0, 1)): @@ -143,7 +145,9 @@ def convert_and_compare(input_data, model_spec, expected_results=None, atol=1e-4, backend=("neuralnetwork", "fp32"), converter_input_type=None, - use_cpu_for_conversion=True): + use_cpu_for_conversion=True, + minimum_deployment_target=None + ): """ If expected results is not set, it will by default be set to the flattened output of the torch model. @@ -169,7 +173,8 @@ def convert_and_compare(input_data, model_spec, expected_results = flatten_and_detach_torch_results(expected_results) mlmodel = convert_to_mlmodel(model_spec, input_data, backend=backend, converter_input_type=converter_input_type, - use_cpu_for_conversion=use_cpu_for_conversion) + use_cpu_for_conversion=use_cpu_for_conversion, + minimum_deployment_target=minimum_deployment_target,) coreml_inputs = convert_to_coreml_inputs(mlmodel.input_description, input_data) @@ -208,9 +213,10 @@ def store_testname_with_args(self, request): def run_compare_torch( input_data, model, expected_results=None, places=5, input_as_shape=True, backend=("neuralnetwork", "fp32"), - rand_range=(0.0, 1.0), use_scripting=False, + rand_range=(-1.0, 1.0), use_scripting=False, converter_input_type=None, use_cpu_for_conversion=True, + minimum_deployment_target=None, ): """ Traces a model and runs a numerical test. @@ -235,7 +241,8 @@ def run_compare_torch( atol=10.0 ** -places, backend=backend, converter_input_type=converter_input_type, use_cpu_for_conversion=use_cpu_for_conversion, + minimum_deployment_target=minimum_deployment_target, ) return model_spec, mlmodel, coreml_inputs, coreml_results, \ - TorchBaseTest.testclassname, TorchBaseTest.testmodelname + TorchBaseTest.testclassname, TorchBaseTest.testmodelname \ No newline at end of file diff --git a/coremltools/converters/mil/mil/block.py b/coremltools/converters/mil/mil/block.py index 4a7c5d397..ba36f8904 100644 --- a/coremltools/converters/mil/mil/block.py +++ b/coremltools/converters/mil/mil/block.py @@ -16,6 +16,8 @@ from .var import Var, InternalVar from .visitors.dot_visitor import DotVisitor +from coremltools import _OPSET +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as _target # BLOCK_STACK[-1] is the current block BLOCK_STACK = [] @@ -26,6 +28,17 @@ def curr_block(): raise ValueError("Must call Builder inside an Function" + " or Block") return BLOCK_STACK[-1] +def curr_opset_version(): + block = curr_block() + while not isinstance(block, Function): + block = block.outer_op.enclosing_block + return block.opset_version + +def is_current_opset_version_compatible_with(opset_version): + if curr_opset_version() is None: + return opset_version <= _target.iOS13 + return curr_opset_version() >= opset_version + class InvalidBlockStateError(Exception): pass @@ -225,6 +238,48 @@ def inputs(self): def outputs(self): return self._outputs + def is_var_visible_in_block(self, var, upto_op_with_id=None): + """ + Checks if a var is visible to ops starting from id=`upto_op_with_id` inside the block. + + Var is visible if + - It is the output of a const op, or + - It is the output of "preceding" operations in that block, or + - It is visible in the enclosing block, or + - It is either a block or a function input + + If upto_op_with_id is None, outputs of all operations inside the block are visible to that block. + """ + + if var in self._internal_vars: + return True + + inputs = self.function_inputs if isinstance(self, Function) else self.inputs + if var in inputs: + return True + + idx = len(self.operations) if upto_op_with_id is None else upto_op_with_id + + for i in range(idx-1, -1, -1): + op_outputs = self.operations[i].outputs + if op_outputs is not None and var in op_outputs: + return True + + if self.outer_op is not None: + enclosing_block = self.outer_op.enclosing_block + outer_op_id = enclosing_block.find_op_id_in_block(self.outer_op) + if enclosing_block.is_var_visible_in_block(var, upto_op_with_id=outer_op_id): + return True + + return False + + def find_op_id_in_block(self, target_op): + try: + idx = self.operations.index(target_op) + except ValueError: + raise ValueError("Op {} not found in {}: {}".format(target_op.name, self.name, self)) + return idx + def set_outputs(self, outputs): """ outputs: list[Var] @@ -233,11 +288,8 @@ def set_outputs(self, outputs): raise ValueError("Outputs must be list of Vars") self.validate() - visible_vars = self._visible_vars_from_enclosing_block() - _, visible_vars_in_block = self._visible_vars_in_block() - visible_vars.update(visible_vars_in_block) for ov in outputs: - if ov not in visible_vars: + if not self.is_var_visible_in_block(ov): msg = ( "Var {} is not visible in block {} and thus cannot " + "be a block output.\n{}" @@ -260,141 +312,10 @@ def __enter__(self): return self def __exit__(self, type, value, traceback): + self._propagate_nonreplaceable_vars() global BLOCK_STACK BLOCK_STACK = BLOCK_STACK[:-1] - def _visible_vars_in_block(self, target_op=None, inclusive=True): - """ - Returns - ------- - - index (int) of target_op in self.operations if target_op not None, - undefined otherwise. - - - visible_vars: set[Var] - Vars returned by ops in the block (self) visible (and equal to - if inclusive==True) target_op. If target_op is not found or None, - include all vars output by self.operations. Examples: - - Raises - ------ - - ValueError if target_op not None and not found in self.operations. - - # main(%a: (1, 2, fp32), - # %b: (1, 2, fp32), - # %c: (1, 2, fp32)) { - # block0() { - # %const1: (1, fp32) = const(...) - # %loop:0: (1, 2, fp32), %loop:1: (1, 2, fp32) = \ - # while_loop(loop_vars=(%a, %b)) - # loop_cond(%a.x, %b.x) { - # %blah: (bool) = some_op(x=%a.x, y=%b.x) - # %cond_var: (bool) = some_op2(x=%a.x, y=%blah) - # } -> (%cond_var) - # loop_body(%a.x, %b.x) { - # %add_0: (1, 2, fp32) = add(x=%a.x, y=%b.x) - # } -> (%add_0, %b.x) - # %linear: (1, fp32) = linear(...) - # } -> (%loop:0, %loop:1) - # } - # - - Let V0 and V1 be the set of internal_vars of block0 and loop_cond - block that supplies const vals (for const). - - Ex1: self = block0, target_op = linear. - idx = 2 - visible_vars = {%const1, %loop:0, %loop:1, %linear, V0} - - Ex2: self = loop_cond, target_op = None. - idx = undefined - visible_vars = {%a.x, %b.x, %blah, %cond_var, V1} - - Ex3: self = loop_cond, target_op = some_op. - idx = 0 - visible_vars = {%a.x, %b.x, %blah, V1} - - Ex4: self = loop_cond, target_op = linear. - raises ValueError (linear not found in loop_cond block) - """ - visible_vars = set(self._internal_vars) - if isinstance(self, Function): - # Function inputs - visible_vars.update(tuple(self.inputs.values())) - else: - # Block inputs - visible_vars.update(self.inputs) - - idx = -1 - # find the location of target_op - for i, op in enumerate(self.operations): - if op == target_op: - if inclusive and op.outputs is not None: - visible_vars.update(op.outputs) - return i, visible_vars - # When op is being constructed (e.g.,type_inference), op.outputs - # is None - if op.outputs is not None: - visible_vars.update(op.outputs) - if target_op is not None: - msg = "Op {} not found in {}: {}" - raise ValueError(msg.format(target_op.name, self.name, self)) - return idx, visible_vars - - def _visible_vars_from_enclosing_block(self): - """ - Returns: - - visible_vars: Vars from lexical scopes visible at the beginning of the - block, up to but not including outputs from before_op. Given program: - - # main(%a: (1, 2, fp32), - # %b: (1, 2, fp32), - # %c: (1, 2, fp32)) { - # block0() { - # %const1: (1, fp32) = const(...) - # %loop:0: (1, 2, fp32), %loop:1: (1, 2, fp32) = \ - # while_loop(loop_vars=(%a, %b)) - # loop_cond(%a.x, %b.x) { - # %blah: (bool) = some_op(x=%a.x, y=%b.x) - # %cond_var: (bool) = some_op2(x=%a.x, y=%blah) - # } -> (%cond_var) - # loop_body(%a.x, %b.x) { - # %add_0: (1, 2, fp32) = add(x=%a.x, y=%b.x) - # } -> (%add_0, %b.x) - # %const2: (1, fp32) = const(...) - # } -> (%loop:0, %loop:1) - # } - - Let V0 be the set of internal_vars of block0 block that supplies const - vals (for const). - - Ex1: self = block0 - visible_vars = {%a, %b, %c} (function input) - - Ex2: self = loop_cond. - visible_vars = {%a, %b, %c, %const1, V0} (Note that %const2 is not - part of the set) - """ - visible_vars = set() - - # function inputs are considered external to the block. - if isinstance(self, Function): - # block in function only has function_inputs as from enclosing - # block (Ex1 above). - visible_vars.update(self.function_inputs) - return visible_vars - - if self.outer_op is not None: - enclosing_block = self.outer_op.enclosing_block - vars_at_start = enclosing_block._visible_vars_from_enclosing_block() - visible_vars.update(vars_at_start) - _, visible_vars_in_block = enclosing_block._visible_vars_in_block( - self.outer_op, inclusive=False - ) - visible_vars.update(visible_vars_in_block) - - return visible_vars - def _insert_op_before(self, new_op, before_op=None): """ A private API used by builder. Please use `builder.YOUR_OP(...,before_op)`. @@ -428,26 +349,16 @@ def _insert_op_before(self, new_op, before_op=None): is not visible before op0. """ self.validate() - visible_vars = self._visible_vars_from_enclosing_block() - if before_op is not None: - idx, visible_vars_in_block = self._visible_vars_in_block( - before_op, inclusive=True - ) - visible_vars.update(visible_vars_in_block) - else: - _, visible_vars_in_block = self._visible_vars_in_block() - visible_vars.update(visible_vars_in_block) + + idx = len(self.operations) if before_op is None else self.find_op_id_in_block(before_op) # check inputs are visible for k, v in new_op.inputs.items(): if not isinstance(v, (Var, tuple)): continue - if isinstance(v, Var): - vs = [v] - else: - vs = v + vs = [v] if isinstance(v, Var) else v for s in vs: - if s not in visible_vars: + if not self.is_var_visible_in_block(s, upto_op_with_id=idx): before_op_name = before_op.name if before_op is not None else "None" msg = "Op '{}' input {}={} is not in scope of {} before {}" raise ValueError( @@ -532,6 +443,36 @@ def replace_block_output_var( if isinstance(self, Function): new_var.name = old_var.name + def try_replace_uses_of_var_after_op( + self, + anchor_op, + old_var, + new_var, + no_check_var_types=False + ): + """ + :param anchor_op: Operation + :param old_var: Var + :param new_var: Var + :param no_check_var_types: bool + :return: True if the old_var can be replaced by new_var. False otherwsie. + + This helper function guards the replace_uses_of_var_after_op function, + by first checking if the old_var could be replaced by the new_var. + + 1. If old_var can be replaced by new_var, the replace_uses_of_var_after_op is called, and returns True. + 2. Return False if the replacement is not allow. + """ + if not old_var.can_be_replaced_by_var(new_var): + return False + + self.replace_uses_of_var_after_op( + anchor_op=anchor_op, + old_var=old_var, + new_var=new_var, + no_check_var_types=no_check_var_types, + ) + return True def replace_uses_of_var_after_op( self, @@ -541,6 +482,7 @@ def replace_uses_of_var_after_op( no_check_var_visibility=False, end_op=None, no_check_var_types=False, + force_replace=False, ): """ Replace all uses of `old_var` with `new_var` after `anchor_op`, @@ -618,39 +560,35 @@ def replace_uses_of_var_after_op( %5 = op3(%2) # will continue using %2 """ - if not no_check_var_visibility: - self.validate() - # Get visible vars from enclosing block - visible_vars = self._visible_vars_from_enclosing_block() + if not force_replace and old_var.op is not None and new_var.op is not None: + if not old_var.can_be_replaced_by_var(new_var): + old_nonreplaceable_vars = old_var.nonreplaceable_vars_upstream + new_nonreplaceable_vars = new_var.nonreplaceable_vars_upstream + err_var = None + for _var in old_nonreplaceable_vars: + if _var not in new_nonreplaceable_vars: + err_var = _var + break + msg = ( + "var {} cannot be replaced by {}. Since the nonreplaceable var {} might potentially " + "be removed during the replacement of those vars." + ).format(old_var, new_var, err_var) + raise ValueError(msg) - if anchor_op is not None: - # Get visible vars from the current block - idx, block_vars = self._visible_vars_in_block(anchor_op, inclusive=True) - visible_vars.update(block_vars) + start = self.find_op_id_in_block(anchor_op) + 1 if anchor_op is not None else 0 + end_id = self.find_op_id_in_block(end_op) if end_op is not None else -1 - # start from the next op, excluding `anchor_op` - start = idx + 1 - else: - _, block_vars = self._visible_vars_in_block() - visible_vars.update(block_vars) - - visible_vars.update(self._block_inputs) - visible_vars.update(self._internal_vars) - # Perform replacement from beginning - start = 0 - - if not no_check_var_visibility and new_var not in visible_vars: - msg = ( - "new_var '{}' is not visible in block '{}' at or before " - + "anchor_op '{}'" - ) - anchor_op_name = "None" if anchor_op is None else anchor_op.name - raise ValueError(msg.format(new_var.name, self.name, anchor_op_name)) + if not no_check_var_visibility: + self.validate() - if end_op is not None: - end_id, _ = self._visible_vars_in_block(end_op, inclusive=True) - else: - end_id = -1 + idx = start if anchor_op is not None else len(self.operations) + if not self.is_var_visible_in_block(new_var, upto_op_with_id=idx): + msg = ( + "new_var '{}' is not visible in block '{}' at or before " + + "anchor_op '{}'" + ) + anchor_op_name = "None" if anchor_op is None else anchor_op.name + raise ValueError(msg.format(new_var.name, self.name, anchor_op_name)) if end_id != -1 and end_id < start: msg = "end_op '{}' comes before the anchor_op '{}'" @@ -760,6 +698,18 @@ def operations_for_vars(self, end_vs): used_vars.add(input_var) return used_ops[::-1] + + def _propagate_nonreplaceable_vars(self): + def propagate_nonreplaceable_vars_block(block): + for op in list(block.operations): + for b in op.blocks: + propagate_nonreplaceable_vars_block(b) + if op.outputs is None: + continue + for o in op.outputs: + o._reset_nonreplaceable_vars_upstream() + o._set_nonreplaceable_vars_upstream() + propagate_nonreplaceable_vars_block(self) def indented_str(self, indent=None): if indent is None: @@ -842,11 +792,14 @@ class Function(Block): """ """ - def __init__(self, inputs): + def __init__(self, inputs, opset_version=None): """ inputs: str -> placeholder + opset_version: AvailableTarget enum. Describes the opset version of the function """ self.placeholder_inputs = inputs + self.opset_version = opset_version + # str -> Var self._input_dict = OrderedDict() for k, v in self.placeholder_inputs.items(): @@ -868,6 +821,19 @@ def __init__(self, inputs): @property def inputs(self): return self._input_dict + + @property + def opset_version(self): + return self._opset_version + + @opset_version.setter + def opset_version(self, version): + if not ( + isinstance(version, _target) or + version is None + ): + raise ValueError("opset_version must be type of coremltools.AvailableTarget") + self._opset_version = version def __repr__(self): return self.__str__() @@ -876,8 +842,9 @@ def __str__(self): return self.to_str("function") def to_str(self, func_name="function"): + func_name = func_name + "[{}]".format(_OPSET[self.opset_version]) if len(self._input_dict) == 0: - s = func_name + "() {" + s = func_name + "()" else: inputs = [(in_name, ph) for in_name, ph in self._input_dict.items()] s = func_name + "(" + str(inputs[0][1]) diff --git a/coremltools/converters/mil/mil/builder.py b/coremltools/converters/mil/mil/builder.py index e904a1152..f73c4a2d4 100644 --- a/coremltools/converters/mil/mil/builder.py +++ b/coremltools/converters/mil/mil/builder.py @@ -202,7 +202,7 @@ def TensorSpec(shape, dtype=None): return Placeholder(shape, dtype) @staticmethod - def program(input_specs=None): + def program(input_specs=None, opset_version=None): """ The ``mb.program`` decorator creates a MIL program with a single @@ -212,13 +212,16 @@ def program(input_specs=None): ---------- input_specs: TensorSpec - Describes a tensor. + Describes a tensor. + + opset_version: AvailableTarget enum + Describes the opset version of the program Examples -------- - - >>> @mb.program(input_specs=[mb.TensorSpec(shape=(1,2))]) + >>> import coremltools as ct + >>> @mb.program(input_specs=[mb.TensorSpec(shape=(1,2))], opset_version=ct.target.iOS16) >>> def prog(a): >>> return mb.add(x=a, y=2) @@ -238,7 +241,7 @@ def wrapper(main_block): ) ) input_spec_dict = {k: v for k, v in zip(arg_names, input_specs)} - with Function(input_spec_dict) as func: + with Function(input_spec_dict, opset_version) as func: input_vars = [func.inputs[a] for a in arg_names] outputs = main_block(*input_vars) if isinstance(outputs, tuple): @@ -250,7 +253,3 @@ def wrapper(main_block): return program return wrapper - - -# importing ops triggers installation of all ops into Builder -from .ops import defs as _ops diff --git a/coremltools/converters/mil/mil/input_type.py b/coremltools/converters/mil/mil/input_type.py index 386966b36..a170892e0 100644 --- a/coremltools/converters/mil/mil/input_type.py +++ b/coremltools/converters/mil/mil/input_type.py @@ -40,10 +40,10 @@ def items(self): return self._ordered_dict.items() def __add__(self, default_inputs): - self._default_inputs.extend(default_inputs._default_inputs) + new_order_dict = {k: v for k, v in self._ordered_dict.items()} for k, v in default_inputs._default_inputs: - self._ordered_dict[k] = v - return self + new_order_dict[k] = v + return DefaultInputs(**new_order_dict) class InputSpec: @@ -56,10 +56,11 @@ def __init__(self, **kwargs): self._ordered_dict[k] = v def __add__(self, input_spec): - self._input_types.extend(input_spec._input_types) + new_order_dict = {k: v for k, v in self._ordered_dict.items()} for k, v in input_spec._input_types: - self._ordered_dict[k] = v - return self + new_order_dict[k] = v + return InputSpec(**new_order_dict) + @property def input_types(self): @@ -113,13 +114,9 @@ def validate_inputs(self, op_name, op_type, candidate_kvs): if input_type.const and \ not isinstance(input_type, InternalInputType) \ and var.val is None: - - if var.op and var.op.op_type.startswith("constexpr_"): - pass # Output of constexprs qualifies as const but gets materialized after load time - else: - msg = msg_prefix + \ - 'Input {} must be const at compile time' - raise ValueError(msg.format(name), name, var.name) + msg = msg_prefix + \ + 'Input {} must be const at compile time' + raise ValueError(msg.format(name), name, var.name) if not isinstance(var, InternalVar) and \ not input_type.is_compatible(var): diff --git a/coremltools/converters/mil/mil/operation.py b/coremltools/converters/mil/mil/operation.py index aa44e4643..6a52cdf8f 100644 --- a/coremltools/converters/mil/mil/operation.py +++ b/coremltools/converters/mil/mil/operation.py @@ -298,6 +298,9 @@ def type_value_inference(self, overwrite_output=False): msg = 'value_inference differs for var {} in op {}' if not _is_compatible_symbolic_array(sym_val.val, out_var.sym_val): raise ValueError(msg.format(out_var.name, self.name)) + + for o in self.outputs: + o._set_nonreplaceable_vars_upstream() def _auto_val(self, output_types): """ @@ -499,6 +502,14 @@ def outputs(self): @property def op_type(self): return type(self).__name__ + + @property + def opset_version(self): + op_variants = type(self)._op_variants + opset_versions = sorted(list(op_variants.keys())) + for i in opset_versions: + if op_variants[i] == type(self): + return i def remove_from_block(self): """ diff --git a/coremltools/converters/mil/mil/ops/defs/__init__.py b/coremltools/converters/mil/mil/ops/defs/__init__.py index d3830fb5b..264b925fe 100644 --- a/coremltools/converters/mil/mil/ops/defs/__init__.py +++ b/coremltools/converters/mil/mil/ops/defs/__init__.py @@ -3,206 +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 .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 .constexpr_ops import ( - constexpr_affine_dequantize, - constexpr_cast, - constexpr_lut_to_dense, - constexpr_sparse_to_dense, -) - -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, - space_to_batch, - squeeze, - transpose, - pixel_shuffle, - sliding_windows, +from . import( + iOS15, + iOS16, ) diff --git a/coremltools/converters/mil/mil/ops/defs/_utils.py b/coremltools/converters/mil/mil/ops/defs/_utils.py index 4aa960a8d..cd968daf3 100644 --- a/coremltools/converters/mil/mil/ops/defs/_utils.py +++ b/coremltools/converters/mil/mil/ops/defs/_utils.py @@ -4,6 +4,7 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import math +import numpy as np from coremltools.converters.mil.mil import get_new_symbol, types from coremltools.converters.mil.mil.types.symbolic import is_symbolic @@ -303,3 +304,74 @@ def _update_vec(str, vec, map_char_to_int, index): index = _update_vec(output_str, output_vec, map_char_to_int, index) return input1_vec, input2_vec, output_vec + +def solve_slice_by_index_shape(x_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask): + """ + Helper function to solve the shape of tensor slicing. + """ + ret_shape = [] + + if begin is None or len(begin) == 0: + begin = [None] * len(x_shape) + if end is None or len(end) == 0: + end = [None] * len(x_shape) + + # solve for shape inference + for idx in range(len(x_shape)): + # skip if we want to squeeze the dimension + if squeeze_mask[idx]: + continue + + # for those a[:] cases + if begin_mask[idx] and end_mask[idx]: + if is_symbolic(x_shape[idx]): + if stride[idx] == -1 or stride[idx] == 1: + ret_shape.append(x_shape[idx]) + else: + ret_shape.append(get_new_symbol()) + else: + num = np.ceil(float(x_shape[idx]) / abs(stride[idx])).astype( + np.int32 + ) + ret_shape.append(num) + continue + + # for symbolic case + if is_symbolic(x_shape[idx]): + ret_shape.append(get_new_symbol()) + continue + + # when begin and end are not determined + if begin[idx] is None and not begin_mask[idx]: + ret_shape.append(get_new_symbol()) + continue + if end[idx] is None and not end_mask[idx]: + ret_shape.append(get_new_symbol()) + continue + + # parse negative dimention + if begin[idx] is not None and begin[idx] < 0: + begin[idx] = max(0, begin[idx] + x_shape[idx]) + if end[idx] is not None and end[idx] < 0: + end[idx] = max(0, end[idx] + x_shape[idx]) + + # compute shape + low, high = [0, x_shape[idx]] if stride[idx] > 0 else [-1, x_shape[idx] - 1] + begin_idx, end_idx = ( + [begin[idx], end[idx]] if stride[idx] > 0 else [end[idx], begin[idx]] + ) + is_begin_mask, is_end_mask = ( + [begin_mask[idx], end_mask[idx]] + if stride[idx] > 0 + else [end_mask[idx], begin_mask[idx]] + ) + if is_begin_mask: + begin_idx = low + end_idx = high if is_end_mask else min(end_idx, high) + num = np.ceil(float(end_idx - begin_idx) / abs(stride[idx])).astype( + np.int32 + ) + ret_shape.append(max(0, num)) + + return ret_shape + diff --git a/coremltools/converters/mil/mil/ops/defs/iOS15/__init__.py b/coremltools/converters/mil/mil/ops/defs/iOS15/__init__.py new file mode 100644 index 000000000..c987a4ad0 --- /dev/null +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/__init__.py @@ -0,0 +1,204 @@ +# Copyright (c) 2022, 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._deployment_compatibility import AvailableTarget as target + +_IOS15_TARGET = target.iOS15 + +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 ( + 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, + space_to_batch, + squeeze, + transpose, + pixel_shuffle, + sliding_windows, +) \ No newline at end of file diff --git a/coremltools/converters/mil/mil/ops/defs/activation.py b/coremltools/converters/mil/mil/ops/defs/iOS15/activation.py similarity index 97% rename from coremltools/converters/mil/mil/ops/defs/activation.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/activation.py index 2a8fc843c..028defd12 100644 --- a/coremltools/converters/mil/mil/ops/defs/activation.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/activation.py @@ -2,6 +2,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 + import math import numpy as np @@ -21,7 +22,7 @@ from .elementwise_unary import elementwise_unary -@register_op(doc_str="") +@register_op class clamped_relu(Operation): """ If ``x >= 0`` return elementwise ``min(beta, x)``, otherwise return @@ -62,7 +63,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class elu(Operation): """ If ``x > 0`` return elementwise ``x``, otherwise return ``alpha * (e^x - 1)``. @@ -106,7 +107,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class gelu(Operation): """ Return the elementwise Gaussian error linear unit activation function for ``x``. @@ -181,7 +182,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class leaky_relu(Operation): """ If ``x >= 0`` apply ``x`` elementwise, otherwise apply ``alpha * x`` elementwise. @@ -225,7 +226,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class linear_activation(Operation): """ Apply elementwise ``x * alpha + beta``. @@ -269,7 +270,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class prelu(Operation): """ Where ``i = 1 ... C``, if ``x_i > 0``, return ``x_i`` , otherwise return ``alpha_i * x_i``. @@ -326,7 +327,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class relu(elementwise_unary): """ Return elementwise-applied rectified linear activation: ``min(x, 0)``. @@ -353,7 +354,7 @@ def value_inference(self): return np.maximum(self.x.val, 0) -@register_op(doc_str="") +@register_op class relu6(elementwise_unary): """ Return elementwise-applied rectified linear activation: ``min(max(x, 0), 6)``. @@ -380,7 +381,7 @@ def value_inference(self): return np.minimum(np.maximum(self.x.val, 0), 6) -@register_op(doc_str="") +@register_op class scaled_tanh(Operation): """ Return ``alpha * tanh(beta * x)`` elementwise. @@ -427,7 +428,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class sigmoid(elementwise_unary): """ Return ``sigmoid(x)`` elementwise. @@ -454,7 +455,7 @@ def value_inference(self): return 1 / (1 + np.exp(-self.x.val)) -@register_op(doc_str="") +@register_op class sigmoid_hard(Operation): """ Return ``min( max( alpha * x + beta, 0 ), 1 )`` elementwise. @@ -500,9 +501,9 @@ def value_inference(self): def type_inference(self): return self.x.sym_type - - -@register_op(doc_str="") + + +@register_op() class silu(Operation): """ Sigmoid Linear Unit, elementwise apply the SiLU or Swish operation ``x * sigmoid(x)``. @@ -529,7 +530,7 @@ def type_inference(self): return types.tensor(self.x.dtype, tuple(self.x.shape)) -@register_op(doc_str="") +@register_op class softplus(elementwise_unary): """ Return ``log( 1 + e^x )`` elementwise. @@ -556,7 +557,7 @@ def value_inference(self): return np.log(1 + np.exp(-np.abs(self.x.val))) + np.maximum(self.x.val, 0) -@register_op(doc_str="") +@register_op class softplus_parametric(Operation): """ Return ``alpha_i * log( 1 + e^( beta_i * x_i ) )``, where ``i = 1 ... C``. @@ -615,7 +616,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class softmax(Operation): """ Return ``exp(x) / tf.reduce_sum(tf.exp(x), axis)``. @@ -660,8 +661,7 @@ def value_inference(self): temp = np.exp(x - max_vals) return temp / np.sum(temp, axis=axis, keepdims=True) - -@register_op(doc_str="") +@register_op class softsign(elementwise_unary): """ Return ``x / ( 1 + |x| )`` applied elementwise. @@ -688,7 +688,7 @@ def value_inference(self): return self.x.val / (1 + np.abs(self.x.val)) -@register_op(doc_str="") +@register_op class thresholded_relu(Operation): """ Return ``x`` if ``x >= alpha``, otherwise return ``0``. diff --git a/coremltools/converters/mil/mil/ops/defs/classify.py b/coremltools/converters/mil/mil/ops/defs/iOS15/classify.py similarity index 99% rename from coremltools/converters/mil/mil/ops/defs/classify.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/classify.py index a6d63bafe..2897191ec 100644 --- a/coremltools/converters/mil/mil/ops/defs/classify.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/classify.py @@ -15,7 +15,7 @@ from coremltools.converters.mil.mil.types.symbolic import any_symbolic -@register_op(doc_str="") +@register_op class classify(Operation): """ Presence of this op indicates that the model is of type classifier, diff --git a/coremltools/converters/mil/mil/ops/defs/control_flow.py b/coremltools/converters/mil/mil/ops/defs/iOS15/control_flow.py similarity index 97% rename from coremltools/converters/mil/mil/ops/defs/control_flow.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/control_flow.py index 13a386085..f8c82b965 100644 --- a/coremltools/converters/mil/mil/ops/defs/control_flow.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/control_flow.py @@ -48,7 +48,7 @@ from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op -@register_op(doc_str="") +@register_op class cond(Operation): """ Perform a conditional execution. The return types must be identical @@ -231,14 +231,14 @@ def _get_type_val(self, value): return builtin_type, value -@register_op(doc_str="") +@register_op class const(Const): def __init__(self, **kwargs): super().__init__(**kwargs) # Internal const can have symbolic value (for testing purpose) -@register_op(doc_str="") +@register_op class _const_symbolic(const): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -253,7 +253,7 @@ def value_inference(self): return val -@register_op(doc_str="") +@register_op class select(Operation): """ Return the elements selected from either ``a`` or ``b`` depending on the ``cond``. @@ -318,7 +318,7 @@ def value_inference(self): return np.where(self.cond.val, self.a.val, self.b.val) -@register_op(doc_str="") +@register_op class while_loop(Operation): """ Perform the body repeatedly while the condition ``cond`` is true. @@ -448,26 +448,20 @@ def build_nested_blocks(self): return block_inputs = tuple(copy.copy(v) for v in self.loop_vars) - _, visible_vars = self.enclosing_block._visible_vars_in_block() - name_count = {v.name: 1 for v in visible_vars} - seen = set() # Avoid using same name among block inputs + name_count = {v.name: 0 for v in block_inputs} for v in block_inputs: v._op = None v.op_output_idx = None v._child_ops = list() # Get unique name + old_v_name = v.name - if v.name in name_count: - v.name = v.name + "_x" + str(name_count[v.name]) - name_count[old_v_name] += 1 - else: - v.name = v.name + "_x0" - name_count[old_v_name] = 0 + v.name = v.name + "_x" + str(name_count[v.name]) + name_count[old_v_name] += 1 v._sym_val = v._sym_val v.consuming_blocks = list() - seen.add(v.name) cond_block, body_block, exit_vars = self._build_block(block_inputs) @@ -546,7 +540,7 @@ def type_inference(self): return tuple(v.sym_type for v in self.blocks[1].outputs) -@register_op(doc_str="") +@register_op class make_list(Operation): """ Create a list of tensor elements. The elements should have the same shape. @@ -624,7 +618,7 @@ def type_inference(self): ) -@register_op(doc_str="") +@register_op class list_length(Operation): """ Return the length of ``ls``. @@ -658,7 +652,7 @@ def value_inference(self): raise NotImplementedError() -@register_op(doc_str="") +@register_op class list_write(Operation): """ Write a value into index ``index`` of ``ls``. @@ -716,7 +710,7 @@ def type_inference(self): return self.ls.sym_type -@register_op(doc_str="") +@register_op class list_read(Operation): """ Read the value at location ``index`` of ``ls``. @@ -757,7 +751,7 @@ def type_inference(self): return list_elem_type -@register_op(doc_str="") +@register_op class list_gather(Operation): """ Return selected values in ``ls`` as a packed ``Tensor``. @@ -802,7 +796,7 @@ def type_inference(self): return types.tensor(dtype, tuple(ret_shape)) -@register_op(doc_str="") +@register_op class list_scatter(Operation): """ Scatter ``values`` to ``ls`` at locations ``indices``. diff --git a/coremltools/converters/mil/mil/ops/defs/conv.py b/coremltools/converters/mil/mil/ops/defs/iOS15/conv.py similarity index 99% rename from coremltools/converters/mil/mil/ops/defs/conv.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/conv.py index 1629882b5..b5c5a9617 100644 --- a/coremltools/converters/mil/mil/ops/defs/conv.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/conv.py @@ -19,7 +19,7 @@ from coremltools.converters.mil.mil.ops.defs._utils import spatial_dimensions_out_shape -@register_op(doc_str="") +@register_op class conv(Operation): """ Perform convolution over input. Supports 1-D, 2-D, and 3-D convolution. @@ -191,7 +191,7 @@ def type_inference(self): return types.tensor(self.x.dtype, tuple(retshape)) -@register_op(doc_str="") +@register_op class conv_quantized(conv): """ Note: This is experimental and may change in the future. @@ -250,7 +250,7 @@ def __init__(self, **kwargs): super().__init__(**kwargs) -@register_op(doc_str="") +@register_op class conv_transpose(Operation): """ Perform transposed convolution (also known as deconvolution and fractionally diff --git a/coremltools/converters/mil/mil/ops/defs/elementwise_binary.py b/coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_binary.py similarity index 96% rename from coremltools/converters/mil/mil/ops/defs/elementwise_binary.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_binary.py index 1d2473f32..1b95224b6 100644 --- a/coremltools/converters/mil/mil/ops/defs/elementwise_binary.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_binary.py @@ -5,8 +5,6 @@ import numpy as np import operator -from ._op_reqs import register_op -from ._utils import promoted_primitive_type, broadcast_shapes from coremltools.converters.mil.mil import ( InputSpec, Operation, @@ -15,6 +13,8 @@ types ) from coremltools.converters.mil.mil.operation import VALUE +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs._utils import promoted_primitive_type, broadcast_shapes class elementwise_binary(Operation): @@ -82,7 +82,7 @@ def _cast_check_value_inferene(self, a, b): """ -@register_op(doc_str="") +@register_op class add(elementwise_binary): """ Return ``x + y`` element-wise with @@ -112,7 +112,7 @@ def get_operator(self): return operator.add -@register_op(doc_str="") +@register_op class equal(elementwise_binary): """ Return the truth value of ``x == y`` element-wise with @@ -147,7 +147,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class floor_div(elementwise_binary): """ Return ``x / y`` element-wise with @@ -179,7 +179,7 @@ def get_operator(self): return operator.floordiv -@register_op(doc_str="") +@register_op class greater(elementwise_binary): """ Return the truth value of ``x > y`` element-wise with @@ -214,7 +214,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class greater_equal(elementwise_binary): """ Return the truth value of ``x >= y`` element-wise with @@ -249,7 +249,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class less(elementwise_binary): """ Return the truth value of ``x < y`` element-wise with @@ -284,7 +284,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class less_equal(elementwise_binary): """ Return the truth value of ``x <= y`` element-wise with @@ -319,7 +319,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class logical_and(elementwise_binary): """ Return the truth value of ``x AND y`` element-wise with @@ -354,7 +354,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class logical_or(elementwise_binary): """ Return the truth value of ``x OR y`` element-wise with @@ -389,7 +389,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class logical_xor(elementwise_binary): """ Return the truth value of ``x XOR y`` element-wise with @@ -424,7 +424,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class maximum(elementwise_binary): """ Return ``x > y ? x : y`` element-wise with @@ -455,7 +455,7 @@ def get_operator(self): return np.maximum -@register_op(doc_str="") +@register_op class minimum(elementwise_binary): """ Return ``x > y ? y : x`` element-wise with @@ -486,7 +486,7 @@ def get_operator(self): return np.minimum -@register_op(doc_str="") +@register_op class mod(elementwise_binary): """ Return ``x % y`` element-wise with @@ -517,7 +517,7 @@ def get_operator(self): return operator.mod -@register_op(doc_str="") +@register_op class mul(elementwise_binary): """ Return ``x * y`` element-wise with @@ -548,7 +548,7 @@ def get_operator(self): return operator.mul -@register_op(doc_str="") +@register_op class not_equal(elementwise_binary): """ Return the truth value of ``x != y`` element-wise with @@ -583,7 +583,7 @@ def get_dtype(self, promoted_dtype): return types.bool -@register_op(doc_str="") +@register_op class real_div(elementwise_binary): """ Return ``x / y`` element-wise with @@ -621,7 +621,7 @@ def get_operator(self): return operator.truediv -@register_op(doc_str="") +@register_op class pow(elementwise_binary): """ Return ``x ^ y`` element-wise with @@ -652,7 +652,7 @@ def get_operator(self): return operator.pow -@register_op(doc_str="") +@register_op class sub(elementwise_binary): """ Return ``x - y`` element-wise with diff --git a/coremltools/converters/mil/mil/ops/defs/elementwise_unary.py b/coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_unary.py similarity index 96% rename from coremltools/converters/mil/mil/ops/defs/elementwise_unary.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_unary.py index a6968d300..046f7265c 100644 --- a/coremltools/converters/mil/mil/ops/defs/elementwise_unary.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/elementwise_unary.py @@ -43,7 +43,7 @@ def type_inference(self): Elementwise unary op implementation(s) """ -@register_op(doc_str="") +@register_op class abs(elementwise_unary): """ Return the absolute values of the input ``x``, element-wise. @@ -71,7 +71,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class acos(elementwise_unary): """ Return the inverse cosine values of the input ``x``, element-wise. @@ -99,7 +99,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class asin(elementwise_unary): """ Return the inverse sine of the input ``x``, element-wise. @@ -127,7 +127,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class atan(elementwise_unary): """ Return the inverse tangent of the input ``x``, element-wise. @@ -155,7 +155,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class atanh(elementwise_unary): """ Return the inverse hyperbolic tangent values of the input @@ -184,7 +184,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class ceil(elementwise_unary): """ Return the ceil values of the input ``x``, element-wise. @@ -212,7 +212,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class clip(Operation): """ Clip the values in the input ``x`` to ``[alpha, beta]``, element-wise. @@ -252,7 +252,7 @@ def value_inference(self): return np.minimum(np.maximum(self.x.val, self.alpha.val), self.beta.val) -@register_op(doc_str="") +@register_op class cos(elementwise_unary): """ Return cosine of ``x`` element-wise. Input domain is ``(-inf, inf)`` and @@ -280,7 +280,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class cosh(elementwise_unary): """ Return hyperbolic cosine of the input ``x``, element-wise. @@ -308,7 +308,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class erf(elementwise_unary): """ Return the gauss error function of the input ``x``, element-wise. @@ -336,7 +336,7 @@ def value_inference(self): return erf_vector_function(self.x.val) -@register_op(doc_str="") +@register_op class exp(elementwise_unary): """ Return e^x, element-wise. @@ -364,7 +364,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class exp2(elementwise_unary): """ Return 2^x, element-wise. @@ -392,7 +392,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class floor(elementwise_unary): """ Return the floor of the input ``x``, element-wise, the same as rounding @@ -421,7 +421,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class inverse(Operation): """ Return the reciprocal value of the input ``x``, element-wise. @@ -465,7 +465,7 @@ def value_inference(self): return np.reciprocal(self.x.val + self.epsilon.val) -@register_op(doc_str="") +@register_op class log(Operation): """ Return the natural logarithm value of the input ``x``, element-wise. @@ -507,7 +507,7 @@ def value_inference(self): return np.log(self.x.val + self.epsilon.val) -@register_op(doc_str="") +@register_op class logical_not(elementwise_unary): """ Return the value of NOT the input ``x``, element-wise. (``1`` for true, ``0`` @@ -536,7 +536,7 @@ def value_inference(self): return np.logical_not(self.x.val) -@register_op(doc_str="") +@register_op class round(elementwise_unary): """ Return the round value of the input ``x`` to nearest integer, element-wise. @@ -565,7 +565,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class rsqrt(Operation): """ Return the reciprocal value of the square root of the input ``x``, element-wise. @@ -610,7 +610,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class sign(elementwise_unary): """ Return the sign value of the input ``x``, element-wise. @@ -640,7 +640,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class sin(elementwise_unary): """ Return the sine value of the input ``x``, element-wise. @@ -668,7 +668,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class sinh(elementwise_unary): """ Return the hyperbolic sine value of the input ``x``, element-wise. @@ -696,7 +696,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class sqrt(elementwise_unary): """ Returns the square root value of the input ``x``, element-wise. @@ -724,7 +724,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class square(elementwise_unary): """ Return ``x^2``, element-wise. @@ -751,7 +751,7 @@ def value_inference(self): return np.square(self.x.val) -@register_op(doc_str="") +@register_op class tan(elementwise_unary): """ Return the tangent value of the input ``x``, element-wise. Both input and output @@ -780,7 +780,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class tanh(elementwise_unary): """ Return the hyperbolic tangent value of the input ``x``, element-wise. Both input @@ -809,7 +809,7 @@ def value_inference(self): return _maintain_shape(self.x.val, result) -@register_op(doc_str="") +@register_op class threshold(Operation): """ Set a lower bound ``alpha`` to the values in the input ``x``, element-wise. @@ -845,7 +845,7 @@ def value_inference(self): return np.maximum(self.x.val, self.alpha.val) -@register_op(doc_str="") +@register_op class cast(Operation): """ Cast the input ``x`` to the new type ``dtype``. diff --git a/coremltools/converters/mil/mil/ops/defs/image_resizing.py b/coremltools/converters/mil/mil/ops/defs/iOS15/image_resizing.py similarity index 96% rename from coremltools/converters/mil/mil/ops/defs/image_resizing.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/image_resizing.py index 25e02ff9b..7aea4e3a6 100644 --- a/coremltools/converters/mil/mil/ops/defs/image_resizing.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/image_resizing.py @@ -1,10 +1,10 @@ -# Copyright (c) 2020, Apple Inc. All rights reserved. +# Copyright (c) 2022, 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 ._op_reqs import register_op from coremltools.converters.mil.mil import ( BoolInputType, DefaultInputs, @@ -15,148 +15,17 @@ IntTensorInputType, IntOrFloatInputType, Operation, + ScalarOrTensorInputType, StringInputType, TensorInputType, types, ) +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs.iOS15 import _IOS15_TARGET from coremltools.converters.mil.mil.types.symbolic import is_symbolic -@register_op(doc_str="") -class affine(Operation): - """ - Apply a linear affine transform to the input 2D image tensor. The value at the - ``(x, y)`` (i.e., ``(w, h)``) coordinate of the output is computed by first computing - the coordinates ``x’`` and ``y’`` with the following equation, and then computing the - value at the coordinate ``(x’,y’)`` in the input image using either bilinear or - nearest neighbor interpolation. If the ``(x’, y’)`` point falls outside the input - image, then padding information is used to compute the value. - - .. sourcecode:: python - - * x’ = a0 * x + a1 * y + a2 - * y’ = b0 * x + b1 * y + b2 - - - Parameters - ---------- - x: tensor<[B, C, H1, W1], T> - * Must be rank ``4``. - transform_matrix: tensor<[D, 6], T> - * Must be rank ``2``. - * ``D`` can be either ``B`` or 1. - * If ``D == B``, there is a separate transform matrix for each batch. - * If ``D == 1``, the same matrix is used for all input batches. - * For each batch: ``[a0, a1, a2, b0, b1, b2]``. - output_height: const - * Target output height - output_width: const - * Target output width - sampling_mode: const - * Allowed values: ``"bilinear"`` - padding_mode: const - * Allowed values: ``"constant"``. - * Note that the following example is 1D case for brevity. - The op supports only 2D image input. - * If ``padding_mode == "constant"``: - * The input image is assumed to be padded with the padding_value. - * For example, ``|1, 2, 3| -> |0, 0, 0, 1, 2, 3, 0, 0, 0|``. - padding_value: const - * Currently non-zero values are not supported. - * To be used only when ``padding_mode == "constant"``, ignored in other cases. - coordinates_mode: const - * Allowed values: ``"normalized_minus_one_to_one"`` - * If ``coordinates_mode == "normalized_minus_one_to_one"``, in-image values are ``[-1, 1]``. - * For example, if ``coordinates_mode == "normalized_minus_one_to_one"``, - the in range values are ``[-1, 1]``. That is: - * ``(-1, -1)``, i.e. ``(w=-1, h=-1)``, corresponds to the top-left pixel. - * ``(1, -1)``, i.e. ``(w=1, h=-1)``, corresponds to the top-right pixel. - * ``(-1, 1)``, i.e. ``(w=-1, h=1)``, corresponds to the bottom-left pixel. - * ``(1, 1)``, i.e. ``(w=1, h=1)``, corresponds to the bottom-right pixel. - align_corners: const - * Currently ``align_corners=False`` is not supported. - * To be used only when ``coordinates_mode != unnormalized``, ignored otherwise. - * if ``align_corners == True``, the extrema coordinates correspond - to the center of the first and last corner pixels. - * if ``align_corners == False``, the extrema coordinates correspond - to the edge of the first and last corner pixels. - - Returns - ------- - tensor<[B, C, output_height, output_width], T> - - Attributes - ---------- - T: fp16, fp32 - """ - - input_spec = InputSpec( - x=TensorInputType(), - transform_matrix=TensorInputType(), - output_height=IntInputType(const=True), - output_width=IntInputType(const=True), - sampling_mode=StringInputType(const=True), - padding_mode=StringInputType(const=True), - padding_value=FloatInputType(const=True), - coordinates_mode=StringInputType(const=True), - align_corners=BoolInputType(const=True), - ) - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def type_inference(self): - if self.x.rank != 4: - raise ValueError( - 'input "x" to the "affine" op must be a rank 4 tensor. ' - "Got rank {} tensor of shape {}".format( - self.x.rank, self.x.shape - ) - ) - if self.transform_matrix.rank != 2: - raise ValueError( - 'input "transform_matrix" to the "affine" op must be a rank 2 tensor. ' - "Got rank {} tensor of shape {}".format( - self.transform_matrix.rank, self.transform_matrix.shape - ) - ) - if self.sampling_mode.val.lower() != "bilinear": - raise NotImplementedError( - 'input "sampling_mode" to the "affine" not implemented. ' - 'Got "{}"'.format(self.sampling_mode.val) - ) - if self.coordinates_mode.val.lower() != "normalized_minus_one_to_one": - raise NotImplementedError( - 'input "coordinates_mode" to the "affine" not implemented. ' - 'Got "{}"'.format(self.coordinates_mode.val) - ) - if self.padding_mode.val.lower() != "constant" or self.padding_value.val != 0.0: - raise NotImplementedError( - 'input "padding_mode" to the "affine" not implemented. ' - 'Got "{}" with "padding_value={}"'.format( - self.padding_mode.val, self.padding_value.val - ) - ) - - input_shape = self.x.shape - transform_matrix_shape = self.transform_matrix.shape - if ( - not is_symbolic(transform_matrix_shape[-1]) - and transform_matrix_shape[-1] != 6 - ): - raise ValueError( - 'input "transform_matrix" to the "affine" op last dimension must be 6 ' - "[a0, a1, a2, b0, b1, b2], " - "Got {} for last dimension".format(transform_matrix_shape[-1]) - ) - - ret_shape = list(input_shape) - ret_shape[2] = self.output_height.val - ret_shape[3] = self.output_width.val - return types.tensor(self.x.dtype, tuple(ret_shape)) - - -@register_op(doc_str="") +@register_op class upsample_nearest_neighbor(Operation): """ Upsample the spatial dimensions (last two dimensions) of the input @@ -188,8 +57,16 @@ class upsample_nearest_neighbor(Operation): input_spec = InputSpec( x=TensorInputType(), - scale_factor_height=IntOrFloatInputType(const=True, optional=True), - scale_factor_width=IntOrFloatInputType(const=True, optional=True), + scale_factor_height=ScalarOrTensorInputType( + const=True, + optional=True, + type_domain=(np.int32, np.float32) + ), + scale_factor_width=ScalarOrTensorInputType( + const=True, + optional=True, + type_domain=(np.int32, np.float32) + ), ) def default_inputs(self): @@ -213,188 +90,56 @@ def type_inference(self): return types.tensor(self.x.dtype, ret_shape) -@register_op(doc_str="") -class resample(Operation): +@register_op +class resize_nearest_neighbor(Operation): """ - Resample the input image tensor ``x`` at the ``coordinates``. - Resampling is required if the coordinates do not correspond to exact - pixels in the input image. The ``sampling_mode ``determines - the algorithm used for resampling and computing the values. + Resize the spatial (last two) dimensions to the specified target size + using nearest neighbor interpolation. Although this op is similar to + ``upsample_nearest_neighbor``, ``resize_nearest_neighbor`` works with + a target size rather than with scale factors. Parameters ---------- - x: tensor<[B, C, H1, W1], T> - * Must be rank ``4``. - coordinates: tensor<[B, H2, W2, 2], U> - * Must be rank ``4``. - * Coordinates are provided in the order ``(x, y)`` (i.e. ``(w, h)``). - * The value of each output location ``output[b, c, h, w]`` is calculated - by sampling from the input image ``x[b, c, :, :]``. - * The pixel at the ``(x, y)`` location corresponds to the length-2 - vector: ``coordinates[b, h, w, :]``. - * Coordinate (normalized or unnormalized) should be specified according - to ``coordinates_mode``. - sampling_mode: const - * Allowed values: ``"bilinear" , "nearest"`` - padding_mode: const - * Allowed values: ``"constant"``, ``"border"``, ``"reflection"``, ``"symmetric"`` - * Note that the following example is 1D case for brevity. - The op supports only 2D image input. - * If ``padding_mode == "constant"``: - * The input image is assumed to be padded with the ``padding_value``. - * For example: ``|1, 2, 3| -> |0, 0, 0, 1, 2, 3, 0, 0, 0|`` - * if ``padding_mode == "border"``: - * The input image is assumed to be padded with the values replicated - from the values at the edge. This is also referred to as the - "clamped" or "replication" mode, since the padded values are - clamped to the border values. - * For example: ``|1, 2, 3| -> |1, 1, 1, 1, 2, 3, 3, 3, 3|`` - * If ``padding_mode == "reflection"``: - * The border values are reflected, *not* including the values at the edge/border. - * For example: ``|1, 2, 3| -> |2, 3, 2, 1, 2, 3, 2, 1, 2|`` - * If ``padding_mode == "symmetric"``: - * Values are reflected, including the border/edge values. - * For example: ``|1, 2, 3| -> |3, 2, 1 , 1, 2, 3, 3, 2, 1|`` - padding_value: const - * To be used only when ``padding_mode == "constant"``, ignored in other cases. - coordinates_mode: const - * Allowed values: ``"unnormalized"``, ``"normalized_minus_one_to_one"``, - ``"normalized_zero_to_one"`` - * If ``coordinates_mode == "unnormalized"``, the coordinates input values - are interpreted to be in range ``[0, W - 1] / [0, H - 1]``, which - corresponds to the in-image point. - * If ``coordinates_mode == "normalized_minus_one_to_one"``, - the in-image values are ``[-1, 1]``. - * If ``coordinates_mode == "normalized_zero_to_one"``, - in-image values are ``[0, 1]``. - * For example, if ``coordinates_mode == "normalized_minus_one_to_one"``, - the in range values are [-1, 1]. That is: - * ``(-1, -1)``, i.e. ``(w=-1, h=-1)``, corresponds to the top-left pixel. - * ``(1, -1)``, i.e. ``(w=1, h=-1)``, corresponds to the top-right pixel. - * ``(-1, 1)``, i.e. ``(w=-1, h=1)``, corresponds to the bottom-left pixel. - * ``(1, 1)``, i.e. ``(w=1, h=1)``, corresponds to the bottom-right pixel. - align_corners: const - * If ``align_corners == True``, the extrema coordinates correspond - to the center of the first and last corner pixels. - * If ``align_corners == False``, the extrema coordinates correspond - to the edge of the first and last corner pixels. + x: tensor<[\*D, H1, W1], T> (Required) + * Must be at least rank ``3``. + target_size_height: const (Required) + * Target spatial size for the height dimension (``axis=-2``). + target_size_width: const (Required) + * Target spatial size for the width dimension (``axis=-1``). + + Notes + ----- + See ``resize_bilinear`` for examples. + + See Also + -------- + resize_bilinear Returns ------- - tensor<[B, C, H2, W2], T> + tensor<[\*D, H2, W2], T> + * Tensor with same type as the input. + * ``H2`` = ``target_size_height``. + * ``W2`` = ``target_size_width``. Attributes ---------- T: fp16, fp32 - U: fp32, i32 """ input_spec = InputSpec( x=TensorInputType(), - coordinates=TensorInputType(), - sampling_mode=StringInputType(const=True), - padding_mode=StringInputType(const=True), - padding_value=FloatInputType(const=True), - coordinates_mode=StringInputType(const=True), - align_corners=BoolInputType(const=True), + target_size_height=IntInputType(const=True), + target_size_width=IntInputType(const=True), ) def __init__(self, **kwargs): super().__init__(**kwargs) def type_inference(self): - if self.x.rank != 4: - raise ValueError( - 'input "x" to the "resample" op must be a rank 4 tensor. ' - "Got rank {} tensor of shape {}".format( - self.x.rank, self.x.shape - ) - ) - if self.coordinates.rank != 4: + if self.x.rank < 3: raise ValueError( - 'input "coordinates" to the "resample" op must be a rank 4 tensor. ' - "Got rank {} tensor of shape {}".format( - self.coordinates.rank, self.coordinates.shape - ) - ) - - input_shape = self.x.shape - coord_shape = self.coordinates.shape - if ( - not is_symbolic(input_shape[0]) - and not is_symbolic(coord_shape[0]) - and input_shape[0] != coord_shape[0] - ): - raise ValueError( - 'input "x" and "coordinates" to the "resample" must agree on ' - "dimension of batch size: {} vs. {}".format( - input_shape[0], coord_shape[0] - ) - ) - if not is_symbolic(coord_shape[-1]) and coord_shape[-1] != 2: - raise ValueError( - 'input "coordinates" to the "resample" op last dimension must be 2. ' - "Got {} for last dimension".format( - coord_shape[-1] - ) - ) - - ret_shape = list(input_shape) - ret_shape[2] = coord_shape[1] # Output height - ret_shape[3] = coord_shape[2] # Output width - return types.tensor(self.x.dtype, tuple(ret_shape)) - - -@register_op(doc_str="") -class resize_nearest_neighbor(Operation): - """ - Resize the spatial (last two) dimensions to the specified target size - using nearest neighbor interpolation. Although this op is similar to - ``upsample_nearest_neighbor``, ``resize_nearest_neighbor`` works with - a target size rather than with scale factors. - - Parameters - ---------- - x: tensor<[\*D, H1, W1], T> (Required) - * Must be at least rank ``3``. - target_size_height: const (Required) - * Target spatial size for the height dimension (``axis=-2``). - target_size_width: const (Required) - * Target spatial size for the width dimension (``axis=-1``). - - Notes - ----- - See ``resize_bilinear`` for examples. - - See Also - -------- - resize_bilinear - - Returns - ------- - tensor<[\*D, H2, W2], T> - * Tensor with same type as the input. - * ``H2`` = ``target_size_height``. - * ``W2`` = ``target_size_width``. - - Attributes - ---------- - T: fp16, fp32 - """ - - input_spec = InputSpec( - x=TensorInputType(), - target_size_height=IntInputType(const=True), - target_size_width=IntInputType(const=True), - ) - - def __init__(self, **kwargs): - super().__init__(**kwargs) - - def type_inference(self): - if self.x.rank < 3: - raise ValueError( - 'input to the "resize_nearest_neighbor" op must have rank at least 3' + 'input to the "resize_nearest_neighbor" op must have rank at least 3' ) ret_shape = list(self.x.shape) @@ -403,7 +148,7 @@ def type_inference(self): return types.tensor(self.x.dtype, ret_shape) -@register_op(doc_str="") +@register_op class upsample_bilinear(Operation): """ Upsample the spatial dimensions (last two dimensions) of the input @@ -485,10 +230,16 @@ class upsample_bilinear(Operation): input_spec = InputSpec( x=TensorInputType(), - scale_factor_height=IntOrFloatInputType(const=True, - optional=True), - scale_factor_width=IntOrFloatInputType(const=True, - optional=True), + scale_factor_height=ScalarOrTensorInputType( + const=True, + optional=True, + type_domain=(np.int32, np.float32) + ), + scale_factor_width=ScalarOrTensorInputType( + const=True, + optional=True, + type_domain=(np.int32, np.float32) + ), align_corners=BoolInputType(const=True, optional=True), ) @@ -514,7 +265,7 @@ def type_inference(self): return types.tensor(self.x.dtype, ret_shape) -@register_op(doc_str="") +@register_op class resize_bilinear(Operation): """ Resize the spatial (last two) dimensions to the specified target size @@ -652,7 +403,7 @@ def type_inference(self): return types.tensor(self.x.dtype, ret_shape) -@register_op(doc_str="") +@register_op class crop_resize(Operation): """ Resize the spatial dimensions (last two dimensions) of the first input @@ -746,7 +497,7 @@ class crop_resize(Operation): target_height=IntInputType(const=True, optional=True), target_width=IntInputType(const=True, optional=True), normalized_coordinates=BoolInputType(const=True, optional=True), - spatial_scale=FloatInputType(const=True, optional=True), + spatial_scale=ScalarOrTensorInputType(const=True, optional=True, type_domain=(np.float32,)), box_coordinate_mode=StringInputType(const=True, optional=True), sampling_mode=StringInputType(const=True, optional=True), ) @@ -798,7 +549,7 @@ def type_inference(self): return types.tensor(self.x.dtype, ret_shape) -@register_op(doc_str="") +@register_op class crop(Operation): """ Crop the spatial dimensions (last two dimensions) of the input by the @@ -867,3 +618,268 @@ def type_inference(self): + [input_shape[-1] - crop_width[0] - crop_width[1]] ) return types.tensor(self.x.dtype, ret_shape) + +@register_op(opset_version=_IOS15_TARGET) +class affine(Operation): + """ + Apply a linear affine transform to the input 2D image tensor. The value at the + ``(x, y)`` (i.e., ``(w, h)``) coordinate of the output is computed by first computing + the coordinates ``x’`` and ``y’`` with the following equation, and then computing the + value at the coordinate ``(x’,y’)`` in the input image using either bilinear or + nearest neighbor interpolation. If the ``(x’, y’)`` point falls outside the input + image, then padding information is used to compute the value. + + .. sourcecode:: python + + * x’ = a0 * x + a1 * y + a2 + * y’ = b0 * x + b1 * y + b2 + + + Parameters + ---------- + x: tensor<[B, C, H1, W1], T> + * Must be rank ``4``. + transform_matrix: tensor<[D, 6], T> + * Must be rank ``2``. + * ``D`` can be either ``B`` or 1. + * If ``D == B``, there is a separate transform matrix for each batch. + * If ``D == 1``, the same matrix is used for all input batches. + * For each batch: ``[a0, a1, a2, b0, b1, b2]``. + output_height: const + * Target output height + output_width: const + * Target output width + sampling_mode: const + * Allowed values: ``"bilinear"`` + padding_mode: const + * Allowed values: ``"constant"``. + * Note that the following example is 1D case for brevity. + The op supports only 2D image input. + * If ``padding_mode == "constant"``: + * The input image is assumed to be padded with the padding_value. + * For example, ``|1, 2, 3| -> |0, 0, 0, 1, 2, 3, 0, 0, 0|``. + padding_value: const + * Currently non-zero values are not supported. + * To be used only when ``padding_mode == "constant"``, ignored in other cases. + coordinates_mode: const + * Allowed values: ``"normalized_minus_one_to_one"`` + * If ``coordinates_mode == "normalized_minus_one_to_one"``, in-image values are ``[-1, 1]``. + * For example, if ``coordinates_mode == "normalized_minus_one_to_one"``, + the in range values are ``[-1, 1]``. That is: + * ``(-1, -1)``, i.e. ``(w=-1, h=-1)``, corresponds to the top-left pixel. + * ``(1, -1)``, i.e. ``(w=1, h=-1)``, corresponds to the top-right pixel. + * ``(-1, 1)``, i.e. ``(w=-1, h=1)``, corresponds to the bottom-left pixel. + * ``(1, 1)``, i.e. ``(w=1, h=1)``, corresponds to the bottom-right pixel. + align_corners: const + * Currently ``align_corners=False`` is not supported. + * To be used only when ``coordinates_mode != unnormalized``, ignored otherwise. + * if ``align_corners == True``, the extrema coordinates correspond + to the center of the first and last corner pixels. + * if ``align_corners == False``, the extrema coordinates correspond + to the edge of the first and last corner pixels. + + Returns + ------- + tensor<[B, C, output_height, output_width], T> + + Attributes + ---------- + T: fp16, fp32 + """ + + input_spec = InputSpec( + x=TensorInputType(), + transform_matrix=TensorInputType(), + output_height=IntInputType(const=True), + output_width=IntInputType(const=True), + sampling_mode=StringInputType(const=True), + padding_mode=StringInputType(const=True), + padding_value=FloatInputType(const=True), + coordinates_mode=StringInputType(const=True), + align_corners=BoolInputType(const=True), + ) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def type_inference(self): + if self.x.rank != 4: + raise ValueError( + 'input "x" to the "affine" op must be a rank 4 tensor. ' + "Got rank {} tensor of shape {}".format( + self.x.rank, self.x.shape + ) + ) + if self.transform_matrix.rank != 2: + raise ValueError( + 'input "transform_matrix" to the "affine" op must be a rank 2 tensor. ' + "Got rank {} tensor of shape {}".format( + self.transform_matrix.rank, self.transform_matrix.shape + ) + ) + if self.sampling_mode.val.lower() != "bilinear": + raise NotImplementedError( + 'input "sampling_mode" to the "affine" not implemented. ' + 'Got "{}"'.format(self.sampling_mode.val) + ) + if self.coordinates_mode.val.lower() != "normalized_minus_one_to_one": + raise NotImplementedError( + 'input "coordinates_mode" to the "affine" not implemented. ' + 'Got "{}"'.format(self.coordinates_mode.val) + ) + if self.padding_mode.val.lower() != "constant" or self.padding_value.val != 0.0: + raise NotImplementedError( + 'input "padding_mode" to the "affine" not implemented. ' + 'Got "{}" with "padding_value={}"'.format( + self.padding_mode.val, self.padding_value.val + ) + ) + + input_shape = self.x.shape + transform_matrix_shape = self.transform_matrix.shape + if ( + not is_symbolic(transform_matrix_shape[-1]) + and transform_matrix_shape[-1] != 6 + ): + raise ValueError( + 'input "transform_matrix" to the "affine" op last dimension must be 6 ' + "[a0, a1, a2, b0, b1, b2], " + "Got {} for last dimension".format(transform_matrix_shape[-1]) + ) + + ret_shape = list(input_shape) + ret_shape[2] = self.output_height.val + ret_shape[3] = self.output_width.val + return types.tensor(self.x.dtype, tuple(ret_shape)) + + +@register_op(opset_version=_IOS15_TARGET) +class resample(Operation): + """ + Resample the input image tensor ``x`` at the ``coordinates``. + Resampling is required if the coordinates do not correspond to exact + pixels in the input image. The ``sampling_mode ``determines + the algorithm used for resampling and computing the values. + + Parameters + ---------- + x: tensor<[B, C, H1, W1], T> + * Must be rank ``4``. + coordinates: tensor<[B, H2, W2, 2], U> + * Must be rank ``4``. + * Coordinates are provided in the order ``(x, y)`` (i.e. ``(w, h)``). + * The value of each output location ``output[b, c, h, w]`` is calculated + by sampling from the input image ``x[b, c, :, :]``. + * The pixel at the ``(x, y)`` location corresponds to the length-2 + vector: ``coordinates[b, h, w, :]``. + * Coordinate (normalized or unnormalized) should be specified according + to ``coordinates_mode``. + sampling_mode: const + * Allowed values: ``"bilinear" , "nearest"`` + padding_mode: const + * Allowed values: ``"constant"``, ``"border"``, ``"reflection"``, ``"symmetric"`` + * Note that the following example is 1D case for brevity. + The op supports only 2D image input. + * If ``padding_mode == "constant"``: + * The input image is assumed to be padded with the ``padding_value``. + * For example: ``|1, 2, 3| -> |0, 0, 0, 1, 2, 3, 0, 0, 0|`` + * if ``padding_mode == "border"``: + * The input image is assumed to be padded with the values replicated + from the values at the edge. This is also referred to as the + "clamped" or "replication" mode, since the padded values are + clamped to the border values. + * For example: ``|1, 2, 3| -> |1, 1, 1, 1, 2, 3, 3, 3, 3|`` + * If ``padding_mode == "reflection"``: + * The border values are reflected, *not* including the values at the edge/border. + * For example: ``|1, 2, 3| -> |2, 3, 2, 1, 2, 3, 2, 1, 2|`` + * If ``padding_mode == "symmetric"``: + * Values are reflected, including the border/edge values. + * For example: ``|1, 2, 3| -> |3, 2, 1 , 1, 2, 3, 3, 2, 1|`` + padding_value: const + * To be used only when ``padding_mode == "constant"``, ignored in other cases. + coordinates_mode: const + * Allowed values: ``"unnormalized"``, ``"normalized_minus_one_to_one"``, + ``"normalized_zero_to_one"`` + * If ``coordinates_mode == "unnormalized"``, the coordinates input values + are interpreted to be in range ``[0, W - 1] / [0, H - 1]``, which + corresponds to the in-image point. + * If ``coordinates_mode == "normalized_minus_one_to_one"``, + the in-image values are ``[-1, 1]``. + * If ``coordinates_mode == "normalized_zero_to_one"``, + in-image values are ``[0, 1]``. + * For example, if ``coordinates_mode == "normalized_minus_one_to_one"``, + the in range values are [-1, 1]. That is: + * ``(-1, -1)``, i.e. ``(w=-1, h=-1)``, corresponds to the top-left pixel. + * ``(1, -1)``, i.e. ``(w=1, h=-1)``, corresponds to the top-right pixel. + * ``(-1, 1)``, i.e. ``(w=-1, h=1)``, corresponds to the bottom-left pixel. + * ``(1, 1)``, i.e. ``(w=1, h=1)``, corresponds to the bottom-right pixel. + align_corners: const + * If ``align_corners == True``, the extrema coordinates correspond + to the center of the first and last corner pixels. + * If ``align_corners == False``, the extrema coordinates correspond + to the edge of the first and last corner pixels. + + Returns + ------- + tensor<[B, C, H2, W2], T> + + Attributes + ---------- + T: fp16, fp32 + U: fp32, i32 + """ + + input_spec = InputSpec( + x=TensorInputType(), + coordinates=ScalarOrTensorInputType(type_domain=(np.int32, np.float32)), + sampling_mode=StringInputType(const=True), + padding_mode=StringInputType(const=True), + padding_value=FloatInputType(const=True), + coordinates_mode=StringInputType(const=True), + align_corners=BoolInputType(const=True), + ) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def type_inference(self): + if self.x.rank != 4: + raise ValueError( + 'input "x" to the "resample" op must be a rank 4 tensor. ' + "Got rank {} tensor of shape {}".format( + self.x.rank, self.x.shape + ) + ) + if self.coordinates.rank != 4: + raise ValueError( + 'input "coordinates" to the "resample" op must be a rank 4 tensor. ' + "Got rank {} tensor of shape {}".format( + self.coordinates.rank, self.coordinates.shape + ) + ) + + input_shape = self.x.shape + coord_shape = self.coordinates.shape + if ( + not is_symbolic(input_shape[0]) + and not is_symbolic(coord_shape[0]) + and input_shape[0] != coord_shape[0] + ): + raise ValueError( + 'input "x" and "coordinates" to the "resample" must agree on ' + "dimension of batch size: {} vs. {}".format( + input_shape[0], coord_shape[0] + ) + ) + if not is_symbolic(coord_shape[-1]) and coord_shape[-1] != 2: + raise ValueError( + 'input "coordinates" to the "resample" op last dimension must be 2. ' + "Got {} for last dimension".format( + coord_shape[-1] + ) + ) + + ret_shape = list(input_shape) + ret_shape[2] = coord_shape[1] # Output height + ret_shape[3] = coord_shape[2] # Output width + return types.tensor(self.x.dtype, tuple(ret_shape)) diff --git a/coremltools/converters/mil/mil/ops/defs/linear.py b/coremltools/converters/mil/mil/ops/defs/iOS15/linear.py similarity index 98% rename from coremltools/converters/mil/mil/ops/defs/linear.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/linear.py index af65b4889..5849aa2ad 100644 --- a/coremltools/converters/mil/mil/ops/defs/linear.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/linear.py @@ -16,12 +16,12 @@ types, ) from coremltools.converters.mil.mil.operation import VALUE +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs._utils import broadcast_shapes, parse_einsum_equation from coremltools.converters.mil.mil.types.symbolic import is_symbolic -from ._op_reqs import register_op -from ._utils import broadcast_shapes, parse_einsum_equation -@register_op(doc_str="") +@register_op class linear(Operation): """ Perform ``x * weight.T + bias`` where ``weight`` and ``bias`` are constant at @@ -91,7 +91,7 @@ def value_inference(self): return res -@register_op(doc_str="") +@register_op class matmul(Operation): """ Perform N-D batch matrix multiplication with NumPy-style broadcasting @@ -228,7 +228,7 @@ def value_inference(self): return np.matmul(x, y) -@register_op(doc_str="") +@register_op class einsum(Operation): """ Perform tensor multiplication expressed according to the einsum notation. diff --git a/coremltools/converters/mil/mil/ops/defs/normalization.py b/coremltools/converters/mil/mil/ops/defs/iOS15/normalization.py similarity index 98% rename from coremltools/converters/mil/mil/ops/defs/normalization.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/normalization.py index 7c7a46d5e..3bd1db0cc 100644 --- a/coremltools/converters/mil/mil/ops/defs/normalization.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/normalization.py @@ -4,7 +4,6 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np -from ._op_reqs import register_op from coremltools.converters.mil.mil import ( DefaultInputs, FloatInputType, @@ -17,11 +16,12 @@ types, ) from coremltools.converters.mil.mil.operation import VALUE +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="") +@register_op class batch_norm(Operation): """ Normalize input tensor ``x`` by ``mean`` and ``variance``, and optionally apply a @@ -85,7 +85,7 @@ def type_inference(self): return types.tensor(self.x.dtype, tuple(x_shape)) -@register_op(doc_str="") +@register_op class instance_norm(Operation): """ Apply instance normalization to the n-dimensional input tensor. @@ -137,7 +137,7 @@ def type_inference(self): return types.tensor(self.x.dtype, tuple(x_shape)) -@register_op(doc_str="") +@register_op class l2_norm(Operation): """ Apply L2 normalization to the n-dimensional input tensor. That is, divide the input @@ -210,8 +210,7 @@ def value_inference(self): output = np.reshape(output, shape) return output - -@register_op(doc_str="") +@register_op class layer_norm(Operation): """ Apply layer normalization to the n-dimensional input tensor: @@ -326,7 +325,7 @@ def np_layer_norm(x, axes, gamma, beta, epsilon=1e-5): return np_layer_norm(self.x.val, _axes, _gamma, _beta, self.epsilon.val) -@register_op(doc_str="") +@register_op class local_response_norm(Operation): """ Apply local response normalization to the n-dimensional input tensor: diff --git a/coremltools/converters/mil/mil/ops/defs/pool.py b/coremltools/converters/mil/mil/ops/defs/iOS15/pool.py similarity index 99% rename from coremltools/converters/mil/mil/ops/defs/pool.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/pool.py index 94d9eb8bc..bb4288b14 100644 --- a/coremltools/converters/mil/mil/ops/defs/pool.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/pool.py @@ -74,7 +74,7 @@ def type_inference(self): return types.tensor(self.x.dtype, tuple(ret_shape)) -@register_op(doc_str="") +@register_op class avg_pool(Pooling): """ Perform average pooling. Supports 1-D, 2-D, and 3-D pool (1, 2, or 3 spatial dimensions). @@ -173,7 +173,7 @@ def __init__(self, **kwargs): super().__init__(**kwargs) -@register_op(doc_str="") +@register_op class l2_pool(Pooling): """ Perform L2 pooling. Supports 1-D, 2-D, and 3-D pool. @@ -213,7 +213,7 @@ def __init__(self, **kwargs): super().__init__(**kwargs) -@register_op(doc_str="") +@register_op class max_pool(Pooling): """ Perform max pooling. Supports 1-D, 2-D, and 3-D pool. diff --git a/coremltools/converters/mil/mil/ops/defs/random.py b/coremltools/converters/mil/mil/ops/defs/iOS15/random.py similarity index 98% rename from coremltools/converters/mil/mil/ops/defs/random.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/random.py index 563cdaf92..7ce6f2638 100644 --- a/coremltools/converters/mil/mil/ops/defs/random.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/random.py @@ -3,7 +3,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 ._op_reqs import register_op from coremltools.converters.mil.mil import get_new_symbol, get_new_variadic_symbol, types from coremltools.converters.mil.mil.input_type import ( DefaultInputs, @@ -15,6 +14,7 @@ StringInputType ) from coremltools.converters.mil.mil.operation import Operation +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op from coremltools.converters.mil.mil.types.symbolic import any_symbolic @@ -46,7 +46,7 @@ def type_inference(self): """ -@register_op(doc_str="") +@register_op class random_bernoulli(RandomDistribution): r""" Returns a tensor with the specified shape, with random values from a Bernoulli @@ -107,7 +107,7 @@ def type_inference(self): return super().type_inference() -@register_op(doc_str="") +@register_op class random_categorical(Operation): """ Returns random values from a categorical distribution. @@ -165,7 +165,7 @@ def type_inference(self): return types.tensor(self.out_dtype, output_shape) -@register_op(doc_str="") +@register_op class random_normal(RandomDistribution): r""" Returns a tensor with the specified shape, with random values from a normal @@ -226,7 +226,7 @@ def type_inference(self): return super().type_inference() -@register_op(doc_str="") +@register_op class random_uniform(RandomDistribution): r""" Returns a tensor with the specified shape with random values from a uniform diff --git a/coremltools/converters/mil/mil/ops/defs/recurrent.py b/coremltools/converters/mil/mil/ops/defs/iOS15/recurrent.py similarity index 99% rename from coremltools/converters/mil/mil/ops/defs/recurrent.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/recurrent.py index 6962d3caa..62bfcf4d6 100644 --- a/coremltools/converters/mil/mil/ops/defs/recurrent.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/recurrent.py @@ -3,7 +3,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 ._op_reqs import register_op from coremltools.converters.mil.mil import Operation, types from coremltools.converters.mil.mil.input_type import ( BoolInputType, @@ -13,9 +12,10 @@ TensorInputType, StringInputType ) +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op -@register_op(doc_str="") +@register_op class gru(Operation): r""" Gated recurrent unit (GRU). @@ -170,7 +170,7 @@ def type_inference(self): ) -@register_op(doc_str="") +@register_op class lstm(Operation): r""" Single long short-term memory (LSTM) sequence. @@ -398,7 +398,7 @@ def weight_shape_check(wt_ih, wt_hh): ) -@register_op(doc_str="") +@register_op class rnn(Operation): """ Recurrent neural network (RNN). diff --git a/coremltools/converters/mil/mil/ops/defs/reduction.py b/coremltools/converters/mil/mil/ops/defs/iOS15/reduction.py similarity index 97% rename from coremltools/converters/mil/mil/ops/defs/reduction.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/reduction.py index 4328fde32..b76e2e2e8 100644 --- a/coremltools/converters/mil/mil/ops/defs/reduction.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/reduction.py @@ -4,7 +4,6 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import numpy as np -from ._op_reqs import register_op from coremltools.converters.mil.mil import ( Operation, precondition, @@ -19,6 +18,7 @@ TensorInputType ) from coremltools.converters.mil.mil.operation import VALUE +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op class ReductionAxes(Operation): @@ -116,7 +116,7 @@ def get_operator(self): raise NotImplementedError() -@register_op(doc_str="") +@register_op class reduce_arg(ReductionAxis): def __init__(self, **kwargs): super().__init__(**kwargs) @@ -139,7 +139,7 @@ def type_inference(self): Reduction op implementations """ -@register_op(doc_str="") +@register_op class reduce_argmax(reduce_arg): """ Computes the indices of the maximum value across dimensions of a tensor. @@ -177,7 +177,7 @@ def get_operator(self): return np.argmax -@register_op(doc_str="") +@register_op class reduce_argmin(reduce_arg): """ Computes the indices of the minimum value across dimensions of a tensor. @@ -216,7 +216,7 @@ def get_operator(self): return np.argmin -@register_op(doc_str="") +@register_op class reduce_l1_norm(ReductionAxes): """ Computes the L1 normalization of elements across given dimensions of the input tensor. @@ -258,7 +258,7 @@ def l1_norm(x, axis=None, keepdims=False): return l1_norm -@register_op(doc_str="") +@register_op class reduce_l2_norm(ReductionAxes): """ Computes the L2 normalization of elements across given dimensions of the input tensor. @@ -295,7 +295,7 @@ def l2_norm(x, axis=None, keepdims=False): return l2_norm -@register_op(doc_str="") +@register_op class reduce_log_sum(ReductionAxes): """ Computes the natural logarithm of the sum of all the elements across given dimensions @@ -333,7 +333,7 @@ def log_sum(x, axis=None, keepdims=False): return log_sum -@register_op(doc_str="") +@register_op class reduce_log_sum_exp(ReductionAxes): """ Computes the natural logarithm of the sum of the exponentials of the elements across @@ -387,7 +387,7 @@ def operator(a, axis=None, keepdims=False): return operator -@register_op(doc_str="") +@register_op class reduce_max(ReductionAxes): """ Computes the maximum of elements across given dimensions of the input tensor. @@ -421,7 +421,7 @@ def get_operator(self): return np.max -@register_op(doc_str="") +@register_op class reduce_mean(ReductionAxes): """ Computes the mean of elements across given dimensions of the input tensor. @@ -459,7 +459,7 @@ def get_operator(self): return np.mean -@register_op(doc_str="") +@register_op class reduce_min(ReductionAxes): """ Computes the minimum of elements across given dimensions of the input tensor. @@ -493,7 +493,7 @@ def get_operator(self): return np.min -@register_op(doc_str="") +@register_op class reduce_prod(ReductionAxes): """ Computes the product of elements across given dimensions of the input tensor. @@ -528,7 +528,7 @@ def get_operator(self): return np.prod -@register_op(doc_str="") +@register_op class reduce_sum(ReductionAxes): """ Computes the sum of elements across given dimensions of the input tensor. @@ -562,7 +562,7 @@ def get_operator(self): return np.sum -@register_op(doc_str="") +@register_op class reduce_sum_square(ReductionAxes): """ Computes the sum of squares of elements across given dimensions of the input tensor. diff --git a/coremltools/converters/mil/mil/ops/defs/scatter_gather.py b/coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py similarity index 99% rename from coremltools/converters/mil/mil/ops/defs/scatter_gather.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py index 813944993..4cbfb5621 100644 --- a/coremltools/converters/mil/mil/ops/defs/scatter_gather.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/scatter_gather.py @@ -24,7 +24,7 @@ ) -@register_op(doc_str="") +@register_op class gather(Operation): """ Gather slices from input ``x`` along dimension ``axis`` according to ``indices``, @@ -134,7 +134,7 @@ def type_inference(self): return types.tensor(out_type, out_shape) -@register_op(doc_str="") +@register_op class scatter(Operation): """ Scatter ``updates`` to ``data`` at locations ``indices`` at dimension ``axis`` @@ -247,7 +247,7 @@ def type_inference(self): return self.data.sym_type -@register_op(doc_str="") +@register_op class gather_along_axis(Operation): """ Take the values along ``axis`` at locations ``indices``. @@ -323,7 +323,7 @@ def type_inference(self): return types.tensor(self.x.dtype, self.indices.shape) -@register_op(doc_str="") +@register_op class scatter_along_axis(Operation): """ Scatter ``updates`` to ``data`` at locations ``indices`` along ``axis`` dimension @@ -440,7 +440,7 @@ def type_inference(self): return self.data.sym_type -@register_op(doc_str="") +@register_op class gather_nd(Operation): """ Gather slices from ``x`` according to ``indices``, similar to `tf.gather_nd `_. @@ -488,7 +488,7 @@ def type_inference(self): return types.tensor(out_type, out_shape) -@register_op(doc_str="") +@register_op class scatter_nd(Operation): """ Scatter ``updates`` to ``data`` at locations ``indices``. diff --git a/coremltools/converters/mil/mil/ops/defs/tensor_operation.py b/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_operation.py similarity index 98% rename from coremltools/converters/mil/mil/ops/defs/tensor_operation.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/tensor_operation.py index fcd7f703d..1c55dabd0 100644 --- a/coremltools/converters/mil/mil/ops/defs/tensor_operation.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_operation.py @@ -38,12 +38,11 @@ SYMBOL, VALUE ) +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs._utils import promoted_primitive_type, MAX_SIZE_CONSTANT_FOLDING -from ._op_reqs import register_op -from ._utils import promoted_primitive_type, MAX_SIZE_CONSTANT_FOLDING - -@register_op(doc_str="") +@register_op class band_part(Operation): """ Returns a tensor setting everything outside a center band to zeros for the innermost @@ -94,7 +93,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class cumsum(Operation): """ Returns the cumulative sum of the input along the given axis. @@ -169,7 +168,7 @@ def type_inference(self): return self.x.sym_type -@register_op(doc_str="") +@register_op class fill(Operation): """ Returns a tensor with a given shape filled with a constant value. @@ -222,7 +221,7 @@ def value_inference(self): return np.full(shape=self.shape.val, fill_value=self.value.val) -@register_op(doc_str="") +@register_op class non_maximum_suppression(Operation): """ Applies non-maximum suppression (NMS) on the input box coordinates according @@ -299,7 +298,7 @@ def type_inference(self): ) -@register_op(doc_str="") +@register_op class non_zero(Operation): """ Returns the indices of the elements in the given tensor that are non-zero. @@ -335,7 +334,7 @@ def value_inference(self): return np.transpose(np.nonzero(self.x.val)) -@register_op(doc_str="") +@register_op class one_hot(Operation): """ Returns one-hot vectors whose locations represented in ``indices`` take the ``on_value``, @@ -420,7 +419,7 @@ def type_inference(self): return types.tensor(on_type, retshape) -@register_op(doc_str="") +@register_op class pad(Operation): """ Pad a tensor. @@ -527,7 +526,7 @@ def value_inference(self): return np.pad(self.x.val, pad_val, mode) -@register_op(doc_str="") +@register_op class range_1d(Operation): """ Returns a numpy-like 1- range sequence. @@ -596,7 +595,7 @@ def type_inference(self): return types.tensor(self.start.dtype, shape) -@register_op(doc_str="") +@register_op class tile(Operation): """ Returns a new tensor by replicating input ``x`` multiples times. @@ -664,7 +663,7 @@ def value_inference(self): return np.tile(self.x.val, reps=self.reps.val) -@register_op(doc_str="") +@register_op class argsort(Operation): """ Returns a tensor containing the indices of the sorted values along a given axis @@ -717,7 +716,7 @@ def value_inference(self): return np.argsort(-self.x.val, axis=self.axis.val) -@register_op(doc_str="") +@register_op class topk(Operation): """ Returns a tensor containing top or bottom ``k`` values and the corresponding @@ -792,7 +791,7 @@ def value_inference(self): return values, indices -@register_op(doc_str="") +@register_op class flatten2d(Operation): """ Flattens input tensor into 2d tensor by flattening dimensions before and @@ -856,7 +855,7 @@ def value_inference(self): return self.x.val.reshape(dim_pre_axis, dim_post_axis) -@register_op(doc_str="") +@register_op class shape(Operation): """ Returns a 1-dimensional tensor with the shape of the input tensor @@ -895,7 +894,7 @@ def value_inference(self): return np.array(self.x.shape).astype(np.int32) -@register_op(doc_str="") +@register_op class concat(Operation): """ Concatenates tensors along a dimension. @@ -1059,7 +1058,7 @@ def value_inference(self): return np.concatenate(values, axis=self.axis.val) -@register_op(doc_str="") +@register_op class split(Operation): """ Split tensors into a tuple @@ -1188,7 +1187,7 @@ def value_inference(self): return tuple(np.split(self.x.sym_val, split_indices[:-1], axis=self.axis.val)) -@register_op(doc_str="") +@register_op class stack(Operation): """ Concatenates tensors along a dimension. @@ -1256,7 +1255,7 @@ def value_inference(self): # identity is used for renaming and is rarely necessary. See # `loop_invariant_elimination` pass for a rare use case. -@register_op(doc_str="") +@register_op class identity(Operation): """ Returns a tensor with the same shape and contents as input. diff --git a/coremltools/converters/mil/mil/ops/defs/tensor_transformation.py b/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py similarity index 92% rename from coremltools/converters/mil/mil/ops/defs/tensor_transformation.py rename to coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py index 3859e7341..521c2016b 100644 --- a/coremltools/converters/mil/mil/ops/defs/tensor_transformation.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS15/tensor_transformation.py @@ -36,80 +36,9 @@ ) from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs._utils import solve_slice_by_index_shape -def _solve_slice_by_index_shape(x_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask): - """ - Helper function to solve the shape of tensor slicing. - """ - ret_shape = [] - - if begin is None or len(begin) == 0: - begin = [None] * len(x_shape) - if end is None or len(end) == 0: - end = [None] * len(x_shape) - - # solve for shape inference - for idx in range(len(x_shape)): - # skip if we want to squeeze the dimension - if squeeze_mask[idx]: - continue - - # for those a[:] cases - if begin_mask[idx] and end_mask[idx]: - if is_symbolic(x_shape[idx]): - if stride[idx] == -1 or stride[idx] == 1: - ret_shape.append(x_shape[idx]) - else: - ret_shape.append(get_new_symbol()) - continue - else: - num = np.ceil(float(x_shape[idx]) / abs(stride[idx])).astype( - np.int32 - ) - ret_shape.append(num) - continue - - # for symbolic case - if is_symbolic(x_shape[idx]): - ret_shape.append(get_new_symbol()) - continue - - # when begin and end are not determined - if begin[idx] is None and not begin_mask[idx]: - ret_shape.append(get_new_symbol()) - continue - if end[idx] is None and not end_mask[idx]: - ret_shape.append(get_new_symbol()) - continue - - # parse negative dimention - if begin[idx] is not None and begin[idx] < 0: - begin[idx] = max(0, begin[idx] + x_shape[idx]) - if end[idx] is not None and end[idx] < 0: - end[idx] = max(0, end[idx] + x_shape[idx]) - - # compute shape - low, high = [0, x_shape[idx]] if stride[idx] > 0 else [-1, x_shape[idx] - 1] - begin_idx, end_idx = ( - [begin[idx], end[idx]] if stride[idx] > 0 else [end[idx], begin[idx]] - ) - is_begin_mask, is_end_mask = ( - [begin_mask[idx], end_mask[idx]] - if stride[idx] > 0 - else [end_mask[idx], begin_mask[idx]] - ) - if is_begin_mask: - begin_idx = low - end_idx = high if is_end_mask else min(end_idx, high) - num = np.ceil(float(end_idx - begin_idx) / abs(stride[idx])).astype( - np.int32 - ) - ret_shape.append(max(0, num)) - - return ret_shape - - -@register_op(doc_str="") +@register_op class depth_to_space(Operation): """ Rearrange elements in a tensor from depth (channel) into spatial dimensions. @@ -148,7 +77,7 @@ def type_inference(self): return types.tensor(x_type, ret_shape) -@register_op(doc_str="") +@register_op class expand_dims(Operation): """ Insert a single-dimension in a 1-D or higher tensor at each axis in axes. @@ -231,7 +160,7 @@ def reshape_with_symbol(v, shape): return v.reshape(shape) -@register_op(doc_str="") +@register_op class reshape(Operation): """ Return a tensor that has the same values as ``x`` with shape ``shape``. @@ -365,7 +294,7 @@ def enforce_volumetric_constraint(left_volume, inshape): return shape -@register_op(doc_str="") +@register_op class reverse(Operation): """ Reverse the order of the input tensor ``x`` along specified ``axes`` (dimensions). @@ -419,7 +348,7 @@ def value_inference(self): return res -@register_op(doc_str="") +@register_op class reverse_sequence(Operation): """ Reverse variable length slices for specified axes / dimensions of the input @@ -480,7 +409,7 @@ def value_inference(self): raise NotImplementedError("TODO") -@register_op(doc_str="") +@register_op class slice_by_index(Operation): """ Method for numpy style indexing and slicing. @@ -561,7 +490,7 @@ def type_inference(self): # solve shape x_shape = self.x.shape - ret_shape = _solve_slice_by_index_shape(x_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask) + ret_shape = solve_slice_by_index_shape(x_shape, begin, end, stride, begin_mask, end_mask, squeeze_mask) if len(ret_shape) == 0: # Scalar case. @@ -628,7 +557,7 @@ def value_inference(self): return res -@register_op(doc_str="") +@register_op class slice_by_size(Operation): """ Slice input tensor starting from the given ``begin`` index and by @@ -729,7 +658,7 @@ def value_inference(self): return self.x.val[tuple(slices)] -@register_op(doc_str="") +@register_op class space_to_depth(Operation): """ Rearrange elements in a tensor from spatial into depth (channel) dimension. @@ -766,7 +695,7 @@ def type_inference(self): ret_shape = (n, c * (bs * bs), h // bs, w // bs) return types.tensor(x_type, ret_shape) -@register_op(doc_str="") +@register_op class space_to_batch(Operation): """ Rearrange elements in a tensor from spatial into batch dimension. @@ -841,8 +770,7 @@ def type_inference(self): return types.tensor(x_type, ret_shape) - -@register_op(doc_str="") +@register_op() class batch_to_space(Operation): """ Rearrange elements in a tensor from batch into spatial dimension. @@ -922,7 +850,7 @@ def type_inference(self): return types.tensor(x_type, ret_shape) -@register_op(doc_str="") +@register_op class squeeze(Operation): """ Remove single-dimension dimensions in a 1-D or higher tensor. @@ -987,7 +915,7 @@ def value_inference(self): val = np.squeeze(self.x.val, axis=tuple(self.axes.val)) return val if val.shape != () else self.x.val[0] -@register_op(doc_str="") +@register_op class transpose(Operation): """ Permute tensor ``x`` dimensions according to ``perm``. @@ -1040,7 +968,7 @@ def value_inference(self): return np.transpose(self.x.val, axes=self.perm.val) -@register_op(doc_str="") +@register_op class pixel_shuffle(Operation): """ Rearrange elements in a tensor from depth (channel) into spatial dimensions. @@ -1082,7 +1010,7 @@ def type_inference(self): return types.tensor(x_type, ret_shape) -@register_op(doc_str="") +@register_op class sliding_windows(Operation): """ Return a tensor containing all windows of ``size``, separated by stride along the diff --git a/coremltools/converters/mil/mil/ops/defs/iOS16/__init__.py b/coremltools/converters/mil/mil/ops/defs/iOS16/__init__.py new file mode 100644 index 000000000..66f287829 --- /dev/null +++ b/coremltools/converters/mil/mil/ops/defs/iOS16/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2022, 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._deployment_compatibility import AvailableTarget as target + +_IOS16_TARGET = target.iOS16 + +from .constexpr_ops import ( + constexpr_affine_dequantize, + constexpr_cast, + constexpr_lut_to_dense, + constexpr_sparse_to_dense, +) + +from .image_resizing import resample + +from .tensor_operation import topk + +from .tensor_transformation import pixel_unshuffle diff --git a/coremltools/converters/mil/mil/ops/defs/constexpr_ops.py b/coremltools/converters/mil/mil/ops/defs/iOS16/constexpr_ops.py similarity index 93% rename from coremltools/converters/mil/mil/ops/defs/constexpr_ops.py rename to coremltools/converters/mil/mil/ops/defs/iOS16/constexpr_ops.py index 56c589436..55593af58 100644 --- a/coremltools/converters/mil/mil/ops/defs/constexpr_ops.py +++ b/coremltools/converters/mil/mil/ops/defs/iOS16/constexpr_ops.py @@ -6,11 +6,12 @@ InputSpec, ScalarOrTensorInputType, ) +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target from coremltools.converters.mil.mil.operation import Operation from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs.iOS16 import _IOS16_TARGET - -@register_op(doc_str="") +@register_op(opset_version=_IOS16_TARGET) class constexpr_affine_dequantize(Operation): """ A compile-time operation that returns a constant output value upon dequantizing its constant inputs. @@ -113,9 +114,6 @@ def assert_vector_size_same_as_axial_dimension(param, axis_dim_size, name): return types.tensor(dtype, shape) def value_inference(self): - return None # Needs to be None to avoid decompression - - def get_decompressed_value(self): return self.decompress( self.quantized_data.val, self.zero_point.val, @@ -125,16 +123,21 @@ def get_decompressed_value(self): @staticmethod def decompress(quantized_data, zero_point, scale, axis): - axes = tuple( - [i for i in range(len(quantized_data.shape)) if i != axis] - ) - sc = np.expand_dims(scale, axis=axes) - zp = np.expand_dims(zero_point, axis=axes) + + def rank_promoted_to_same_as_quantized_data(param): + if len(param.shape) == 0: + return np.reshape(param, np.ones(len(quantized_data.shape), np.int32)) + else: + axes = [i for i in range(len(quantized_data.shape)) if i != axis] + return np.expand_dims(param, axis=tuple(axes)) + + sc = rank_promoted_to_same_as_quantized_data(scale) + zp = rank_promoted_to_same_as_quantized_data(zero_point) val = sc * (quantized_data.astype(np.float32) - zp.astype(np.float32)) return val.astype(scale.dtype) -@register_op(doc_str="") +@register_op(opset_version=_IOS16_TARGET) class constexpr_cast(Operation): """ A compile-time operation that returns a constant output value upon casting its constant input. @@ -175,10 +178,10 @@ def type_inference(self): return types.tensor(dtype, shape) def value_inference(self): - return None # Needs to be None to avoid decompression + return np.float32(self.source_val.val) -@register_op(doc_str="") +@register_op(opset_version=_IOS16_TARGET) class constexpr_lut_to_dense(Operation): """ A compile-time operation that returns a constant output value upon decompressing look-up-table to a dense tensor. @@ -259,9 +262,6 @@ def assert_is_vector(param, name): return types.tensor(dtype, shape) def value_inference(self): - return None # Needs to be None to avoid decompression - - def get_decompressed_value(self): return self.decompress( self.lut.val, self.indices.val, @@ -287,7 +287,7 @@ def decompress(lut, indices, shape): return flatten_val.reshape(shape) -@register_op(doc_str="") +@register_op(opset_version=_IOS16_TARGET) class constexpr_sparse_to_dense(Operation): """ A compile-time operation that returns a constant output value upon de-sparsification of its constant inputs. @@ -363,9 +363,6 @@ def assert_is_vector(param, name): return types.tensor(dtype, shape) def value_inference(self): - return None # Needs to be None to avoid decompression - - def get_decompressed_value(self): return self.decompress(self.nonzero_data.val, self.mask.val, self.shape.val) @staticmethod diff --git a/coremltools/converters/mil/mil/ops/defs/iOS16/image_resizing.py b/coremltools/converters/mil/mil/ops/defs/iOS16/image_resizing.py new file mode 100644 index 000000000..b65a4fb23 --- /dev/null +++ b/coremltools/converters/mil/mil/ops/defs/iOS16/image_resizing.py @@ -0,0 +1,42 @@ +# Copyright (c) 2022, 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 types +from coremltools.converters.mil.mil.input_type import ( + BoolInputType, + FloatInputType, + InputSpec, + ScalarOrTensorInputType, + StringInputType, + TensorInputType, +) +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target +from coremltools.converters.mil.mil.operation import Operation +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs.iOS15.image_resizing import resample as _resample_iOS15 +from coremltools.converters.mil.mil.ops.defs.iOS16 import _IOS16_TARGET + +@register_op(opset_version=_IOS16_TARGET) +class resample(_resample_iOS15): + """ + iOS16 version of resample supports float16 coordinates + """ + input_spec = InputSpec( + x=TensorInputType(), + coordinates=ScalarOrTensorInputType(type_domain=(np.int32, np.float32, np.float16)), + sampling_mode=StringInputType(const=True), + padding_mode=StringInputType(const=True), + padding_value=FloatInputType(const=True), + coordinates_mode=StringInputType(const=True), + align_corners=BoolInputType(const=True), + ) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def type_inference(self): + return super().type_inference() diff --git a/coremltools/converters/mil/mil/ops/defs/iOS16/tensor_operation.py b/coremltools/converters/mil/mil/ops/defs/iOS16/tensor_operation.py new file mode 100644 index 000000000..fd264d1f8 --- /dev/null +++ b/coremltools/converters/mil/mil/ops/defs/iOS16/tensor_operation.py @@ -0,0 +1,74 @@ +# Copyright (c) 2022, 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 coremltools.converters.mil.mil import types + +from coremltools.converters.mil.mil.input_type import ( + InputSpec, + TensorInputType, + ScalarOrTensorInputType, + BoolInputType, + DefaultInputs, +) +from coremltools.converters.mil.mil.operation import ( + Operation, + precondition, + VALUE +) +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs.iOS15.tensor_operation import topk as _topk_iOS15 +from coremltools.converters.mil.mil.ops.defs.iOS16 import _IOS16_TARGET + + +@register_op(opset_version=_IOS16_TARGET) +class topk(_topk_iOS15): + """ + An iOS16 version of topk + + Additional Parameters + ---------- + * sort: const (Optional) + * Default to ``True`` + * If true, top-k elements are themselves sorted. + Otherwise, no particular ordering is guaranteed. + * return_indices: const (Optional) + # Default to ``True`` + # If true, returns both values and indices. Otherwise, returns only the top-k values. + + Returns + ------- + tensor<\*?, T> + * Values of top/bottom ``k`` elements. + + tensor<\*?, int32> + * Only returned when ``return_indices = True`` + * Indices of the top/bottom ``k`` elements along axis. + + Attributes + ---------- + T: fp32, int32 + """ + + input_spec = _topk_iOS15.input_spec + InputSpec( + sort=BoolInputType(const=True, optional=True), + return_indices=BoolInputType(const=True, optional=True), + ) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def default_inputs(self): + return super().default_inputs() + DefaultInputs(sort=True, return_indices=True) + + def type_inference(self): + value_type, indices_type = super().type_inference() + if not self.return_indices.val: + return value_type + return value_type, indices_type + + @precondition(allow=VALUE) + def value_inference(self): + values, indices = super().value_inference() + if not self.return_indices.val: + return values + return values, indices diff --git a/coremltools/converters/mil/mil/ops/defs/iOS16/tensor_transformation.py b/coremltools/converters/mil/mil/ops/defs/iOS16/tensor_transformation.py new file mode 100644 index 000000000..7ec8dcd73 --- /dev/null +++ b/coremltools/converters/mil/mil/ops/defs/iOS16/tensor_transformation.py @@ -0,0 +1,59 @@ +# Copyright (c) 2022, 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 coremltools.converters.mil.mil import types + +import numpy as np + +from coremltools.converters.mil.mil.input_type import( + InputSpec, + TensorInputType, + ScalarOrTensorInputType, +) +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target +from coremltools.converters.mil.mil import types +from coremltools.converters.mil.mil.operation import Operation +from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op +from coremltools.converters.mil.mil.ops.defs.iOS16 import _IOS16_TARGET + +@register_op(opset_version=_IOS16_TARGET) +class pixel_unshuffle(Operation): + """ + Rearrange elements in a tensor from spatial dimensions into depth (channel). + It is basically the inverse operation of pixel_shuffle. + Equivalent to PyTorch's ``PixelUnshuffle``. + + Parameters + ---------- + x: tensor<[n, C, H / f , W / f], T> (Required) + * Input tensor of rank ``4``. + downscale_factor: const + * Factor to decrease spatial resolution by. + + Returns + ------- + tensor<[n, C * f^2, H, W], T> + * Where ``f`` is the downscale factor. + + Attributes + ---------- + T: fp32 + + References + ---------- + `torch.nn.PixelUnshuffle `_ + """ + + input_spec = InputSpec( + x=TensorInputType(), downscale_factor=ScalarOrTensorInputType(const=True, type_domain=(np.uint32,)), + ) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def type_inference(self): + x_type = self.x.dtype + n, c, h, w = self.x.shape + f = self.downscale_factor.val + ret_shape = (n, c * f * f, h / f, w / f) + return types.tensor(x_type, ret_shape) diff --git a/coremltools/converters/mil/mil/ops/helper.py b/coremltools/converters/mil/mil/ops/helper.py new file mode 100644 index 000000000..a86ddeb3f --- /dev/null +++ b/coremltools/converters/mil/mil/ops/helper.py @@ -0,0 +1,24 @@ +# Copyright (c) 2022, 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 _get_version_of_op(op_variants, opset_version): + """ + A utility function that retrieves an op cls given a dictionary of op variants and target version + """ + assert isinstance(op_variants, dict) + opset_versions = list(op_variants.keys()) + opset_versions.sort() + if opset_version is None: + op_cls = op_variants[opset_versions[0]] + else: + if opset_version not in op_variants: + op_type = list(op_variants.values())[0].__name__ + msg = ( + "No available version for {} in the {!s} opset. Please update the " + "minimum_deployment_target to at least {!s}" + ).format(op_type, opset_version, opset_versions[0]) + raise ValueError(msg) + op_cls = op_variants[opset_version] + return op_cls diff --git a/coremltools/converters/mil/mil/ops/registry.py b/coremltools/converters/mil/mil/ops/registry.py index e2d790956..edf87ea7d 100644 --- a/coremltools/converters/mil/mil/ops/registry.py +++ b/coremltools/converters/mil/mil/ops/registry.py @@ -4,65 +4,177 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import logging +from .helper import _get_version_of_op from ..builder import Builder from collections import defaultdict - +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as target +from coremltools.converters.mil.mil.block import curr_opset_version class SSAOpRegistry: - # ops is 3 nested dicts: - # namespace (str) -> {op_type (str) -> {op_class, doc_str}} - ops = defaultdict(dict) + + """ + There are three kinds of operations that we could register: + + (1) core_ops: dict[str, dict[Operation]] + - These are the core ops in PyMIL, which have a direct mapping to the backend in neural_network or mlprogram + - The registered op is considered a core op if the namespace is not provided + - coreml_ops[op_type] is a dict that tracks different opset versions for an op. For instance + - ``core_ops[op_1] = { + ct.target.iOS13: op_1_iOS13, + ct.target.iOS14: op_1_iOS13, + ct.target.iOS15: op_1_iOS13, + ct.target.iOS16: op_1_iOS13, + }`` + . Only one version of op type ``op_1`` is registered, and it is defined in iOS13, which both + neural_network and mlprogram backend support + - ``core_ops[op_2] = { + ct.target.iOS13: op_2_iOS13, + ct.target.iOS14: op_2_iOS13, + ct.target.iOS15: op_2_iOS13, + ct.target.iOS16: op_2_iOS16, + }`` + . Two versions of op type ``op_2`` are registered, one defined in iOS13 and one for iOS16 + . The builder picks up correct version of the op according to curr_opset_version(), which returns the opset version of + the current function. + -- If ``curr_opset_version()`` is ``None`` (the version of the function is not set), ``mb.op_2`` would call the oldest version of the op by default, which is ``op_2_ios13`` + -- Otherwise, the builder would pick up core_ops[op_2][curr_opset_version()] + - In the highest level, users can choose the desired version by specifying the ``minum_deployment_target`` argument in ``coremltools.convert`` + - The default ``opset_version`` for the core ops would be set to iOS13, for which neural_network backend supports + + (2) dialect_ops: dict[str, Operation] + - These are the ops that are created for specific frontend framework, for instance: ``tf_lstm_block, torch_upsample_nearest_neighbor`` + - A graph pass must be customized by the developer to translate a dialect_ops into core ops + + (3) custom_ops: dict[str, Operation] + - These are the custom ops, in which an additional ``bindings`` which should be specificed in operator + """ + SUPPORTED_OPSET_VERSIONS = ( + target.iOS13, + target.iOS14, + target.iOS15, + target.iOS16 + ) + core_ops = defaultdict(dict) + dialect_ops = {} custom_ops = {} + + @staticmethod + def _get_core_op_cls(op_type=None): + """ + A utility function that retrieves an op cls using the curr_opset_version + """ + if op_type not in SSAOpRegistry.core_ops: + raise ValueError("op {} not registered.".format(op_type)) + candidate_ops = SSAOpRegistry.core_ops[op_type] + return _get_version_of_op(candidate_ops, curr_opset_version()) @staticmethod - def register_op(doc_str="", is_custom_op=False, namespace="core", allow_override=False): + def register_op(_cls=None, is_custom_op=False, namespace=None, opset_version=target.iOS13, allow_override=False): """ Registration routine for MIL Program operators - is_custom_op: (Boolean) [Default=False] - If True, maps current operator to `custom_op` - `custom_op` requires additional `bindings` which should be - specified in operator. - 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. - """ + + Parameters + ---------- + is_custom_op: boolean + - If ``True``, maps current operator to ``custom_op``. ``custom_op`` requires additional ``bindings`` which should be specified in operator + - Default ``False`` + + namespace: str + - If provided, the op is registered as a dialect op + - Otherwise is considered as a core op + opset_version: int + - Specify the minimum spec version that supports this op + - Default to ``ct.target.iOS13``, which is for the neural_network backend + + allow_override: boolean + - If True, it is allowed for an operation to override the previous operation with the same registered name + - Default ``False`` + """ def class_wrapper(op_cls): op_type = op_cls.__name__ + op_cls.__name__ = op_type - # Operation specific to custom op - op_msg = "Custom op" if is_custom_op else "op" - op_reg = ( - SSAOpRegistry.custom_ops - if is_custom_op - else SSAOpRegistry.ops[namespace] - ) - + # debug message + op_msg = "op" + is_dialect_op = (namespace is not None) + if is_custom_op: + op_msg = "Custom op" + elif is_dialect_op: + op_msg = "Dialect op" logging.debug("Registering {} {}".format(op_msg, op_type)) - if op_type in op_reg and not allow_override: - raise ValueError( - "SSA {} {} already registered.".format(op_msg, op_type) - ) - - if namespace != "core": + # pick the right dict for registration + if is_custom_op: + op_reg = SSAOpRegistry.custom_ops + elif is_dialect_op: + op_reg = SSAOpRegistry.dialect_ops # Check that op_type is prefixed with namespace if op_type[: len(namespace)] != namespace: msg = ( - "Op type {} registered under {} namespace must " + "Dialect pp type {} registered under {} namespace must " + "prefix with {}" ) raise ValueError(msg.format(op_type, namespace, namespace)) + else: + op_reg = SSAOpRegistry.core_ops - op_reg[op_type] = {"class": op_cls, "doc_str": doc_str} + # verify that the op have not been registered before if allow_override = False + msg = "SSA {} {} already registered.".format(op_msg, op_type) + if is_custom_op or is_dialect_op: + if op_type in op_reg and not allow_override: + raise ValueError(msg) + else: + if opset_version in op_reg[op_type] and not allow_override: + if opset_version - 1 not in op_reg[op_type] or (op_reg[op_type][opset_version - 1] != op_reg[op_type][opset_version]): + raise ValueError(msg) + # add the op to op_reg + if is_custom_op or is_dialect_op: + op_reg[op_type] = op_cls + else: + msg = "Older version of op {} must be registered before a newer version.".format(op_type) + if opset_version in op_reg[op_type]: + old_op_cls = op_reg[op_type][opset_version] + for i in range(opset_version, SSAOpRegistry.SUPPORTED_OPSET_VERSIONS[-1] + 1): + if op_reg[op_type][i] != old_op_cls: + raise ValueError(msg) + elif len(op_reg[op_type]) != 0: + raise ValueError(msg) + idx = SSAOpRegistry.SUPPORTED_OPSET_VERSIONS.index(opset_version) + for i in range(idx, len(SSAOpRegistry.SUPPORTED_OPSET_VERSIONS)): + op_reg[op_type][SSAOpRegistry.SUPPORTED_OPSET_VERSIONS[i]] = op_cls + + # add the version information to the op cls + op_cls._op_variants = op_reg[op_type] + @classmethod def add_op(cls, **kwargs): - return cls._add_op(op_cls, **kwargs) + """ + An utility function that help the builder to pickup the correct op class when calling ``mb.op`` + + There are two cases: + + (1) custom op / dialect op: + If the op is a custom op or a dialect op, we could directly pick up the op class through + ``SSAOpRegistry.custom_ops[op_type]`` or ``SSAOpRegistry.dialect_ops[op_type]`` + + (2) core op: + For the core op, the builder would pick up the correct version according to ``curr_opset_version()`` + """ + op_cls_to_add = None + is_core_op = (op_reg == SSAOpRegistry.core_ops) + if is_core_op: + op_cls_to_add = SSAOpRegistry._get_core_op_cls(op_type) + else: + op_cls_to_add = op_reg[op_type] + + return cls._add_op(op_cls_to_add, **kwargs) setattr(Builder, op_type, add_op) return op_cls - - return class_wrapper + + if _cls is None: + return class_wrapper + + return class_wrapper(_cls) diff --git a/coremltools/converters/mil/mil/ops/tests/test_activation.py b/coremltools/converters/mil/mil/ops/tests/test_activation.py index 8b1b6a0c2..ada3b0f50 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_activation.py +++ b/coremltools/converters/mil/mil/ops/tests/test_activation.py @@ -682,8 +682,6 @@ class TestSiLU: "use_cpu_only, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_only, backend): - if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") x_val = np.array([-1.1, 2.2, -3.3, 4.4], dtype=np.float32).reshape((1, 2, 1, 2)) diff --git a/coremltools/converters/mil/mil/ops/tests/test_const.py b/coremltools/converters/mil/mil/ops/tests/test_const.py index 3d6ade7d6..2361f2d68 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_const.py +++ b/coremltools/converters/mil/mil/ops/tests/test_const.py @@ -36,10 +36,8 @@ class TestConst: ) ) def test_builder_to_backend_smoke(self, use_cpu_for_conversion, backend, dtype): - 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])") if backend[0] == "mlprogram" and dtype in [np.uint8, np.int8, np.uint32]: - pytest.xfail("Data type not supported") + pytest.skip("Data type not supported") t = np.random.randint(0, 5, (4, 2)).astype(np.float32) constant = np.random.randint(0, 5, (4, 2)).astype(dtype) diff --git a/coremltools/converters/mil/mil/ops/tests/test_constexpr_ops.py b/coremltools/converters/mil/mil/ops/tests/test_constexpr_ops.py index 48474c53e..034b77f37 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_constexpr_ops.py +++ b/coremltools/converters/mil/mil/ops/tests/test_constexpr_ops.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020, Apple Inc. All rights reserved. +# Copyright (c) 2022, 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 @@ -11,6 +11,7 @@ import coremltools as ct from coremltools.converters.mil.mil import Builder as mb, types from coremltools.converters.mil.mil.ops.tests.testing_utils import run_compare_builder +from coremltools.converters.mil.testing_utils import get_op_types_in_program, ssa_fn backends = [("mlprogram", "fp32"), ("mlprogram", "fp16")] @@ -49,7 +50,7 @@ def build(x): expected_output_types = (1, 1, 2, 2, types.fp32) expected_outputs = t + decompressed_constant.astype(np.float32) - run_compare_builder( + mlmodel = run_compare_builder( build, input_placeholders, input_values, @@ -62,6 +63,20 @@ def build(x): minimum_deployment_target=ct.target.iOS16, ) + # validate that the constexpr op is not removed by any graph pass + prog = mlmodel._mil_program + assert "constexpr_affine_dequantize" in get_op_types_in_program(prog) + + @ssa_fn + def test_builder_eval(self): + v = mb.constexpr_affine_dequantize( + quantized_data=np.array([[1, 2, 3], [1, 2, 3]]).astype(np.uint8), + zero_point=np.uint8(1), + scale=np.float32(2), + axis=0 + ) + np.testing.assert_allclose(np.float32([[0, 2, 4], [0, 2, 4]]), v.val) + @pytest.mark.skipif( ct.utils._macos_version() < (13, 0), reason="ConstExpr ops available from macOS13 onwards.", @@ -87,7 +102,7 @@ def build(x): expected_output_types = (4, 1, types.fp32) expected_outputs = t + decompressed_constant.astype(np.float32) - run_compare_builder( + mlmodel = run_compare_builder( build, input_placeholders, input_values, @@ -100,6 +115,14 @@ def build(x): minimum_deployment_target=ct.target.iOS16, ) + # validate that the constexpr op is not removed by any graph pass + prog = mlmodel._mil_program + assert "constexpr_cast" in get_op_types_in_program(prog) + + @ssa_fn + def test_builder_eval(self): + v = mb.constexpr_cast(source_val=np.float16([1, 2]), output_dtype="fp32") + np.testing.assert_allclose(np.float32([1, 2]), v.val) @pytest.mark.skipif( ct.utils._macos_version() < (13, 0), @@ -147,7 +170,7 @@ def build(x): expected_output_types = (4, 1, types.fp32) expected_outputs = t + decompressed_constant.astype(np.float32) - run_compare_builder( + mlmodel = run_compare_builder( build, input_placeholders, input_values, @@ -160,6 +183,19 @@ def build(x): minimum_deployment_target=ct.target.iOS16, ) + # validate that the constexpr op is not removed by any graph pass + prog = mlmodel._mil_program + assert "constexpr_lut_to_dense" in get_op_types_in_program(prog) + + @ssa_fn + def test_builder_eval(self): + v = mb.constexpr_lut_to_dense( + lut=np.array([1., 2., 3., 4.]), + indices=np.array([10, 4]).astype(np.uint8), + shape=np.array([5,]).astype(np.uint32), + ) + np.testing.assert_allclose(np.float32([3, 3, 1, 1, 1]).astype(np.float32), v.val) + @pytest.mark.skipif( ct.utils._macos_version() < (13, 0), reason="ConstExpr ops available from macOS13 onwards.", @@ -189,7 +225,7 @@ def build(x): expected_output_types = (4, 1, types.fp32) expected_outputs = t + decompressed_constant.astype(np.float32) - run_compare_builder( + mlmodel = run_compare_builder( build, input_placeholders, input_values, @@ -201,3 +237,16 @@ def build(x): converter=ct.convert, minimum_deployment_target=ct.target.iOS16, ) + + # validate that the constexpr op is not removed by any graph pass + prog = mlmodel._mil_program + assert "constexpr_sparse_to_dense" in get_op_types_in_program(prog) + + @ssa_fn + def test_builder_eval(self): + v = mb.constexpr_sparse_to_dense( + nonzero_data=np.array([1., 2., 4.]), + mask=np.array([11]).astype(np.uint8), + shape=np.array([4,]).astype(np.uint32), + ) + np.testing.assert_allclose(np.float32([1., 2., 0., 4.]), v.val) 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 9d1febbd9..36822c26e 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_control_flow.py +++ b/coremltools/converters/mil/mil/ops/tests/test_control_flow.py @@ -109,11 +109,6 @@ class TestCond: "use_cpu_for_conversion, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_for_conversion, backend): - if backend[0] == "mlprogram": - pytest.skip("rdar://81169758 (TestCond hangs 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])") - input_placeholders = { "a": mb.placeholder(shape=(1,), dtype=types.bool), "b": mb.placeholder(shape=(1,)), @@ -249,7 +244,7 @@ def cond(res, bx): ) def test_builder_to_backend_nested(self, use_cpu_only, backend): if backend[0] == 'neuralnetwork': - pytest.xfail("neuralnetwork backend add const has issue") + pytest.xfail("rdar://96862073 (test_control_folw::TestWhileLoop::test_builder_to_backend_nested failing on nnv1)") input_placeholders = { "x": mb.placeholder(shape=(1,)), diff --git a/coremltools/converters/mil/mil/ops/tests/test_conv.py b/coremltools/converters/mil/mil/ops/tests/test_conv.py index f4111c379..0f285db9a 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_conv.py +++ b/coremltools/converters/mil/mil/ops/tests/test_conv.py @@ -464,10 +464,10 @@ def test_builder_to_backend_stress_weights_input( symbolic = config["symbolic"] if backend[0] == "neuralnetwork" and groups > 1: - pytest.xfail("dynamic conv with groups > 1 is not supported on the neuralnetwork backend") - - if backend[0] == "mlprogram" and conv_dim == "conv1d" and not use_cpu_only: - pytest.xfail("rdar://90819258, mlprogram fails with dynamic weights conv1d on the GPU") + pytest.skip("dynamic conv with groups > 1 is not supported on the neuralnetwork backend") + + if backend[0] == "mlprogram" and not use_cpu_only: + pytest.xfail("rdar://97398343 (test_builder_to_backend_stress_weights_input is failing on mlprogram + GPU)") D, H, W, Kd, Kh, Kw = DHWKdKhKw N, C_in, C_out = 1, 1 * groups, 2 * groups 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 f66d9c3d1..df200a5d0 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_elementwise_binary.py +++ b/coremltools/converters/mil/mil/ops/tests/test_elementwise_binary.py @@ -473,9 +473,6 @@ class TestNotEqual: "use_cpu_for_conversion, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_for_conversion, 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])") - x = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) y = np.array([[-1, 2, -3], [4, -5, 6]], dtype=np.float32) input_placeholders = { diff --git a/coremltools/converters/mil/mil/ops/tests/test_elementwise_unary.py b/coremltools/converters/mil/mil/ops/tests/test_elementwise_unary.py index 6252b9cbc..e5b6e278f 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_elementwise_unary.py +++ b/coremltools/converters/mil/mil/ops/tests/test_elementwise_unary.py @@ -62,9 +62,6 @@ class TestElementwiseUnary: ), ) def test_builder_to_backend_smoke(self, use_cpu_for_conversion, backend, mode): - if backend[0] == "mlprogram" and not use_cpu_for_conversion and mode == "cast": - pytest.xfail("rdar://78343191 ((MIL GPU) Core ML Tools Unit Test failures [failure to load or Seg fault])") - if mode == "abs": val = np.array([[-1, 2, -3], [4, -5, 6]], dtype=np.float32) expected_outputs = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) @@ -622,9 +619,6 @@ def build(x): def test_builder_to_backend_stress_log( self, use_cpu_only, backend, epsilon ): - if backend[0] == "mlprogram" and not use_cpu_only: - pytest.xfail("rdar://78343225 ((MIL GPU) Core ML Tools Unit Test failures [numerical error])") - x = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) numpy_pred = np.log(x + epsilon) @@ -657,9 +651,6 @@ def build(x): def test_builder_to_backend_stress_cast( self, use_cpu_for_conversion, backend, src_dst ): - 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])") - src_dtype, dst_dtype = src_dst x = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) numpy_pred = x.astype(dtype=np.float16) 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 c838ccca3..f03b95cdc 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_image_resizing.py +++ b/coremltools/converters/mil/mil/ops/tests/test_image_resizing.py @@ -10,6 +10,8 @@ import pytest from .testing_utils import run_compare_builder + +import coremltools as ct from coremltools.converters.mil import testing_reqs from coremltools.converters.mil.mil import ( Builder as mb, @@ -29,11 +31,7 @@ class TestAffine: ) def test_builder_to_backend_smoke(self, use_cpu_only, backend): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") - if backend[0] == "mlprogram" and backend[1] == "fp16": - pytest.xfail("rdar://86653285 ([ MIL ] TestAffine::test_builder_to_backend_smoke[[use_cpu_only=False]-mlprogram-fp16] CI Failure)") - if backend[0] == "mlprogram" and backend[1] == "fp32": - pytest.xfail("rdar://88039548 (test_image_resizing.py::TestAffine::test_builder_to_backend_smoke is failing)") + pytest.skip("nn backend not supported") x_val = np.array([11.0, 22.0, 33.0, 44.0], dtype=np.float32).reshape( [1, 1, 2, 2] @@ -102,11 +100,16 @@ def build(x, transform_matrix): class TestResample: @pytest.mark.parametrize( - "use_cpu_only, backend", itertools.product([True, False], backends,) + "use_cpu_only, backend, minimum_deployment_target", + itertools.product( + [True, False], + backends, + [ct.target.iOS15, ct.target.iOS16], + ) ) - def test_builder_to_backend_smoke(self, use_cpu_only, backend): + def test_builder_to_backend_smoke(self, use_cpu_only, backend, minimum_deployment_target): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") + pytest.skip("nn backend not supported") x_ = np.array([11.0, 22.0, 33.0, 44.0], dtype=np.float32).reshape([1, 1, 2, 2]) coordinates_ = np.array( @@ -189,7 +192,7 @@ def build_3(x, coordinates): expected_output_3, ], ): - run_compare_builder( + mlmodel = run_compare_builder( build, input_placeholder_dict, input_value_dict, @@ -197,7 +200,17 @@ def build_3(x, coordinates): expected_output, use_cpu_only=use_cpu_only, backend=backend, + minimum_deployment_target=minimum_deployment_target, ) + prog = mlmodel._mil_program + number_of_cast = len(prog["main"].find_ops(op_type="cast")) + # for the new iOS16 resample op, the coordinates is cast to fp16 + if minimum_deployment_target == ct.target.iOS15: + assert number_of_cast == 2 + elif minimum_deployment_target == ct.target.iOS16: + assert number_of_cast == 3 + else: + raise ValueError("Unrecognized target {}".format(minimum_deployment_target)) class TestResizeNearestNeighbor: @@ -247,10 +260,10 @@ class TestUpsampleNearestNeighborFractionalScales: ) def test_builder_to_backend_smoke(self, use_cpu_for_conversion, backend): if backend[0] == "neuralnetwork": - pytest.xfail("nn backend not supported") - + pytest.skip("nn backend not supported") + if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("rdar://78343225 ((MIL GPU) Core ML Tools Unit Test failures [numerical error])") + pytest.xfail("rdar://97398448 (TestUpsampleNearestNeighborFractionalScales failing on GPU)") x_val = np.array([1.5, -2.5, 3.5], dtype=np.float32).reshape([1, 1, 1, 3]) input_placeholder_dict = {"x": mb.placeholder(shape=x_val.shape)} @@ -303,7 +316,7 @@ def test_builder_to_backend_smoke(self, use_cpu_only, backend): if backend[0] == "mlprogram": pytest.xfail("Seg fault: rdar://78343191 ((MIL GPU) Core ML Tools Unit Test failures [failure to load or Seg fault])") - if backend[0] == "neuralnetwork": + if backend[0] == "neuralnetwork" and use_cpu_only: pytest.xfail("rdar://85318710 (Coremltools Smoke test on ResizeBilinear failing on NNv1 backend.)") x = np.array([0, 1], dtype=np.float32).reshape(1, 1, 2) @@ -407,9 +420,6 @@ class TestUpsampleBilinear: "use_cpu_only, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_only, backend): - if backend[0] == "mlprogram" and not use_cpu_only: - pytest.xfail("test failing on gpu with nan output") - x = np.array([0, 1], dtype=np.float32).reshape(1, 1, 2) input_placeholder_dict = {"x": mb.placeholder(shape=x.shape)} input_value_dict = {"x": x} @@ -653,10 +663,8 @@ class TestCropResize: itertools.product([True, False], backends, [True, False]), ) def test_builder_to_backend_smoke(self, use_cpu_only, backend, is_symbolic): - - if backend[0] == "mlprogram": - pytest.xfail("rdar://78343191 ((MIL GPU) Core ML Tools Unit Test failures [failure to load or Seg fault])") - + if backend[0] == "mlprogram" and not use_cpu_only: + pytest.xfail("rdar://97398582 (TestCropResize failing on mlprogram + GPU)") x = np.array( [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], dtype=np.float32, diff --git a/coremltools/converters/mil/mil/ops/tests/test_linear.py b/coremltools/converters/mil/mil/ops/tests/test_linear.py index 515bb06d4..be7eca0a9 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_linear.py +++ b/coremltools/converters/mil/mil/ops/tests/test_linear.py @@ -57,9 +57,8 @@ def test_builder_eval(self): itertools.product([True, False], backends, [2, 3, 5]), ) def test_builder_to_backend_stress(self, use_cpu_only, backend, rank): - if backend[0] == "mlprogram" and rank == 5 and not use_cpu_only: - pytest.xfail("rdar://94199353 (TestLinear.test_builder_to_backend_stress failing on the CI)") - + if backend[0] == "mlprogram" and not use_cpu_only: + pytest.xfail("rdar://97398733 (TestLinear failing on mlprogram + GPU)") x_shape = np.random.randint(low=1, high=3, size=(rank,)) x_val = np.random.rand(*x_shape) out_channels = 3 diff --git a/coremltools/converters/mil/mil/ops/tests/test_normalization.py b/coremltools/converters/mil/mil/ops/tests/test_normalization.py index b8229f54d..d38a4e0ba 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_normalization.py +++ b/coremltools/converters/mil/mil/ops/tests/test_normalization.py @@ -372,9 +372,6 @@ def _np_layer_norm(x, axes, gamma=None, beta=None, epsilon=1e-5): "use_cpu_only, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_only, backend): - if backend[0] == "mlprogram" and backend[1] == "fp32": - pytest.xfail("rdar://88039548 (test_image_resizing.py::TestAffine::test_builder_to_backend_smoke is failing)") - x_val = np.array([[[1.0, -7.0], [5.0, -6.0], [-3.0, -5.0]]], dtype=np.float32) input_placeholders = {"x": mb.placeholder(shape=x_val.shape)} input_values = {"x": x_val} @@ -439,9 +436,6 @@ def build(x): "use_cpu_only, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke_rank_2(self, use_cpu_only, backend): - if backend[0] == "mlprogram" and backend[1] == "fp32": - pytest.xfail("rdar://88039548 (test_image_resizing.py::TestAffine::test_builder_to_backend_smoke is failing)") - x_val = np.array([[1.0, -7.0], [5.0, -6.0], [-3.0, -5.0]], dtype=np.float32) gamma_val = np.array([1.0, 1.0], dtype=np.float32) beta_val = np.array([1.0, 0.0], dtype=np.float32) diff --git a/coremltools/converters/mil/mil/ops/tests/test_reduction.py b/coremltools/converters/mil/mil/ops/tests/test_reduction.py index 96cecc4db..d5522e2db 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_reduction.py +++ b/coremltools/converters/mil/mil/ops/tests/test_reduction.py @@ -133,6 +133,35 @@ def test_builder_to_backend_global_pool_2d(self, use_cpu_only, backend, mode): backend=backend, ) + @pytest.mark.parametrize( + "use_cpu_only, backend, mode", + itertools.product([True, False], backends, ["max", "mean"]), + ) + def test_builder_to_backend_global_pool_none(self, use_cpu_only, backend, mode): + # test lowering to spatial reduction to global_pool path for axis = None + val = np.array([[[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]]], dtype=np.float32) + input_placeholders = {"x": mb.placeholder(shape=val.shape)} + input_values = {"x": val} + + expected_output_types = (1, 1, 1, 1, types.fp32) + + if mode == "max": + build = lambda x: mb.reduce_max(x=x, axes=None, keep_dims=True) + expected_outputs = np.array([[[[6.0]]]], dtype=np.float32) + elif mode == "mean": + build = lambda x: mb.reduce_mean(x=x, axes=None, keep_dims=True) + expected_outputs = np.array([[[[3.5]]]], dtype=np.float32) + + run_compare_builder( + build, + input_placeholders, + input_values, + expected_output_types, + expected_outputs, + use_cpu_only=use_cpu_only, + backend=backend, + ) + @pytest.mark.parametrize( "use_cpu_only, backend, mode", itertools.product([True, False], backends, ["max", "mean"]), @@ -162,6 +191,7 @@ def test_builder_to_backend_global_pool_3d(self, use_cpu_only, backend, mode): backend=backend, ) + @pytest.mark.parametrize( ["axis", "keep_dims"], itertools.product([1, -3], [True, False]) ) 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 65f6542b1..becf4f65f 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_scatter_gather.py +++ b/coremltools/converters/mil/mil/ops/tests/test_scatter_gather.py @@ -452,10 +452,6 @@ class TestGatherAlongAxis: "use_cpu_only, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_only, backend): - - if backend[0] == "mlprogram" and not use_cpu_only: - pytest.xfail("rdar://80710279 (Crash in GatherAlongAxis unit test on MIL GPU context)") - x = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32) indices = np.array([[1, 0, 1], [1, 1, 0]], dtype=np.int32) input_placeholders = { @@ -518,8 +514,7 @@ def test_builder_eval(self): ) def test_builder_to_backend_programmatic(self, use_cpu_for_conversion, backend, rank_axis): if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("rdar://78343225 ((MIL GPU) Core ML Tools Unit Test failures [numerical error])") - + pytest.xfail("rdar://97398875 (TestGatherAlongAxis failing on mlprgram + GPU)") rank, axis = rank_axis x_shape = np.random.randint(low=2, high=8, size=rank) indices_shape = np.copy(x_shape) 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 97caac4ee..cebd9f40f 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_tensor_operation.py +++ b/coremltools/converters/mil/mil/ops/tests/test_tensor_operation.py @@ -6,6 +6,7 @@ import numpy as np import pytest +import coremltools as ct from coremltools.converters.mil import testing_reqs from coremltools.converters.mil.mil import ( Builder as mb, @@ -1153,6 +1154,53 @@ def build(x): backend=backend, ) + @pytest.mark.parametrize( + "use_cpu_only, backend, return_indices, sort", + itertools.product( + [True, False], + backends, + [True, False], + [True, False], + ) + ) + def test_builder_to_backend_smoke_iOS16(self, use_cpu_only, backend, return_indices, sort): + if backend[0] == "neuralnetwork": + pytest.skip("nn backend not supported") + + if not return_indices: + pytest.xfail("rdar://92880117 (Topk with return_indices = False error out at the MIL->EIR stage)") + + val = np.array([[-1.0, 2.0, -3.0], [4.0, -5.0, 6.0]], dtype=np.float32) + input_placeholders = {"x": mb.placeholder(shape=val.shape)} + input_values = {"x": val} + + def build(x): + return mb.topk(x=x, k=2, axis=1, return_indices=return_indices, sort=sort) + + expected_output_types = [ + (2, 2, types.fp32), + (2, 2, types.int32), + ] + expected_outputs = [ + np.array([[2.0, -1.0], [6.0, 4.0]], dtype=np.float32), + np.array([[1, 0], [2, 0]], dtype=np.float32), + ] + + if not return_indices: + expected_output_types = expected_output_types[:1] + expected_outputs = expected_outputs[:1] + + run_compare_builder( + build, + input_placeholders, + input_values, + expected_output_types, + expected_outputs, + use_cpu_only=use_cpu_only, + backend=backend, + minimum_deployment_target=ct.target.iOS16, + ) + @ssa_fn def test_builder_eval(self): def np_topk(x, k, axis, ascending=False): 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 a0b9ad910..a498edad8 100644 --- a/coremltools/converters/mil/mil/ops/tests/test_tensor_transformation.py +++ b/coremltools/converters/mil/mil/ops/tests/test_tensor_transformation.py @@ -7,6 +7,7 @@ import numpy as np from .testing_utils import UNK_SYM, UNK_VARIADIC, run_compare_builder +import coremltools as ct from coremltools._deps import _HAS_TORCH from coremltools.converters.mil import testing_reqs from coremltools.converters.mil.mil import ( @@ -683,9 +684,6 @@ class TestSqueeze: "use_cpu_for_conversion, backend", itertools.product([True, False], backends,) ) def test_builder_to_backend_smoke(self, use_cpu_for_conversion, backend): - if backend[0] == "mlprogram" and not use_cpu_for_conversion: - pytest.xfail("rdar://78343225 ((MIL GPU) Core ML Tools Unit Test failures [numerical error])") - x = np.array([[[[1], [2], [3]]]], dtype=np.float32) input_placeholders = {"x": mb.placeholder(shape=x.shape)} @@ -886,6 +884,72 @@ def build(x): backend=backend, ) +class TestPixelUnshuffle: + @pytest.mark.parametrize( + "use_cpu_only, backend", itertools.product([True, False], backends,) + ) + def test_builder_to_backend_smoke(self, use_cpu_only, backend): + if backend[0] == "neuralnetwork": + pytest.skip("nn backend not supported") + + val = np.array([[[[9.0, 5.0], [1.0, 3.0]]]], dtype=np.float32) + input_placeholders = {"x": mb.placeholder(shape=val.shape)} + input_values = {"x": val} + + def build(x): + return [mb.pixel_unshuffle(x=x, downscale_factor=np.uint32(2))] + + expected_output_types = (1, 4, 1, 1, types.fp32) + expected_outputs = np.array([[[[9.0]], [[5.0]], [[1.0]], [[3.0]]]], dtype=np.float32) + + run_compare_builder( + build, + input_placeholders, + input_values, + expected_output_types, + expected_outputs, + use_cpu_only=use_cpu_only, + backend=backend, + minimum_deployment_target=ct.target.iOS16, + ) + + @pytest.mark.skipif(not testing_reqs._HAS_TORCH, reason="PyTorch not found.") + @pytest.mark.parametrize( + "use_cpu_only, backend, shape, downscale_factor", + itertools.product( + [True, False], + backends, + [(1, 2, 4, 4), (2, 1, 8, 4)], + [2, 4], + ), + ) + def test_builder_to_backend_stress( + self, use_cpu_only, backend, shape, downscale_factor, + ): + if backend[0] == "neuralnetwork": + pytest.skip("nn backend not supported") + + val = np.random.rand(*shape) + input_placeholders = {"x": mb.placeholder(shape=val.shape)} + input_values = {"x": val} + + def build(x): + return [mb.pixel_unshuffle(x=x, downscale_factor=np.uint32(downscale_factor))] + + torch_pixel_unshuffle = torch.nn.PixelUnshuffle(downscale_factor) + expected_outputs = [torch_pixel_unshuffle(torch.Tensor(val)).numpy()] + expected_output_types = [o.shape[:] + (types.fp32,) for o in expected_outputs] + run_compare_builder( + build, + input_placeholders, + input_values, + expected_output_types, + expected_outputs, + use_cpu_only=use_cpu_only, + backend=backend, + minimum_deployment_target=ct.target.iOS16, + ) + class TestSlidingWindows: @pytest.mark.parametrize( diff --git a/coremltools/converters/mil/mil/ops/tests/testing_utils.py b/coremltools/converters/mil/mil/ops/tests/testing_utils.py index 6831399ef..a97111420 100644 --- a/coremltools/converters/mil/mil/ops/tests/testing_utils.py +++ b/coremltools/converters/mil/mil/ops/tests/testing_utils.py @@ -71,7 +71,7 @@ def run_compare_builder( expected_outputs = [expected_outputs] prog = Program() - with Function(input_placeholders) as ssa_func: + with Function(input_placeholders, opset_version=minimum_deployment_target) as ssa_func: output_vars = build(**ssa_func.inputs) if isinstance(output_vars, tuple): output_vars = list(output_vars) 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 5bb03418f..68465276d 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 @@ -4,8 +4,9 @@ # 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 from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil.types.symbolic import any_symbolic @@ -33,6 +34,7 @@ def _match_pattern(op): and op.output_shape is None \ and not any_symbolic(op.outputs[0].shape) +@block_context_manager def _handle_block(block): for op in list(block.operations): for b in op.blocks: @@ -42,12 +44,11 @@ def _handle_block(block): continue # matched pattern - with block: - x = mb.conv_transpose( - **op.inputs, output_shape=op.outputs[0].shape, \ - name=op.name+'_has_output_shape', before_op=op - ) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=x - ) - block.remove_ops([op]) + x = mb.conv_transpose( + **op.inputs, output_shape=op.outputs[0].shape, \ + name=op.name+'_has_output_shape', before_op=op + ) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=x + ) + block.remove_ops([op]) 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 6866121e9..a1d593997 100644 --- a/coremltools/converters/mil/mil/passes/apply_common_pass_pipeline.py +++ b/coremltools/converters/mil/mil/passes/apply_common_pass_pipeline.py @@ -86,6 +86,8 @@ def _apply(passes, name="common"): _apply([p], type(p).__name__) cleanup_passes = [ + "common::dead_code_elimination", + "common::const_elimination", "common::cast_optimization", "common::const_elimination", "common::loop_invariant_elimination", diff --git a/coremltools/converters/mil/mil/passes/cast_optimization.py b/coremltools/converters/mil/mil/passes/cast_optimization.py index c8835537f..8e4a55d2e 100644 --- a/coremltools/converters/mil/mil/passes/cast_optimization.py +++ b/coremltools/converters/mil/mil/passes/cast_optimization.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 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.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass @register_pass(namespace="common") class cast_optimization(AbstractGraphPass): @@ -31,32 +32,7 @@ class cast_optimization(AbstractGraphPass): """ 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) + _fuse_or_cancel_consecutive_casts_block_wrapper(f) # main function's output_vars are treated differently, which are not handled by the method # above, "_fuse_or_cancel_consecutive_casts_block". @@ -65,6 +41,7 @@ def apply(self, prog): while block_changed: block_changed = _cancel_consecutive_casts_connected_to_outputs(prog.functions["main"]) + class Node: def __init__(self, op_type, match_criterion=None): """ @@ -174,28 +151,56 @@ def _try_to_transform(root_op, cached_vars): cast_2.enclosing_block.remove_ops([cast_2]) return True +@block_context_manager +def _fuse_or_cancel_consecutive_casts_block_wrapper(block): -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: - nested_block_changed = True - 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) - - if len(op.blocks) > 0: - continue + 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: + nested_block_changed = True + 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) - # start pattern match if cast op is encountered - if op.op_type == "cast": - with block: + if len(op.blocks) > 0: + continue + + # start pattern match if cast op is encountered + if op.op_type == "cast": 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 + # has to break as the downstream iterator is affected. + if block_changed: + return block_changed + return block_changed + + 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(block, cached_vars) def _cancel_consecutive_casts_connected_to_outputs(block): diff --git a/coremltools/converters/mil/mil/passes/compression_passes.py b/coremltools/converters/mil/mil/passes/compression_passes.py index f6008b8c6..e6153e38f 100644 --- a/coremltools/converters/mil/mil/passes/compression_passes.py +++ b/coremltools/converters/mil/mil/passes/compression_passes.py @@ -1,43 +1,23 @@ -# Copyright (c) 2022, Apple Inc. All rights reserved. +# Copyright (c) 2022, 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 enum import Enum import logging as _logging +import numpy as np + from coremltools.converters.mil.backend.mil.load import should_use_weight_file from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil.passes.quantization_passes import AbstractQuantizationPass from coremltools.models.neural_network.quantization_utils import _get_kmeans_lookup_table_and_weight -from coremltools.converters.mil.mil.ops.defs.constexpr_ops import ( +from coremltools.converters.mil.mil.ops.defs.iOS16 import ( constexpr_affine_dequantize, constexpr_lut_to_dense, constexpr_sparse_to_dense, ) -class CompressionWeightMode(Enum): - @classmethod - def has_mode(cls, value): - if not isinstance(value, str): - return False - return value.upper() in cls._member_names_ - - @classmethod - def get_mode(cls): - return list(cls._member_names_) - -class SparseMode(CompressionWeightMode): - THRESHOLD_BASED = 1 - PERCENTILE_BASED = 2 - -class AffineQuantizeMode(CompressionWeightMode): - LINEAR = 1 - LINEAR_SYMMETRIC = 2 - -class PalettizeMode(CompressionWeightMode): - KMEANS = 1 - UNIFORM = 2 - UNIQUE = 3 - CUSTOM = 4 class SparseParams: def __init__(self, nonzero_data=None, mask=None, shape=None): @@ -54,20 +34,22 @@ class WeightSparsifier(AbstractQuantizationPass): - If fake_compression=True, Zeroed-Out Value is encoded via const op - Old const is replaced by a new operation with zeroed-out value. """ + WEIGHT_SPARSIFICATION_MODES = ["THRESHOLD_BASED", "PERCENTILE_BASED"] def __init__(self, mode="threshold_based", threshold=1e-3, target_percentile=1.0, fake_compression=False, op_selector=None): super().__init__(op_selector=op_selector) self.fake_compression = fake_compression - self.mode = mode + self.mode = mode.upper() self.threshold = threshold self.target_percentile = target_percentile - if not SparseMode.has_mode(self.mode): - msg = "Only mode {} supported for weight sparsification. Got mode {}.".format(SparseMode.get_mode(), self.mode) + if not self.mode in WeightSparsifier.WEIGHT_SPARSIFICATION_MODES: + msg = "Only mode {} supported for weight sparsification. Got mode {}.".format( + WeightSparsifier.WEIGHT_SPARSIFICATION_MODES, + self.mode + ) raise ValueError(msg) - self.mode = SparseMode[self.mode.upper()] - if self.target_percentile < 0 or self.target_percentile > 1: raise ValueError("Invalid value of target_percentile: {}. Needs to be in [0, 1]".format(self.target_percentile)) @@ -80,7 +62,9 @@ def is_valid_op(self, op): return False @staticmethod - def compress(val, mode, target_percentile, threshold): + def compress(val, mode, target_percentile=None, threshold=None): + + mode = mode.upper() def sparsify_with_percentile(val, target_percentile): q = target_percentile * 100 @@ -94,9 +78,9 @@ def sparsify_with_thresohld(val, threshold): flattened_val = val.flatten() - if mode == SparseMode.PERCENTILE_BASED: + if mode == "PERCENTILE_BASED": flattened_val = sparsify_with_percentile(flattened_val, target_percentile) - elif mode == SparseMode.THRESHOLD_BASED: + elif mode == "THRESHOLD_BASED": flattened_val = sparsify_with_thresohld(flattened_val, threshold) params = SparseParams() @@ -115,31 +99,30 @@ def transform_op(self, op): block = op.enclosing_block sparse_params = self.compress(op.val.val, self.mode, self.target_percentile, self.threshold) - with block: - if not self.fake_compression: - new_var = mb.constexpr_sparse_to_dense( - nonzero_data=sparse_params.nonzero_data, - mask=sparse_params.mask, - shape=np.uint32(sparse_params.shape), - before_op=op, - name=op.name + "_sparsified", - ) - else: - decompressed_val = self.decompress(sparse_params) - new_var = mb.const( - val=decompressed_val, - before_op=op, - name=op.name + "_fake_sparsified", - ) - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, - old_var=op.outputs[0], - new_var=new_var, - no_check_var_types=True, + if not self.fake_compression: + new_var = mb.constexpr_sparse_to_dense( + nonzero_data=sparse_params.nonzero_data, + mask=sparse_params.mask, + shape=np.uint32(sparse_params.shape), + before_op=op, + name=op.name + "_sparsified", ) - - block.remove_ops([op]) + else: + decompressed_val = self.decompress(sparse_params) + new_var = mb.const( + val=decompressed_val, + before_op=op, + name=op.name + "_fake_sparsified", + ) + + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, + old_var=op.outputs[0], + new_var=new_var, + no_check_var_types=True, + ) + + block.remove_ops([op]) class LutParams: @@ -157,36 +140,37 @@ class WeightPalettizer(AbstractQuantizationPass): - If fake_compression=True, compressed value is decompressed and then encoded via const op - Old const op is replaced by a newly created operation. """ - + WEIGHT_PALETTIZATION_MODES = ["KMEANS", "UNIFORM", "UNIQUE", "CUSTOM"] def __init__(self, nbits, fake_compression=False, op_selector=None, mode="kmeans", lut_function=None): super().__init__(op_selector=op_selector) self.fake_compression = fake_compression self.nbits = nbits - self.mode = mode + self.mode = mode.upper() self.lut_function = lut_function - if not PalettizeMode.has_mode(self.mode): - msg = "Only mode {} supported for weight palettization. Got mode {}.".format(PalettizeMode.get_mode(), self.mode) + if not self.mode in WeightPalettizer.WEIGHT_PALETTIZATION_MODES: + msg = "Only mode {} supported for weight palettization. Got mode {}.".format( + WeightPalettizer.WEIGHT_PALETTIZATION_MODES, + self.mode + ) raise ValueError(msg) - self.mode = PalettizeMode[self.mode.upper()] - - if nbits is None and self.mode in (PalettizeMode.KMEANS, PalettizeMode.UNIFORM): + if nbits is None and self.mode in ("KMEANS", "UNIFORM"): msg = "nbits must be provided for mode {}".format(mode) raise ValueError(msg) - if nbits is not None and self.mode in (PalettizeMode.UNIQUE, PalettizeMode.CUSTOM): + if nbits is not None and self.mode in ("UNIQUE", "CUSTOM"): msg = "nbits must NOT be provided for mode {}".format(mode) raise ValueError(msg) if self.nbits is not None and self.nbits not in (1, 2, 4, 6, 8): raise ValueError("Invalid value of nbits ({}) for palettization. Supported bits are {1, 2, 4, 6, 8}".format(nbits)) - if (self.mode == PalettizeMode.CUSTOM) ^ (lut_function is not None): + if (self.mode == "CUSTOM") ^ (lut_function is not None): msg = "lut_function must be None if mode is not custom, and that it cannot be None when the mode is custom." raise ValueError(msg) - if self.mode == PalettizeMode.CUSTOM and not callable(self.lut_function): + if self.mode == "CUSTOM" and not callable(self.lut_function): msg = "A function object must be provided as lut_function. Got a lut_functions as type {}".format(type(self.lut_function)) raise ValueError(msg) @@ -196,7 +180,9 @@ def is_valid_op(self, op): return False @staticmethod - def compress(val, nbits, mode, lut_function): + def compress(val, mode, nbits=None, lut_function=None): + + mode = mode.upper() def compress_kmeans(val, nbits): lut, indices = _get_kmeans_lookup_table_and_weight(nbits, val) @@ -269,16 +255,16 @@ def check_lut_parameters_are_valid(val, lut, indices): if not isinstance(val, (np.ndarray, np.generic)): raise ValueError("Only numpy arrays are supported") - if mode == PalettizeMode.KMEANS: + if mode == "KMEANS": lut, indices = compress_kmeans(val, nbits) - elif mode == PalettizeMode.UNIFORM: + elif mode == "UNIFORM": lut, indices = compress_uniform(val, nbits) - elif mode == PalettizeMode.UNIQUE: + elif mode == "UNIQUE": nbits = get_nbits_for_unique_mode(val) if nbits is None: return None lut, indices = compress_unique(val, nbits) - elif mode == PalettizeMode.CUSTOM: + elif mode == "CUSTOM": lut, indices = lut_function(val) check_lut_parameters_are_valid(val, lut, indices) @@ -298,36 +284,35 @@ def decompress(params): def transform_op(self, op): block = op.enclosing_block - lut_params = self.compress(op.val.val, self.nbits, self.mode, self.lut_function) + lut_params = self.compress(op.val.val, self.mode, self.nbits, self.lut_function) if lut_params is None: return - with block: - if not self.fake_compression: - new_var = mb.constexpr_lut_to_dense( - indices=lut_params.indices, - lut=lut_params.lut, - shape=np.uint32(lut_params.shape), - before_op=op, - name=op.name + "_palettized", - ) - else: - decompressed_val = self.decompress(lut_params) - new_var = mb.const( - val=decompressed_val, - before_op=op, - name=op.name + "_fake_palettized", - ) - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, - old_var=op.outputs[0], - new_var=new_var, - no_check_var_types=True, + if not self.fake_compression: + new_var = mb.constexpr_lut_to_dense( + indices=lut_params.indices, + lut=lut_params.lut, + shape=np.uint32(lut_params.shape), + before_op=op, + name=op.name + "_palettized", + ) + else: + decompressed_val = self.decompress(lut_params) + new_var = mb.const( + val=decompressed_val, + before_op=op, + name=op.name + "_fake_palettized", ) - block.remove_ops([op]) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, + old_var=op.outputs[0], + new_var=new_var, + no_check_var_types=True, + ) + + block.remove_ops([op]) class AffineQuantParams: @@ -346,17 +331,19 @@ class WeightAffineQuantizer(AbstractQuantizationPass): - If fake_compression=True, compressed value is decompressed and then encoded via const op - Old const is replaced by a newly created operation. """ - + WEIGHT_AFFINE_QUANTIZATION_MODES = ["LINEAR_SYMMETRIC", "LINEAR"] def __init__(self, fake_compression=False, op_selector=None, mode="linear"): super().__init__(op_selector=op_selector) self.fake_compression = fake_compression - self.mode = mode + self.mode = mode.upper() - if not AffineQuantizeMode.has_mode(self.mode): + if not self.mode in WeightAffineQuantizer.WEIGHT_AFFINE_QUANTIZATION_MODES: msg = "Only mode {} supported for weight affine quantization. Got mode {}."\ - .format(AffineQuantizeMode.get_mode(), self.mode) + .format( + WeightAffineQuantizer.WEIGHT_AFFINE_QUANTIZATION_MODES, + self.mode + ) raise ValueError(msg) - self.mode = AffineQuantizeMode[self.mode.upper()] def is_valid_op(self, op): if op.op_type == "const" and should_use_weight_file(op.val.val): @@ -373,6 +360,7 @@ def _get_axis(op): @staticmethod def compress(val, axis, mode): + mode = mode.upper() if not isinstance(val, (np.ndarray, np.generic)): raise ValueError("Only numpy arrays are supported") @@ -382,7 +370,7 @@ def compress(val, axis, mode): val_max = np.amax(val, axis=axes, keepdims=True) val_range = 255 - if mode == AffineQuantizeMode.LINEAR_SYMMETRIC: + if mode == "LINEAR_SYMMETRIC": # For the linear_symmetric mode, the range is symmetrical to 0 max_abs = np.maximum(np.abs(val_min), np.abs(val_max)) val_min = -max_abs @@ -410,32 +398,31 @@ def transform_op(self, op): block = op.enclosing_block quant_params = self.compress(op.val.val, self._get_axis(op), self.mode) - with block: - if not self.fake_compression: - new_var = mb.constexpr_affine_dequantize( - quantized_data=quant_params.quantized_data, - zero_point=quant_params.zero_point, - scale=quant_params.scale, - axis=quant_params.axis, - before_op=op, - name=op.name + "_affine_quantized", - ) - else: - decompressed_val = self.decompress(quant_params) - new_var = mb.const( - val=decompressed_val, - before_op=op, - name=op.name + "_fake_affine_quantized", - ) - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, - old_var=op.outputs[0], - new_var=new_var, - no_check_var_types=True, + if not self.fake_compression: + new_var = mb.constexpr_affine_dequantize( + quantized_data=quant_params.quantized_data, + zero_point=quant_params.zero_point, + scale=quant_params.scale, + axis=quant_params.axis, + before_op=op, + name=op.name + "_affine_quantized", + ) + else: + decompressed_val = self.decompress(quant_params) + new_var = mb.const( + val=decompressed_val, + before_op=op, + name=op.name + "_fake_affine_quantized", ) - block.remove_ops([op]) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, + old_var=op.outputs[0], + new_var=new_var, + no_check_var_types=True, + ) + + block.remove_ops([op]) class WeightDecompressor(AbstractQuantizationPass): @@ -456,19 +443,19 @@ def is_valid_op(self, op): def transform_op(self, op): block = op.enclosing_block - with block: - decompressed_val = op.get_decompressed_value() - new_var = mb.const( - val=decompressed_val, - before_op=op, - name=op.name, - ) - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, - old_var=op.outputs[0], - new_var=new_var, - no_check_var_types=True, - ) + decompressed_val = op.value_inference() + new_var = mb.const( + val=decompressed_val, + before_op=op, + name=op.name, + ) - block.remove_ops([op]) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, + old_var=op.outputs[0], + new_var=new_var, + no_check_var_types=True, + force_replace=True, + ) + + block.remove_ops([op]) 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 eebd25234..ba5373f77 100644 --- a/coremltools/converters/mil/mil/passes/concat_to_pixel_shuffle.py +++ b/coremltools/converters/mil/mil/passes/concat_to_pixel_shuffle.py @@ -4,8 +4,9 @@ # 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 from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _match_pattern(op): @@ -57,27 +58,25 @@ def _match_pattern(op): def _replace_ops(block, w_concat, h_concat_0, h_concat_1): - with block: - h_concat_0_inputs = list(h_concat_0.inputs["values"]) - h_concat_1_inputs = list(h_concat_1.inputs["values"]) - - all_inputs = [h_concat_0_inputs[0], h_concat_1_inputs[0], h_concat_0_inputs[1], h_concat_1_inputs[1]] - - # Concatenate all 4 inputs on the channel axis - x = mb.concat(values=all_inputs, axis=1, before_op=h_concat_0, interleave=True) - # Shuffle into place - x = mb.pixel_shuffle(x=x, upscale_factor=2, before_op=h_concat_0) + h_concat_0_inputs = list(h_concat_0.inputs["values"]) + h_concat_1_inputs = list(h_concat_1.inputs["values"]) - w_concat.enclosing_block.replace_uses_of_var_after_op( - anchor_op=h_concat_0, old_var=w_concat.outputs[0], new_var=x - ) + all_inputs = [h_concat_0_inputs[0], h_concat_1_inputs[0], h_concat_0_inputs[1], h_concat_1_inputs[1]] - block.remove_ops([w_concat, h_concat_0, h_concat_1]) + # Concatenate all 4 inputs on the channel axis + x = mb.concat(values=all_inputs, axis=1, before_op=h_concat_0, interleave=True) + # Shuffle into place + x = mb.pixel_shuffle(x=x, upscale_factor=2, before_op=h_concat_0) + w_concat.enclosing_block.replace_uses_of_var_after_op( + anchor_op=h_concat_0, old_var=w_concat.outputs[0], new_var=x + ) + block.remove_ops([w_concat, h_concat_0, h_concat_1]) + +@block_context_manager def _concat_to_pixel_shuffle_block(block): for op in list(block.operations): - layers = _match_pattern(op) if layers: _replace_ops(block, layers[0], layers[1], layers[2]) diff --git a/coremltools/converters/mil/mil/passes/const_elimination.py b/coremltools/converters/mil/mil/passes/const_elimination.py index ebb3d993b..82080631a 100644 --- a/coremltools/converters/mil/mil/passes/const_elimination.py +++ b/coremltools/converters/mil/mil/passes/const_elimination.py @@ -2,9 +2,48 @@ # # 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 from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass + +@block_context_manager +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_replaced = True + for i, o in enumerate(op.outputs): + if o.val is not None: + res = mb.const( + val=o.val, + before_op=op, + # same var name, but different python + # instance does not violate SSA property. + name=o.name, + ) + + if op.enclosing_block.try_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_replaced = False + else: + all_outputs_are_replaced = False + + if all_outputs_are_replaced: + op.remove_from_block() @register_pass(namespace="common") class const_elimination(AbstractGraphPass): @@ -24,46 +63,6 @@ class const_elimination(AbstractGraphPass): # %4 = other_op(%2_const, %3) # """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass - - def _const_elimination_block(self, block): - # shallow copy hides changes on f.operations during the loop - for op in list(block.operations): - - if op.op_type == "const" or op in self.ops_to_skip: - continue - - for b in op.blocks: - self._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() - - def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): - self._const_elimination_block(f) + _const_elimination_block(f) diff --git a/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py b/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py index 3481c9ba1..471a4fd35 100644 --- a/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py +++ b/coremltools/converters/mil/mil/passes/conv_batchnorm_fusion.py @@ -7,6 +7,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -105,12 +106,49 @@ def _try_to_transform(conv_op, bn_op, block): else: x = mb.conv(**conv_kargs) - bn_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=bn_op, old_var=bn_op.outputs[0], new_var=x - ) - # Remove all the ops at once - block.remove_ops([conv_op, bn_op]) - return True + if bn_op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=bn_op, + old_var=bn_op.outputs[0], + new_var=x, + ): + bn_op.enclosing_block.remove_ops([conv_op, bn_op]) + return True + return False + +@block_context_manager +def _fuse_conv_batchnorm_block(block): + + def _match_pattern(op): + if op.op_type == "conv" or op.op_type == "conv_transpose": + # abort fusion if op output is also a block output + if op.outputs[0] in op.enclosing_block.outputs: + return None + # find batch_norm op + child_ops = op.outputs[0].child_ops + if len(child_ops) == 1: + bn_op_candidate = list(child_ops)[0] + if bn_op_candidate.op_type == "batch_norm": + return bn_op_candidate + return None + + fusion_occurred = False + for op in list(block.operations): + for b in op.blocks: + block_changed = True + while block_changed: + block_changed = _fuse_conv_batchnorm_block(b) + if len(op.blocks) > 0: + # This op can't be conv or conv_transpose + continue + + bn_op = _match_pattern(op) + if bn_op is not None: + fusion_occurred = _try_to_transform(op, bn_op, block) + # has to break as the downstream iterator is affected. + if fusion_occurred: + return fusion_occurred + return fusion_occurred + @register_pass(namespace="common") class fuse_conv_batchnorm(AbstractGraphPass): @@ -128,50 +166,9 @@ class fuse_conv_batchnorm(AbstractGraphPass): ... """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass - - def _fuse_conv_batchnorm_block(self, block): - - def _match_pattern(op): - if op.op_type == "conv" or op.op_type == "conv_transpose": - # abort fusion if op output is also a block output - if op.outputs[0] in op.enclosing_block.outputs: - return None - # find batch_norm op - child_ops = op.outputs[0].child_ops - if len(child_ops) == 1: - bn_op_candidate = list(child_ops)[0] - if bn_op_candidate.op_type == "batch_norm": - return bn_op_candidate - return None - - fusion_occurred = False - for op in list(block.operations): - for b in op.blocks: - block_changed = True - while block_changed: - block_changed = self._fuse_conv_batchnorm_block(b) - if len(op.blocks) > 0: - # This op can't be conv or conv_transpose - continue - if op in self.ops_to_skip: - continue - bn_op = _match_pattern(op) - if bn_op is not None: - with block: - fusion_occurred = _try_to_transform(op, bn_op, block) - # has to break as the downstream iterator is affected. - if fusion_occurred: - return fusion_occurred - return fusion_occurred def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): block_changed = True while block_changed: - block_changed = self._fuse_conv_batchnorm_block(f) + 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 3ed8b1927..6bf2c542d 100644 --- a/coremltools/converters/mil/mil/passes/conv_bias_fusion.py +++ b/coremltools/converters/mil/mil/passes/conv_bias_fusion.py @@ -6,10 +6,11 @@ import logging import numpy as np -from .helper import _check_child_op_type +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.helper import _check_child_op_type, block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass -from coremltools.converters.mil.mil import Builder as mb, types + child_op_types = ["add", "sub"] @@ -26,6 +27,253 @@ def _match_pattern(op): if add_op_candidate.op_type in child_op_types: return add_op_candidate return None + +def _try_to_transform_transpose_pattern(conv_op, block): + + ops_to_remove = [] + + # conv layer + if conv_op.op_type != "conv" and conv_op.op_type != "conv_transpose": + return False + is_deconv = conv_op.op_type == "conv_transpose" + ops_to_remove.append(conv_op) + + # transpose layer + if not _check_child_op_type(conv_op, "transpose"): + return False + transpose_op = list(conv_op.outputs[0].child_ops)[0] + ops_to_remove.append(transpose_op) + + # add/sub layer + if not _check_child_op_type(transpose_op, "add") and not _check_child_op_type(transpose_op, "sub"): + return False + add_or_sub_op = list(transpose_op.outputs[0].child_ops)[0] + + ops_to_remove.append(add_or_sub_op) + + # get the bias + if add_or_sub_op.x.val is None and add_or_sub_op.y.val is None: + return False + bias = add_or_sub_op.x.val if add_or_sub_op.x.val is not None else add_or_sub_op.y.val + is_first_input = add_or_sub_op.y.val is not None + is_sub = add_or_sub_op.op_type == "sub" + + + # get the conv bias/weight + conv_shape = conv_op.outputs[0].shape + Cout = conv_shape[1] + conv_weight = conv_op.weight.val + conv_weight_type = conv_weight.dtype + conv_bias = np.zeros(Cout).astype(conv_weight_type) if conv_op.bias is None else conv_op.bias.val + + # check if the bias is compatible for fusion + is_bias_scalar = True + if isinstance(bias, np.ndarray): + if bias.shape == (): + bias = bias.tolist() + elif np.prod(bias.shape) == 1: + bias = np.squeeze(bias).tolist() + else: + is_bias_scalar = False + + if not is_bias_scalar: + if np.prod(bias.shape) != Cout: + return False + rank = transpose_op.outputs[0].rank + cout_dim = transpose_op.perm.val.tolist().index(1) - rank + if bias.shape[cout_dim] != Cout: + return False + bias = np.reshape(bias, (Cout)) + + # compute the new bias + if is_sub: + if is_first_input: + bias = -bias + else: + conv_bias = -conv_bias + + new_bias = conv_bias + bias + + # compute the new weight + if is_sub and not is_first_input: + new_weight = -conv_weight + else: + new_weight = conv_weight + + # check that none of the op in this pattern is connected to the output + # (except the last op) + for op in ops_to_remove[:-1]: + for out in op.outputs: + if out in block.outputs: + return False + + # create a new conv op with the new weight, bias value, copying rest of the attributes + conv_kargs = {"weight": new_weight, "bias": new_bias, "before_op": conv_op} + + for k, v in conv_op.inputs.items(): + if k in ["weight", "bias"]: + continue + conv_kargs[k] = v + + if is_deconv: + x = mb.conv_transpose(**conv_kargs) + else: + x = mb.conv(**conv_kargs) + + # create a new transpose op + out_name = add_or_sub_op.outputs[0].name + tranpose_kargs = {"x": x, "name": out_name, "before_op": transpose_op} + for k, v in transpose_op.inputs.items(): + if k == "x": + continue + tranpose_kargs[k] = v + x = mb.transpose(**tranpose_kargs) + + if add_or_sub_op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=add_or_sub_op, + old_var=add_or_sub_op.outputs[0], + new_var=x, + ): + add_or_sub_op.enclosing_block.remove_ops(ops_to_remove) + return True + return False + + +def _try_to_transform(conv_op, add_op, block): + + if add_op.op_type == "sub": + bias_var = add_op.y + else: + bias_var = add_op.x if add_op.x.val is not None else add_op.y + bias_value = bias_var.val + + is_conv_op = (conv_op.op_type == "conv") + + # check that the bias value is a constant array or a scalar constant + if not isinstance(bias_value, (np.ndarray, np.generic)): + return False + + is_bias_scalar = False + if not isinstance(bias_value, np.ndarray): + is_bias_scalar = True + + # find rank of the conv input + rank = conv_op.x.rank + if rank is None: + return False + if not (rank == 3 or rank == 4 or rank == 5): + return False + + # check compatibility of bias value with the rank of the conv op + # either bias value should be a scalar or: + # rank=3 ==> (B,C,D), which means bias must be (1,C,1) or (C,1) + # rank=4 ==> (B,C,D1,D2), which means bias must be (1,C,1,1) or (C,1,1) + # rank=5 ==> (B,C,D1,D2,D3), which means bias must be (1,C,1,1,1) or (C,1,1,1) + + if is_bias_scalar: + bias_value = np.array([bias_value]) + else: + # check that there is at most one dimension in the shape that is not 1 + if len(np.squeeze(bias_value).shape) > 1: + return False + # check that addition is not happening on the batch dimension + if len(bias_value) == rank: + if bias_value.shape[0] != 1: + return False + # check that last rank-2 entries in the shape vector are all 1s + if np.prod(bias_value.shape[-(rank - 2) :]) != 1: + return False + bias_value = np.squeeze(bias_value) + + if add_op.op_type == "sub": + bias_value *= -1 + + # everything looks good, now find the new updated bias + old_bias = conv_op.inputs.get("bias", None) + old_bias_value = None + if old_bias is not None and old_bias.val is not None: + old_bias_value = old_bias.val + if old_bias is None: + # need to create a fresh numpy array for bias + if np.prod(bias_value.shape) == 1: + # its a scalar bias + # need to find the value of Cout to form a new bias + if conv_op.weight.val is None: + return False + # conv_transpose has weight format [K, C_out, spatial dims] + # conv has weight format [C_out, K, spatial dims] + Cout = conv_op.weight.val.shape[0 if is_conv_op else 1] + new_bias_value = np.broadcast_to(bias_value, (Cout,)) + else: + new_bias_value = bias_value + else: + # just need to update the existing bias array + try: + new_bias_value = old_bias_value + bias_value + except: + return False + + # create a new conv op with the new bias value, copying rest of the attributes + out_name = add_op.outputs[0].name + if new_bias_value.dtype != np.float32 and new_bias_value.dtype != np.float16: + # cast the bias to match the weight type + weight_np_type = types.nptype_from_builtin(conv_op.inputs["weight"].sym_type.get_primitive()) + logging.warning("conv_bias_fusion pass: casting bias " + "from {} to {} to match the dtype of the weight of the conv layer".format( + new_bias_value.dtype, weight_np_type + ) + ) + new_bias_value = new_bias_value.astype(weight_np_type) + new_bias_var = mb.const(val=new_bias_value, before_op=conv_op) + + conv_kargs = {"bias": new_bias_var, "name": out_name, "before_op": conv_op} + + for k, v in conv_op.inputs.items(): + if k == "bias": + continue + conv_kargs[k] = v + + if is_conv_op: + x = mb.conv(**conv_kargs) + else: + x = mb.conv_transpose(**conv_kargs) + + if add_op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=add_op, + old_var=add_op.outputs[0], + new_var=x, + ): + add_op.enclosing_block.remove_ops([conv_op, add_op]) + return True + return False + +@block_context_manager +def _fuse_conv_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_conv_bias_block(b) + if len(op.blocks) > 0: + # This op can't be conv or conv_transpose + continue + + # pattern 1 : conv + add/sub + add_op = _match_pattern(op) + if add_op is not None: + fusion_status = _try_to_transform(op, add_op, block) + # has to break as the downstream iterator is affected. + if fusion_status: + return fusion_status + + # pattern 2 : conv + transpose + add/sub + fusion_status = _try_to_transform_transpose_pattern(op, block) + if fusion_status: + return fusion_status + + return fusion_status + @register_pass(namespace="common") class fuse_conv_bias(AbstractGraphPass): @@ -61,266 +309,9 @@ class fuse_conv_bias(AbstractGraphPass): ... """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): block_changed = True while block_changed: - block_changed = self._fuse_conv_bias_block(f) - - def _try_to_transform_transpose_pattern(self, conv_op, block): - if conv_op in self.ops_to_skip: - return False - - ops_to_remove = [] - - # conv layer - if conv_op.op_type != "conv" and conv_op.op_type != "conv_transpose": - return False - is_deconv = conv_op.op_type == "conv_transpose" - ops_to_remove.append(conv_op) - - # transpose layer - if not _check_child_op_type(conv_op, "transpose"): - return False - transpose_op = list(conv_op.outputs[0].child_ops)[0] - ops_to_remove.append(transpose_op) - - # add/sub layer - if not _check_child_op_type(transpose_op, "add") and not _check_child_op_type(transpose_op, "sub"): - return False - add_or_sub_op = list(transpose_op.outputs[0].child_ops)[0] - - if add_or_sub_op in self.ops_to_skip: - return False - - ops_to_remove.append(add_or_sub_op) - - # get the bias - if add_or_sub_op.x.val is None and add_or_sub_op.y.val is None: - return False - bias = add_or_sub_op.x.val if add_or_sub_op.x.val is not None else add_or_sub_op.y.val - is_first_input = add_or_sub_op.y.val is not None - is_sub = add_or_sub_op.op_type == "sub" - - - # get the conv bias/weight - conv_shape = conv_op.outputs[0].shape - Cout = conv_shape[1] - conv_weight = conv_op.weight.val - conv_weight_type = conv_weight.dtype - conv_bias = np.zeros(Cout).astype(conv_weight_type) if conv_op.bias is None else conv_op.bias.val - - # check if the bias is compatible for fusion - is_bias_scalar = True - if isinstance(bias, np.ndarray): - if bias.shape == (): - bias = bias.tolist() - elif np.prod(bias.shape) == 1: - bias = np.squeeze(bias).tolist() - else: - is_bias_scalar = False - - if not is_bias_scalar: - if np.prod(bias.shape) != Cout: - return False - rank = transpose_op.outputs[0].rank - cout_dim = transpose_op.perm.val.tolist().index(1) - rank - if bias.shape[cout_dim] != Cout: - return False - bias = np.reshape(bias, (Cout)) - - # compute the new bias - if is_sub: - if is_first_input: - bias = -bias - else: - conv_bias = -conv_bias - - new_bias = conv_bias + bias - - # compute the new weight - if is_sub and not is_first_input: - new_weight = -conv_weight - else: - new_weight = conv_weight - - # check that none of the op in this pattern is connected to the output - # (except the last op) - for op in ops_to_remove[:-1]: - for out in op.outputs: - if out in block.outputs: - return False - - # create a new conv op with the new weight, bias value, copying rest of the attributes - conv_kargs = {"weight": new_weight, "bias": new_bias, "before_op": conv_op} - - for k, v in conv_op.inputs.items(): - if k in ["weight", "bias"]: - continue - conv_kargs[k] = v - - if is_deconv: - x = mb.conv_transpose(**conv_kargs) - else: - x = mb.conv(**conv_kargs) - - # create a new transpose op - out_name = add_or_sub_op.outputs[0].name - tranpose_kargs = {"x": x, "name": out_name, "before_op": transpose_op} - for k, v in transpose_op.inputs.items(): - if k == "x": - continue - tranpose_kargs[k] = v - x = mb.transpose(**tranpose_kargs) - - add_or_sub_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=add_or_sub_op, old_var=add_or_sub_op.outputs[0], new_var=x - ) - - # Remove all the ops at once - block.remove_ops(ops_to_remove) - return True - - - def _try_to_transform(self, conv_op, add_op, block): - if conv_op in self.ops_to_skip or add_op in self.ops_to_skip: - return False - - if add_op.op_type == "sub": - bias_var = add_op.y - else: - bias_var = add_op.x if add_op.x.val is not None else add_op.y - bias_value = bias_var.val - - is_conv_op = (conv_op.op_type == "conv") - - # check that the bias value is a constant array or a scalar constant - if not isinstance(bias_value, (np.ndarray, np.generic)): - return False - - is_bias_scalar = False - if not isinstance(bias_value, np.ndarray): - is_bias_scalar = True - - # find rank of the conv input - rank = conv_op.x.rank - if rank is None: - return False - if not (rank == 3 or rank == 4 or rank == 5): - return False - - # check compatibility of bias value with the rank of the conv op - # either bias value should be a scalar or: - # rank=3 ==> (B,C,D), which means bias must be (1,C,1) or (C,1) - # rank=4 ==> (B,C,D1,D2), which means bias must be (1,C,1,1) or (C,1,1) - # rank=5 ==> (B,C,D1,D2,D3), which means bias must be (1,C,1,1,1) or (C,1,1,1) - - if is_bias_scalar: - bias_value = np.array([bias_value]) - else: - # check that there is at most one dimension in the shape that is not 1 - if len(np.squeeze(bias_value).shape) > 1: - return False - # check that addition is not happening on the batch dimension - if len(bias_value) == rank: - if bias_value.shape[0] != 1: - return False - # check that last rank-2 entries in the shape vector are all 1s - if np.prod(bias_value.shape[-(rank - 2) :]) != 1: - return False - bias_value = np.squeeze(bias_value) - - if add_op.op_type == "sub": - bias_value *= -1 - - # everything looks good, now find the new updated bias - old_bias = conv_op.inputs.get("bias", None) - old_bias_value = None - if old_bias is not None and old_bias.val is not None: - old_bias_value = old_bias.val - if old_bias is None: - # need to create a fresh numpy array for bias - if np.prod(bias_value.shape) == 1: - # its a scalar bias - # need to find the value of Cout to form a new bias - if conv_op.weight.val is None: - return False - # conv_transpose has weight format [K, C_out, spatial dims] - # conv has weight format [C_out, K, spatial dims] - Cout = conv_op.weight.val.shape[0 if is_conv_op else 1] - new_bias_value = np.broadcast_to(bias_value, (Cout,)) - else: - new_bias_value = bias_value - else: - # just need to update the existing bias array - try: - new_bias_value = old_bias_value + bias_value - except: - return False - - # create a new conv op with the new bias value, copying rest of the attributes - out_name = add_op.outputs[0].name - if new_bias_value.dtype != np.float32 and new_bias_value.dtype != np.float16: - # cast the bias to match the weight type - weight_np_type = types.nptype_from_builtin(conv_op.inputs["weight"].sym_type.get_primitive()) - logging.warning("conv_bias_fusion pass: casting bias " - "from {} to {} to match the dtype of the weight of the conv layer".format( - new_bias_value.dtype, weight_np_type - ) - ) - new_bias_value = new_bias_value.astype(weight_np_type) - new_bias_var = mb.const(val=new_bias_value, before_op=conv_op) - - conv_kargs = {"bias": new_bias_var, "name": out_name, "before_op": conv_op} - - for k, v in conv_op.inputs.items(): - if k == "bias": - continue - conv_kargs[k] = v - - if is_conv_op: - x = mb.conv(**conv_kargs) - else: - x = mb.conv_transpose(**conv_kargs) - - add_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=add_op, old_var=add_op.outputs[0], new_var=x - ) - # Remove all the ops at once - block.remove_ops([conv_op, add_op]) - return True - - def _fuse_conv_bias_block(self, block): - fusion_status = False - for op in list(block.operations): - for b in op.blocks: - block_changed = True - while block_changed: - block_changed = self._fuse_conv_bias_block(b) - if len(op.blocks) > 0: - # This op can't be conv or conv_transpose - continue - - # pattern 1 : conv + add/sub - add_op = _match_pattern(op) - if add_op is not None: - with block: - fusion_status = self._try_to_transform(op, add_op, block) - # has to break as the downstream iterator is affected. - if fusion_status: - return fusion_status - - # pattern 2 : conv + transpose + add/sub - with block: - fusion_status = self._try_to_transform_transpose_pattern(op, block) - if fusion_status: - return fusion_status - - return fusion_status + 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 154534374..3ac5ea18a 100644 --- a/coremltools/converters/mil/mil/passes/conv_scale_fusion.py +++ b/coremltools/converters/mil/mil/passes/conv_scale_fusion.py @@ -5,9 +5,11 @@ import numpy as np +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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass -from coremltools.converters.mil.mil import Builder as mb + def _try_to_transform(conv_op, scale_op, block): @@ -124,12 +126,49 @@ def _try_to_transform(conv_op, scale_op, block): else: x = mb.conv(**conv_kargs) - scale_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=scale_op, old_var=scale_op.outputs[0], new_var=x - ) - # Remove all the ops at once - block.remove_ops([conv_op, scale_op]) - return True + if scale_op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=scale_op, + old_var=scale_op.outputs[0], + new_var=x, + ): + scale_op.enclosing_block.remove_ops([conv_op, scale_op]) + return True + return False + +@block_context_manager +def _fuse_conv_scale_block(block): + def _match_pattern(op): + if op.op_type == "conv" or op.op_type == "conv_transpose": + # abort fusion if op output is also a block output + if op.outputs[0] in op.enclosing_block.outputs: + return None + # find batch_norm op + child_ops = op.outputs[0].child_ops + if len(child_ops) == 1: + scale_op_candidate = list(child_ops)[0] + if scale_op_candidate.op_type in ["mul", "real_div"]: + return scale_op_candidate + return None + + fusion_occurred = False + for op in list(block.operations): + for b in op.blocks: + block_changed = True + while block_changed: + block_changed = _fuse_conv_scale_block(b) + if len(op.blocks) > 0: + # This op can't be conv or conv_transpose + continue + + scale_op = _match_pattern(op) + + if scale_op is not None: + fusion_occurred = _try_to_transform(op, scale_op, block) + # has to break as the downstream iterator is affected. + if fusion_occurred: + return fusion_occurred + return fusion_occurred + @register_pass(namespace="common") class fuse_conv_scale(AbstractGraphPass): @@ -151,53 +190,9 @@ class fuse_conv_scale(AbstractGraphPass): ... """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass - - def _fuse_conv_scale_block(self, block): - - def _match_pattern(op): - if op.op_type == "conv" or op.op_type == "conv_transpose": - # abort fusion if op output is also a block output - if op.outputs[0] in op.enclosing_block.outputs: - return None - # find batch_norm op - child_ops = op.outputs[0].child_ops - if len(child_ops) == 1: - scale_op_candidate = list(child_ops)[0] - if scale_op_candidate.op_type in ["mul", "real_div"]: - return scale_op_candidate - return None - - fusion_occurred = False - for op in list(block.operations): - for b in op.blocks: - block_changed = True - while block_changed: - block_changed = self._fuse_conv_scale_block(b) - if len(op.blocks) > 0: - # This op can't be conv or conv_transpose - continue - - scale_op = _match_pattern(op) - - if op in self.ops_to_skip or scale_op in self.ops_to_skip: - continue - - if scale_op is not None: - with block: - fusion_occurred = _try_to_transform(op, scale_op, block) - # has to break as the downstream iterator is affected. - if fusion_occurred: - return fusion_occurred - return fusion_occurred def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): block_changed = True while block_changed: - block_changed = self._fuse_conv_scale_block(f) + block_changed = _fuse_conv_scale_block(f) diff --git a/coremltools/converters/mil/mil/passes/detect_concat_interleave.py b/coremltools/converters/mil/mil/passes/detect_concat_interleave.py index 58993a454..2055a35be 100644 --- a/coremltools/converters/mil/mil/passes/detect_concat_interleave.py +++ b/coremltools/converters/mil/mil/passes/detect_concat_interleave.py @@ -6,8 +6,9 @@ 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 +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil.types.symbolic import any_symbolic @@ -118,7 +119,7 @@ def _try_to_transform(concat_op, add_op, block): block.remove_ops(all_ops) return True - +@block_context_manager def _fuse_concat_interleave(block): fusion_status = False for op in list(block.operations): @@ -131,8 +132,7 @@ def _fuse_concat_interleave(block): 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 diff --git a/coremltools/converters/mil/mil/passes/divide_to_multiply.py b/coremltools/converters/mil/mil/passes/divide_to_multiply.py index 70a991b98..345555cd0 100644 --- a/coremltools/converters/mil/mil/passes/divide_to_multiply.py +++ b/coremltools/converters/mil/mil/passes/divide_to_multiply.py @@ -5,12 +5,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 +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil import types as _types - +@block_context_manager def _divide_to_multiply_block(block): for op in list(block.operations): for b in op.blocks: @@ -25,18 +26,17 @@ def _divide_to_multiply_block(block): # a floating point number, then the original type # signature (with integer output) would not be preserved. if op.op_type == "real_div" and op.y.val is not None and _types.is_float(op.x.dtype): - with block: - new_y_val = np.array(1.0, dtype=op.y.val.dtype) / op.y.val - if not np.isfinite(new_y_val).all(): - continue - - x = mb.mul( - x=op.x, y=new_y_val, name="_inversed_" + op.name, before_op=op - ) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=x - ) - block.remove_ops([op]) + new_y_val = np.array(1.0, dtype=op.y.val.dtype) / op.y.val + if not np.isfinite(new_y_val).all(): + continue + + x = mb.mul( + x=op.x, y=new_y_val, name="_inversed_" + op.name, before_op=op + ) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=x + ) + block.remove_ops([op]) @register_pass(namespace="common") diff --git a/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py b/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py index c794718d4..7a67269cf 100644 --- a/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py +++ b/coremltools/converters/mil/mil/passes/elementwise_batchnorm_fusion.py @@ -7,6 +7,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -87,7 +88,7 @@ def _try_to_transform(mul_op, add_op, block): block.remove_ops([mul_op, add_op]) return True - +@block_context_manager def _fuse_elementwise_to_batchnorm_block(block): fusion_status = False for op in list(block.operations): @@ -101,8 +102,7 @@ def _fuse_elementwise_to_batchnorm_block(block): 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 diff --git a/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py b/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py index 15d78788d..2cb232b47 100644 --- a/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py +++ b/coremltools/converters/mil/mil/passes/gelu_exact_fusion.py @@ -3,10 +3,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.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 +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import ( + _check_child_op_type, + _check_var_scalar_value, + block_context_manager, +) +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _try_to_transform(op, block): ops_to_remove = [] @@ -84,7 +88,7 @@ def _try_to_transform(op, block): block.remove_ops(ops_to_remove) return True - +@block_context_manager def _fuse_gelu_exact_block(block): fusion_occurred = False for op in list(block.operations): @@ -97,11 +101,10 @@ def _fuse_gelu_exact_block(block): continue if op.op_type in ["mul", "real_div"]: - with block: - fusion_occurred = _try_to_transform(op, block) - # has to break as the downstream iterator is affected. - if fusion_occurred: - return fusion_occurred + fusion_occurred = _try_to_transform(op, block) + # has to break as the downstream iterator is affected. + if fusion_occurred: + return fusion_occurred return fusion_occurred @register_pass(namespace="common") diff --git a/coremltools/converters/mil/mil/passes/graph_pass.py b/coremltools/converters/mil/mil/passes/graph_pass.py index 7c64f59e8..db3954c65 100644 --- a/coremltools/converters/mil/mil/passes/graph_pass.py +++ b/coremltools/converters/mil/mil/passes/graph_pass.py @@ -8,9 +8,6 @@ class AbstractGraphPass(ABC): - def __init__(self, minimun_deployment_target=target.iOS13): - self._minimum_deployment_target = minimun_deployment_target - def __call__(self, prog): if not prog.skip_all_passes: self.apply(prog) @@ -18,17 +15,6 @@ def __call__(self, prog): def __str__(self): return type(self).__name__ - @property - def minimun_deployment_target(self): - return self._minimum_deployment_target - - @minimun_deployment_target.setter - def minimun_deployment_target(self, t): - if not isinstance(t, target): - raise TypeError("minimun_deployment_target must be an enumeration from Enum class AvailableTarget") - self._minimum_deployment_target = t - - @abstractmethod def apply(self, prog): pass diff --git a/coremltools/converters/mil/mil/passes/helper.py b/coremltools/converters/mil/mil/passes/helper.py index 389f0e6fc..ad22a6885 100644 --- a/coremltools/converters/mil/mil/passes/helper.py +++ b/coremltools/converters/mil/mil/passes/helper.py @@ -5,7 +5,41 @@ import numpy as np -from coremltools.converters.mil.mil import Var +from coremltools.converters.mil.mil import Block, Var + +def block_context_manager(func): + """ + This decorator executes a function under the context manager `with block`. + For instance, given a function `func` with an input block and other arguments: + + def func(block, *args): + ... + with block: + op_1 = mb.add(...) + ... + with block: + op_2 = mb.relu...() + + It can be be streamlined as: + + @block_context_manager + def func(block, *args): + ... + op_1 = mb.add(...) + ... + op_2 = mb.relu...() + + Note that, the first argument of the function must have type Block. + It is highly recommended to decorate a function with block_context_manager if it is calling `with block` multiple times, + since when the code exit `block`, an expensive _propagate_nonreplaceable_vars() is invoked. + The decorator reduces the amount of calling `with block` overally. + """ + def wrapper(*args): + if not isinstance(args[0], Block): + raise ValueError("The function decorated with block_context_manager must have a Block type argument as the first input.") + with args[0]: + return func(*args) + return wrapper def _check_child_op_type(op, child_op_type): """ 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 7508f3047..2b5665f1d 100644 --- a/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py +++ b/coremltools/converters/mil/mil/passes/layernorm_instancenorm_pattern_fusion.py @@ -13,8 +13,9 @@ Operation, Var ) -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.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass DEBUG = False # set to true to plot the block before and after the transformation @@ -783,7 +784,7 @@ def _try_match_and_transform_pattern_4(reduce_op: Operation, block: Block) -> bo reduce_op, block, gamma_var, beta_var, epsilon_var, add_op, ops_to_remove ) - +@block_context_manager def _fuse_layernorm_or_instancenorm_block(block: Block): fusion_status = False for i, op in enumerate(list(block.operations)): @@ -796,20 +797,18 @@ def _fuse_layernorm_or_instancenorm_block(block: Block): # start pattern match if reduce_mean op is encountered if op.op_type == "reduce_mean": - with block: - if fusion_status is False: - fusion_status = _try_match_and_transform_pattern_1(op, block) - if fusion_status is False: - fusion_status = _try_match_and_transform_pattern_2(op, block) - if fusion_status is False: - fusion_status = _try_match_and_transform_pattern_3(op, block) + if fusion_status is False: + fusion_status = _try_match_and_transform_pattern_1(op, block) + if fusion_status is False: + fusion_status = _try_match_and_transform_pattern_2(op, block) + if fusion_status is False: + fusion_status = _try_match_and_transform_pattern_3(op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status elif op.op_type == "reduce_sum": - with block: - if fusion_status is False: - fusion_status = _try_match_and_transform_pattern_4(op, block) + if fusion_status is False: + fusion_status = _try_match_and_transform_pattern_4(op, block) # has to break as the downstream iterator is affected. if fusion_status: return fusion_status diff --git a/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py b/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py index 409d4e1c4..c1fc404a8 100644 --- a/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py +++ b/coremltools/converters/mil/mil/passes/leaky_relu_fusion.py @@ -3,11 +3,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.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 +from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import ( + _check_child_op_type, + _check_var_scalar_value_in_interval, + block_context_manager, +) +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _try_to_transform(mul_op, block): @@ -50,7 +53,7 @@ def _try_to_transform(mul_op, block): block.remove_ops(ops_to_remove) return True - +@block_context_manager def _fuse_leaky_relu_block(block): fusion_status = False for i, op in enumerate(list(block.operations)): @@ -63,8 +66,7 @@ def _fuse_leaky_relu_block(block): # start pattern match if mul op is encountered if op.op_type == "mul": - 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 diff --git a/coremltools/converters/mil/mil/passes/linear_bias_fusion.py b/coremltools/converters/mil/mil/passes/linear_bias_fusion.py index 7a8fe53d5..bc7671ce0 100644 --- a/coremltools/converters/mil/mil/passes/linear_bias_fusion.py +++ b/coremltools/converters/mil/mil/passes/linear_bias_fusion.py @@ -7,6 +7,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -54,13 +55,49 @@ def _try_to_transform(linear_op, add_or_sub_op, block): linear_kargs[k] = v x = mb.linear(**linear_kargs) + + if add_or_sub_op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=add_or_sub_op, + old_var=add_or_sub_op.outputs[0], + new_var=x, + ): + add_or_sub_op.enclosing_block.remove_ops([linear_op, add_or_sub_op]) + return True + return False + +@block_context_manager +def _fuse_linear_bias_block(block): + + def _find_candicate_op(op): + if op.op_type != "linear": + return None + # abort fusion if op output is also a block output + if op.outputs[0] in op.enclosing_block.outputs: + return None + # find add/sub op + child_ops = op.outputs[0].child_ops + if len(child_ops) == 1: + op_candidate = list(child_ops)[0] + if op_candidate.op_type in ["add", "sub"]: + return op_candidate + + fusion_occurred = False + for op in list(block.operations): + for b in op.blocks: + block_changed = True + while block_changed: + block_changed = _fuse_linear_bias_block(b) + if len(op.blocks) > 0: + # This op can't be conv or conv_transpose + continue - add_or_sub_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=add_or_sub_op, old_var=add_or_sub_op.outputs[0], new_var=x - ) - # Remove all the ops at once - block.remove_ops([linear_op, add_or_sub_op]) - return True + add_or_sub_op = _find_candicate_op(op) + if add_or_sub_op is not None: + fusion_occurred = _try_to_transform(op, add_or_sub_op, block) + # has to break as the downstream iterator is affected. + if fusion_occurred: + return fusion_occurred + return fusion_occurred @register_pass(namespace="common") class fuse_linear_bias(AbstractGraphPass): @@ -93,51 +130,9 @@ class fuse_linear_bias(AbstractGraphPass): prog: Program """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): block_changed = True while block_changed: - block_changed = self._fuse_linear_bias_block(f) - - def _fuse_linear_bias_block(self, block): - - def _find_candicate_op(op): - if op.op_type != "linear": - return None - # abort fusion if op output is also a block output - if op.outputs[0] in op.enclosing_block.outputs: - return None - # find add/sub op - child_ops = op.outputs[0].child_ops - if len(child_ops) == 1: - op_candidate = list(child_ops)[0] - if op_candidate.op_type in ["add", "sub"]: - return op_candidate - - fusion_occurred = False - for op in list(block.operations): - for b in op.blocks: - block_changed = True - while block_changed: - block_changed = self._fuse_linear_bias_block(b) - if len(op.blocks) > 0: - # This op can't be conv or conv_transpose - continue - if op in self.ops_to_skip: - continue - - add_or_sub_op = _find_candicate_op(op) - if add_or_sub_op is not None and add_or_sub_op not in self.ops_to_skip: - with block: - fusion_occurred = _try_to_transform(op, add_or_sub_op, block) - # has to break as the downstream iterator is affected. - if fusion_occurred: - return fusion_occurred - return fusion_occurred + 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 dcbe5235c..74ea8f9df 100644 --- a/coremltools/converters/mil/mil/passes/loop_invariant_elimination.py +++ b/coremltools/converters/mil/mil/passes/loop_invariant_elimination.py @@ -4,8 +4,10 @@ # 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 from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass + def _detect_loop_invariants(while_op): block = while_op.blocks[1] # body block @@ -14,8 +16,11 @@ def _detect_loop_invariants(while_op): vx_out = block.outputs[i] # first output is cond var. return_input_as_output = vx_in == vx_out # this block output is a var from outside of the block + + enclosing_block = while_op.enclosing_block + while_op_id = enclosing_block.find_op_id_in_block(while_op) output_from_outside_of_block = ( - vx_out in block._visible_vars_from_enclosing_block() + True if enclosing_block.is_var_visible_in_block(vx_out, upto_op_with_id=while_op_id) else False ) if return_input_as_output or output_from_outside_of_block: loop_invariant_ids.append(i) @@ -24,7 +29,7 @@ def _detect_loop_invariants(while_op): # need to move computation out of while loop. return loop_invariant_ids - +@block_context_manager def _loop_invariant_elimination_block(block): # Phase 1: Find vars needed to be renamed. # @@ -58,11 +63,10 @@ def _loop_invariant_elimination_block(block): for v_src, v_tgt, op in output_rename: if v_tgt in block.outputs: # rename the loop output to existing block output names - with block: - res = mb.identity(x=v_src, before_op=op, name=v_tgt.name) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=v_tgt, new_var=res - ) + res = mb.identity(x=v_src, before_op=op, name=v_tgt.name) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=v_tgt, new_var=res + ) # Phase 3: Perform loop invariant elimination without fear! for op in list(block.operations): 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 1470bd6c2..6930444d3 100644 --- a/coremltools/converters/mil/mil/passes/matmul_weight_bias_fusion.py +++ b/coremltools/converters/mil/mil/passes/matmul_weight_bias_fusion.py @@ -7,6 +7,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -55,6 +56,10 @@ def _try_to_transform(matmul_op, add_op, block): weight, linear_x = matmul_op.y, matmul_op.x transpose_weight = matmul_op.transpose_y.val transpose_x = matmul_op.transpose_x.val + + # We potentially are going to transpose the weight, so if the weight itself is not removable, we skip this path + if len(weight.nonreplaceable_vars_upstream) > 0: + return False if linear_x.rank < 2 or weight.rank != 2: # We don't support these cases yet. @@ -114,12 +119,35 @@ def _try_to_transform(matmul_op, add_op, block): name=out_name, ) - add_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=add_op, old_var=add_op.outputs[0], new_var=x - ) - # Remove all the ops at once - block.remove_ops([matmul_op, add_op]) - return True + if add_op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=add_op, + old_var=add_op.outputs[0], + new_var=x, + ): + add_op.enclosing_block.remove_ops([matmul_op, add_op]) + return True + return False + +@block_context_manager +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) + if len(op.blocks) > 0: + # This op can't be matmul + continue + + add_op = _find_candidate_op(op) + + if add_op is not None: + 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") class fuse_matmul_weight_bias(AbstractGraphPass): @@ -140,41 +168,9 @@ class fuse_matmul_weight_bias(AbstractGraphPass): prog: Program """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): block_changed = True while block_changed: - block_changed = self._fuse_matmul_weight_bias_block(f) - - def _fuse_matmul_weight_bias_block(self, block): - fusion_status = False - for op in list(block.operations): - for b in op.blocks: - block_changed = True - while block_changed: - block_changed = self._fuse_matmul_weight_bias_block(b) - if len(op.blocks) > 0: - # This op can't be matmul - continue - - if op in self.ops_to_skip: - continue - - add_op = _find_candidate_op(op) - if add_op in self.ops_to_skip: - continue - - if add_op is not None: - with 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 + 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 3c3320e7f..b5a1da987 100644 --- a/coremltools/converters/mil/mil/passes/merge_consecutive_paddings.py +++ b/coremltools/converters/mil/mil/passes/merge_consecutive_paddings.py @@ -6,9 +6,12 @@ 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 -from .helper import _check_child_op_type +from coremltools.converters.mil.mil.passes.helper import ( + _check_child_op_type, + block_context_manager, +) +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _match_pattern(block, padding_op): @@ -47,20 +50,18 @@ def _match_pattern(block, padding_op): 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, + 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 + ) - 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, - 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 - ) - - block.remove_ops([padding_op, child_padding_op]) + block.remove_ops([padding_op, child_padding_op]) return True - +@block_context_manager def _merge_padding_block(block): for op in list(block.operations): result = _match_pattern(block, op) diff --git a/coremltools/converters/mil/mil/passes/noop_elimination.py b/coremltools/converters/mil/mil/passes/noop_elimination.py index 5029b9020..9249e02dd 100644 --- a/coremltools/converters/mil/mil/passes/noop_elimination.py +++ b/coremltools/converters/mil/mil/passes/noop_elimination.py @@ -5,8 +5,9 @@ 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.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _remove_elementwise_binary(op, block, x, y): # We remove the ops that has op.x == x or op.y == y @@ -25,13 +26,15 @@ def _remove_elementwise_binary(op, block, x, y): # We might be using elementwise as broadcasting if input_shape != output_shape: return False - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=input_op, old_var=op.outputs[0], new_var=input_var - ) - block.remove_ops([op]) - - return True + + if op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=input_op, + old_var=op.outputs[0], + new_var=input_var, + ): + op.enclosing_block.remove_ops([op]) + return True + return False def remove_elementwise(op, block): @@ -46,38 +49,65 @@ def remove_elementwise(op, block): else: return False -def remove_same_shape(op, block): +def remove_slice_by_index(op ,block): input_shape = op.x.sym_type output_shape = op.outputs[0].sym_type if input_shape != output_shape: return False + if op.stride is not None and op.stride.val is not None: + stride = op.stride.val.flatten().tolist() + if any([x < 0 for x in stride]): + return False + input_var = op.x input_op = input_var.op - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=input_op, old_var=op.outputs[0], new_var=input_var - ) + if op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=input_op, + old_var=op.outputs[0], + new_var=input_var, + ): + op.enclosing_block.remove_ops([op]) + return True + return False - # Remove all the ops at once - block.remove_ops([op]) - return True -def remove_linear(op, block): - if op.alpha.val != 1 or op.beta.val != 0: +def remove_same_shape(op, block): + input_shape = op.x.sym_type + output_shape = op.outputs[0].sym_type + + if input_shape != output_shape: return False input_var = op.x input_op = input_var.op - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=input_op, old_var=op.outputs[0], new_var=input_var - ) + if op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=input_op, + old_var=op.outputs[0], + new_var=input_var, + ): + op.enclosing_block.remove_ops([op]) + return True + return False - # Remove all the ops at once - block.remove_ops([op]) - return True +def remove_linear(op, block): + if op.alpha.val != 1 or op.beta.val != 0: + return False + + input_var = op.x + input_op = input_var.op + + if op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=input_op, + old_var=op.outputs[0], + new_var=input_var, + ): + op.enclosing_block.remove_ops([op]) + return True + return False def remove_transpose(op, block): perm = np.array([p if p >= 0 else p+len(op.perm.val) for p in op.perm.val]) @@ -87,14 +117,16 @@ def remove_transpose(op, block): input_var = op.x input_op = input_var.op - - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=input_op, old_var=op.outputs[0], new_var=input_var - ) - - # Remove all the ops at once - block.remove_ops([op]) - return True + + if op.enclosing_block.try_replace_uses_of_var_after_op( + anchor_op=input_op, + old_var=op.outputs[0], + new_var=input_var, + ): + op.enclosing_block.remove_ops([op]) + return True + return False + _SUPPORTED_OPS = { "add", "mul", @@ -125,7 +157,7 @@ def remove_transpose(op, block): "sub": remove_elementwise, "reshape": remove_same_shape, "split": remove_same_shape, - "slice_by_index": remove_same_shape, + "slice_by_index": remove_slice_by_index, "slice_by_size": remove_same_shape, "pad": remove_same_shape, "tile": remove_same_shape, @@ -149,6 +181,31 @@ def _match_pattern(op): return op_to_removal_fn[op.op_type] return None + +@block_context_manager +def _noop_elimination_block_wrapper(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) + if len(op.blocks) > 0: + continue + + remove_fn = _match_pattern(op) + if remove_fn is not None: + status = remove_fn(op, block) + # has to break as the downstream iterator is affected. + if status: + return status + return False + + block_changed = True + while block_changed: + block_changed = _noop_elimination_block(block) + @register_pass(namespace="common") class noop_elimination(AbstractGraphPass): @@ -168,33 +225,7 @@ class noop_elimination(AbstractGraphPass): ... """ - def __init__(self): - self.ops_to_skip = set() - - def set_ops_to_skip(self, prog): - pass def apply(self, prog): - self.set_ops_to_skip(prog) for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = self._noop_elimination_block(f) - - def _noop_elimination_block(self, block): - for op in list(block.operations): - for b in op.blocks: - block_changed = True - while block_changed: - block_changed = self._noop_elimination_block(b) - if len(op.blocks) > 0: - continue - - remove_fn = _match_pattern(op) - if remove_fn is not None and op not in self.ops_to_skip: - with block: - status = remove_fn(op, block) - # has to break as the downstream iterator is affected. - if status: - return status - return False \ No newline at end of file + _noop_elimination_block_wrapper(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 bead0e468..1a2fe82a7 100644 --- a/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.py +++ b/coremltools/converters/mil/mil/passes/onehot_matmul_to_gather.py @@ -5,8 +5,9 @@ from .helper import _check_child_op_type, _check_var_scalar_value 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.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _try_to_transform(onehot_op, block): @@ -60,7 +61,7 @@ def _try_to_transform(onehot_op, block): block.remove_ops([onehot_op, matmul_op]) return True - +@block_context_manager def _fuse_onehot_matmul_to_gather_block(block): fusion_status = False for i, op in enumerate(list(block.operations)): @@ -74,8 +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 diff --git a/coremltools/converters/mil/mil/passes/pad_conv_connect.py b/coremltools/converters/mil/mil/passes/pad_conv_connect.py index 867e3e044..f2c8fd298 100644 --- a/coremltools/converters/mil/mil/passes/pad_conv_connect.py +++ b/coremltools/converters/mil/mil/passes/pad_conv_connect.py @@ -8,6 +8,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -82,7 +83,7 @@ def _compute_new_pad_values(transpose_op): return True - +@block_context_manager def _pad_conv_connect_block(block): fusion_status = False for op in list(block.operations): @@ -96,8 +97,7 @@ def _pad_conv_connect_block(block): 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 diff --git a/coremltools/converters/mil/mil/passes/quantization_passes.py b/coremltools/converters/mil/mil/passes/quantization_passes.py index 057a16322..ab12faa1e 100644 --- a/coremltools/converters/mil/mil/passes/quantization_passes.py +++ b/coremltools/converters/mil/mil/passes/quantization_passes.py @@ -7,8 +7,12 @@ import numpy as np -from coremltools.converters.mil.mil import Builder as mb +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.helper import block_context_manager from coremltools.converters.mil.mil.program import Program @@ -61,7 +65,8 @@ def apply(self, prog): raise TypeError( 'Transform "{}" can only be applied on PyMIL programs.'.format(self) ) - + + @block_context_manager def apply_block(block): for op in list(block.operations): for b in op.blocks: @@ -149,19 +154,11 @@ def is_valid_op(self, op): return True def is_valid_parameter(self, op, param_name): - - if op.op_type in ["resample"] and param_name == "coordinates": - return False - - if op.op_type in ["crop_resize"] and param_name == "spatial_scale": - return False - - if op.op_type in ["upsample_nearest_neighbor"] and param_name in ["scale_factor_height", "scale_factor_width"]: - return False - - if op.op_type in ["upsample_bilinear"] and param_name in ["scale_factor_height", "scale_factor_width"]: - return False - + type_domain = getattr(op.input_spec.input_types[param_name], "type_domain", None) + if type_domain is not None: + if len(type_domain) == 0: + return True + return types.fp16 in type_domain return True def _check_underflow_to_zero(self, new_var, var): @@ -207,19 +204,19 @@ def transform_op(self, op): continue inputs_modified = True - with block: - casted_var_name = var.name + "_to_fp16" - if len(var._child_ops) > 1 and casted_var_name in self.cache_vars and (self.cache_vars[casted_var_name] in block._visible_vars_in_block()[1]): - casted_inputs[param][i] = self.cache_vars[casted_var_name] - else: - x = mb.cast( - x=var, dtype="fp16", name=casted_var_name, before_op= op - ) - self._check_underflow_to_zero(x, var) - - casted_inputs[param][i] = x - if len(var._child_ops) > 1: - self.cache_vars[casted_var_name] = casted_inputs[param][i] + casted_var_name = var.name + "_to_fp16" + if len(var._child_ops) > 1 and casted_var_name in self.cache_vars and ( + block.is_var_visible_in_block(self.cache_vars[casted_var_name])): + casted_inputs[param][i] = self.cache_vars[casted_var_name] + else: + x = mb.cast( + x=var, dtype="fp16", name=casted_var_name, before_op= op + ) + self._check_underflow_to_zero(x, var) + + casted_inputs[param][i] = x + if len(var._child_ops) > 1: + self.cache_vars[casted_var_name] = casted_inputs[param][i] if not is_list_input: casted_inputs[param] = casted_inputs[param][0] @@ -230,8 +227,7 @@ def transform_op(self, op): ) casted_inputs["name"] = op.name + "_cast" casted_inputs["before_op"] = op - with block: - quant_output = getattr(mb, op.op_type)(**casted_inputs) + quant_output = getattr(mb, op.op_type)(**casted_inputs) if not isinstance(quant_output, (list, tuple)): quant_output = [quant_output] @@ -240,19 +236,24 @@ def transform_op(self, op): if old_output_var.is_tensor_or_scalar_of(dtype="fp32") and ( not new_output_var.is_tensor_or_scalar_of(dtype="fp32") ): - with block: - x = mb.cast( - x=new_output_var, - dtype="fp32", - name=new_output_var.name + "_to_fp32", - before_op=op, - ) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=old_output_var, new_var=x - ) + x = mb.cast( + x=new_output_var, + dtype="fp32", + name=new_output_var.name + "_to_fp32", + before_op=op, + ) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, + old_var=old_output_var, + new_var=x, + force_replace=True, + ) else: op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=old_output_var, new_var=new_output_var + anchor_op=op, + old_var=old_output_var, + new_var=new_output_var, + force_replace=True, ) block.remove_ops([op]) 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 878d77d68..19a00b424 100644 --- a/coremltools/converters/mil/mil/passes/rank0_expand_dims_swap.py +++ b/coremltools/converters/mil/mil/passes/rank0_expand_dims_swap.py @@ -5,6 +5,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -80,7 +81,7 @@ def _try_to_transform(op, block): block.remove_ops(ops_to_remove) return True - +@block_context_manager def _rank0_expand_dims_swap(block): fusion_occurred = False for op in list(block.operations): @@ -93,11 +94,10 @@ def _rank0_expand_dims_swap(block): continue if op.op_type in ["add", "sub", "mul", "real_div", "floor_div"]: - with block: - fusion_occurred = _try_to_transform(op, block) - # has to break as the downstream iterator is affected. - if fusion_occurred: - return fusion_occurred + fusion_occurred = _try_to_transform(op, block) + # has to break as the downstream iterator is affected. + if fusion_occurred: + return fusion_occurred return fusion_occurred @register_pass(namespace="common") diff --git a/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py b/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py index a7b8d4a76..684ae1ea3 100644 --- a/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py +++ b/coremltools/converters/mil/mil/passes/reduce_mean_fusion.py @@ -3,11 +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 -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.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import ( + _check_child_op_type, + _check_var_scalar_value, + block_context_manager, +) +from coremltools.converters.mil.mil.passes.pass_registry import register_pass from coremltools.converters.mil.mil.types.symbolic import is_symbolic -from .helper import _check_var_scalar_value, _check_child_op_type def _try_to_transform(reduce_sum_op, block): @@ -68,7 +72,7 @@ def _try_to_transform(reduce_sum_op, block): block.remove_ops(ops_to_remove) return True - +@block_context_manager def _fuse_reduce_mean_block(block): fusion_status = False for i, op in enumerate(list(block.operations)): @@ -81,8 +85,7 @@ def _fuse_reduce_mean_block(block): # start pattern match if mul op is encountered if op.op_type == "reduce_sum": - 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 diff --git a/coremltools/converters/mil/mil/passes/reduce_transposes.py b/coremltools/converters/mil/mil/passes/reduce_transposes.py index 6e60c8dcd..9ad3b7d6a 100644 --- a/coremltools/converters/mil/mil/passes/reduce_transposes.py +++ b/coremltools/converters/mil/mil/passes/reduce_transposes.py @@ -764,14 +764,13 @@ def _add_output_sinks(self): new_outputs = [] output_sinks_var = {} for out_var in self.block.outputs: - with self.block: - if out_var not in output_sinks_var: - out_sink = mb.identity(x=out_var) - output_sinks_var[out_var] = out_sink - else: - out_sink = output_sinks_var[out_var] - new_outputs.append(out_sink) - self.output_sink_ops.append(out_sink.op) + if out_var not in output_sinks_var: + out_sink = mb.identity(x=out_var) + output_sinks_var[out_var] = out_sink + else: + out_sink = output_sinks_var[out_var] + new_outputs.append(out_sink) + self.output_sink_ops.append(out_sink.op) self.block.set_outputs(new_outputs) def _visit_unary_like_op(self, op, input_var=None): @@ -1103,17 +1102,15 @@ def _remove_transpose_ops(self, starting_transpose_op): # If the same input_var is in output twice, we can't rename it twice, therefore we initiate an # Identity op to match the name if input_var in self.block.inputs.values(): - with self.block: - input_var = mb.identity(x=input_var, before_op=op, name=output_var.name) - parent_op = None # set anchor op as None. + input_var = mb.identity(x=input_var, before_op=op, name=output_var.name) + parent_op = None # set anchor op as None. elif input_var not in name_changed_vars: input_var.name = output_var.name input_var.op.name = output_var.op.name name_changed_vars.update([input_var]) else: - with self.block: - input_var = mb.identity(x=input_var, before_op=op, name=output_var.name) - parent_op = input_var.op + input_var = mb.identity(x=input_var, before_op=op, name=output_var.name) + parent_op = input_var.op # connect all the child ops of the output_var to the parent of the transpose op. self.block.replace_uses_of_var_after_op( @@ -1144,28 +1141,27 @@ def _remove_transpose_ops(self, starting_transpose_op): continue self.materialized_ops_handled.add((op, input_var)) - with self.block: - if input_var == starting_transpose_op_out_var: - # materialize op is connected to the starting transpose op - # in this case, connect to its parent - if op in self.output_sink_ops: - continue - i1 = starting_transpose_op_input_var - else: - i1 = input_var - + if input_var == starting_transpose_op_out_var: + # materialize op is connected to the starting transpose op + # in this case, connect to its parent if op in self.output_sink_ops: - # The input_var of output sink is itself a output. We can safely - # modify the name of the input_var since it should only be consumed - # by block output here. - if i1 not in name_changed_vars: - x = mb.transpose(x=i1, perm=perm, before_op=op, name=i1.name) - i1.name = '_before_transpose_op_' + x.op.name - i1.op.name = '_before_transpose_op_' + x.op.name - else: - x = mb.transpose(x=i1, perm=perm, before_op=op, name=self.old_output_vars[i1]) + continue + i1 = starting_transpose_op_input_var + else: + i1 = input_var + + if op in self.output_sink_ops: + # The input_var of output sink is itself a output. We can safely + # modify the name of the input_var since it should only be consumed + # by block output here. + if i1 not in name_changed_vars: + x = mb.transpose(x=i1, perm=perm, before_op=op, name=i1.name) + i1.name = '_before_transpose_op_' + x.op.name + i1.op.name = '_before_transpose_op_' + x.op.name else: - x = mb.transpose(x=i1, perm=perm, before_op=op) + x = mb.transpose(x=i1, perm=perm, before_op=op, name=self.old_output_vars[i1]) + else: + x = mb.transpose(x=i1, perm=perm, before_op=op) self.block.replace_uses_of_var_after_op( anchor_op=x.op, @@ -1241,10 +1237,11 @@ def _reduce_transposes_block(block): for op in list(block.operations): if len(op.blocks) > 0: return - - opt_transposes = TransposeOptimization(block) - opt_transposes.block_traversal() - opt_transposes.apply_transform() + + with block: + opt_transposes = TransposeOptimization(block) + opt_transposes.block_traversal() + opt_transposes.apply_transform() @register_pass(namespace="common") diff --git a/coremltools/converters/mil/mil/passes/remove_redundant_ops.py b/coremltools/converters/mil/mil/passes/remove_redundant_ops.py index 2bada8cf9..17e88491c 100644 --- a/coremltools/converters/mil/mil/passes/remove_redundant_ops.py +++ b/coremltools/converters/mil/mil/passes/remove_redundant_ops.py @@ -7,6 +7,7 @@ from .helper import _are_ops_identical from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -106,42 +107,46 @@ def _try_to_transform(parent_var): block_changed = True return block_changed +@block_context_manager +def _remove_redundant_ops_in_block_wrapper(block): -def _remove_redundant_ops_in_block(block): + 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.") + 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: + # iterate over the block inputs + for input_var in block_input_var_list: + if len(input_var.child_ops) > 1: _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 + # 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) + 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: + if len(op.outputs) > 0 and len(op.outputs[0].child_ops) > 1: # 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 + # has to break as the downstream iterator is affected. + if graph_updated: + return graph_updated + return graph_updated + + block_changed = True + while block_changed: + block_changed = _remove_redundant_ops_in_block(block) @register_pass(namespace="common") @@ -183,6 +188,4 @@ class remove_redundant_ops(AbstractGraphPass): """ def apply(self, prog): for f in prog.functions.values(): - block_changed = True - while block_changed: - block_changed = _remove_redundant_ops_in_block(f) + _remove_redundant_ops_in_block_wrapper(f) diff --git a/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py b/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py index c41605d3d..0be499072 100644 --- a/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py +++ b/coremltools/converters/mil/mil/passes/remove_symbolic_reshape.py @@ -11,10 +11,12 @@ num_symbolic, ) 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.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass +@block_context_manager def _remove_symbolic_reshape_block(block): num_changes = 0 for op in list(block.operations): @@ -44,18 +46,17 @@ def _remove_symbolic_reshape_block(block): continue # Convert the one symbol to -1 integer_shape = [-1 if is_symbolic(i) else i for i in shape] - with block: - shape_const = mb.const( - val=integer_shape, - name=op.shape.name + "x", - before_op=op, - ) - reshaped = mb.reshape(x=op.x, shape=shape_const, name=op.name, before_op=op) - op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=op, old_var=op.outputs[0], new_var=reshaped - ) - # Remove all the ops at once - block.remove_ops([op, op.shape.op]) + shape_const = mb.const( + val=integer_shape, + name=op.shape.name + "x", + before_op=op, + ) + reshaped = mb.reshape(x=op.x, shape=shape_const, name=op.name, before_op=op) + op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=op, old_var=op.outputs[0], new_var=reshaped + ) + # Remove all the ops at once + block.remove_ops([op, op.shape.op]) num_changes += 1 return 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 3c78227ad..c257c5861 100644 --- a/coremltools/converters/mil/mil/passes/replace_stack_reshape.py +++ b/coremltools/converters/mil/mil/passes/replace_stack_reshape.py @@ -4,9 +4,9 @@ # 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 from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass - +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass def _match_operation(stack_op): @@ -72,15 +72,14 @@ def _replace_stack_reshape_ops(block, stack_op, reshape_op): interleave = (concat_axis == stack_axis_val - 1) - with block: - x = mb.concat(values=stack_op.values, axis=concat_axis, before_op=stack_op, interleave=interleave) - - reshape_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=stack_op, old_var=reshape_op.outputs[0], new_var=x - ) - block.remove_ops([stack_op, reshape_op]) + x = mb.concat(values=stack_op.values, axis=concat_axis, before_op=stack_op, interleave=interleave) + reshape_op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=stack_op, old_var=reshape_op.outputs[0], new_var=x + ) + block.remove_ops([stack_op, reshape_op]) +@block_context_manager def _replace_stack_reshape_block(block): for op in list(block.operations): diff --git a/coremltools/converters/mil/mil/passes/test_compression_passes.py b/coremltools/converters/mil/mil/passes/test_compression_passes.py new file mode 100644 index 000000000..27ec586e0 --- /dev/null +++ b/coremltools/converters/mil/mil/passes/test_compression_passes.py @@ -0,0 +1,78 @@ +# Copyright (c) 2022, 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 itertools + +import numpy as np +import pytest + +import coremltools as ct +from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.testing_utils import get_op_types_in_program + +from .compression_passes import ( + WeightSparsifier, + WeightPalettizer, + WeightAffineQuantizer, +) + +np.random.seed(1984) + +def _get_conv_program(): + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 30, 10, 10))], opset_version=ct.target.iOS16) + def prog(x): + conv_weight = np.random.rand(90, 30, 2, 2).astype(np.float32) + x = mb.conv(x=x, weight=conv_weight) + return x + return prog + +class TestBasicCompressionGraphPass: + # Most of the numerical tests are already convered in coremltools.tests.ml_program.test_compression_utils + # This test is checking the basic behavior of the graph pass classes + + @staticmethod + @pytest.mark.parametrize( + "fake_compression", + [True, False], + ) + def test_affine_quantizer(fake_compression): + quantizer = WeightAffineQuantizer(fake_compression=fake_compression, op_selector=lambda const: True) + prog = _get_conv_program() + quantizer.apply(prog) + expected_ops = ["constexpr_affine_dequantize", "conv"] if not fake_compression else ["conv"] + assert get_op_types_in_program(prog) == expected_ops + + @staticmethod + @pytest.mark.parametrize( + "fake_compression", + [True, False], + ) + def test_weight_sparsifier(fake_compression): + quantizer = WeightSparsifier( + fake_compression=fake_compression, + op_selector=lambda const: True, + mode="percentile_based", + target_percentile=0.75) + prog = _get_conv_program() + quantizer.apply(prog) + expected_ops = ["constexpr_sparse_to_dense", "conv"] if not fake_compression else ["conv"] + assert get_op_types_in_program(prog) == expected_ops + + @staticmethod + @pytest.mark.parametrize( + "fake_compression", + [True, False], + ) + def test_weight_palettization(fake_compression): + quantizer = WeightPalettizer( + fake_compression=fake_compression, + op_selector=lambda const: True, + mode="uniform", + nbits=4, + ) + prog = _get_conv_program() + quantizer.apply(prog) + expected_ops = ["constexpr_lut_to_dense", "conv"] if not fake_compression else ["conv"] + assert get_op_types_in_program(prog) == expected_ops diff --git a/coremltools/converters/mil/mil/passes/test_noop_elimination.py b/coremltools/converters/mil/mil/passes/test_noop_elimination.py index 151958649..01988c076 100644 --- a/coremltools/converters/mil/mil/passes/test_noop_elimination.py +++ b/coremltools/converters/mil/mil/passes/test_noop_elimination.py @@ -200,6 +200,30 @@ def prog(x): expected_output_shapes={block.outputs[0].name: (2, 4)}, ) +def test_slicebyindex_negative_stride(): + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4))]) + def prog(x): + r1 = mb.slice_by_index( + x=x, + begin=[0, 0], + end=[0, 0], + stride=[1, -1], + begin_mask=[True, True], + end_mask=[True, True] + ) + return mb.relu(x=r1) + + prev_prog, prev_block, block = apply_pass_and_basic_check( + prog, "common::noop_elimination" + ) + assert get_op_types_in_program(prev_prog) == ["slice_by_index", "relu"] + assert get_op_types_in_program(prog) == ["slice_by_index", "relu"] + assert_model_is_valid( + prog, + {"x": (2, 4)}, + expected_output_shapes={block.outputs[0].name: (2, 4)}, + ) + @pytest.mark.parametrize("begin_mask, end_mask", itertools.product(itertools.product([True, False],[True, False]), diff --git a/coremltools/converters/mil/mil/passes/test_passes.py b/coremltools/converters/mil/mil/passes/test_passes.py index 3ba008ae6..07bc10811 100644 --- a/coremltools/converters/mil/mil/passes/test_passes.py +++ b/coremltools/converters/mil/mil/passes/test_passes.py @@ -29,6 +29,12 @@ from coremltools.converters.mil.mil.passes.helper import _check_var_scalar_value from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY +from .compression_passes import ( + WeightSparsifier, + WeightPalettizer, + WeightAffineQuantizer, +) + np.random.seed(1984) validate_model = True @@ -204,8 +210,8 @@ def prog(a, b): assert len(while_op.blocks[0].inputs) == 2 assert len(while_op.outputs) == 2 assert len(while_op.loop_vars) == 2 - assert while_op.blocks[0].inputs[0].name == "a_x1" - assert while_op.blocks[0].inputs[1].name == "b_x1" + assert while_op.blocks[0].inputs[0].name == "a_x0" + assert while_op.blocks[0].inputs[1].name == "b_x0" prev_prog = copy.deepcopy(prog) PASS_REGISTRY["common::loop_invariant_elimination"](prog) @@ -215,7 +221,7 @@ def prog(a, b): assert len(while_op.blocks[0].inputs) == 1 assert len(while_op.outputs) == 1 assert len(while_op.loop_vars) == 1 - assert while_op.blocks[0].inputs[0].name == "a_x1" + assert while_op.blocks[0].inputs[0].name == "a_x0" if validate_model: assert_model_is_valid(prog, {"a": (1, 2), "b": (1, 2)}) @@ -245,8 +251,8 @@ def cond(a, bx): assert len(while_op.blocks[0].inputs) == 2 assert len(while_op.outputs) == 2 assert len(while_op.loop_vars) == 2 - assert while_op.blocks[0].inputs[0].name == "a_x1" - assert while_op.blocks[0].inputs[1].name == "b_x1" + assert while_op.blocks[0].inputs[0].name == "a_x0" + assert while_op.blocks[0].inputs[1].name == "b_x0" prev_prog = copy.deepcopy(prog) PASS_REGISTRY["common::loop_invariant_elimination"](prog) @@ -256,7 +262,7 @@ def cond(a, bx): assert len(while_op.blocks[0].inputs) == 1 assert len(while_op.outputs) == 1 assert len(while_op.loop_vars) == 1 - assert while_op.blocks[0].inputs[0].name == "a_x1" + assert while_op.blocks[0].inputs[0].name == "a_x0" if validate_model: assert_model_is_valid(prog, {"a": (1, 2), "b": (1, 2)}) @@ -1694,7 +1700,6 @@ def prog(x): expected_output_shapes={block.outputs[0].name: (B, H, W, C)}, ) - class TestUpdateOutputDtypePass: def test_single_output(self): @@ -1770,4 +1775,295 @@ def prog(input): assert block.outputs[1].name == "split_1" +def _get_constexpr_cast(shape): + val = np.random.rand(*shape).astype(np.float16) + return mb.constexpr_cast(source_val=val, output_dtype="fp32") + +def _get_constexpr_sparse_to_dense(shape): + val = np.random.rand(*shape) + sparse_params = WeightSparsifier.compress(val=val, mode="PERCENTILE_MODE", target_percentile=0.4) + return mb.constexpr_sparse_to_dense( + nonzero_data=sparse_params.nonzero_data, + mask=sparse_params.mask, + shape=np.uint32(sparse_params.shape),) + +def _get_constexpr_lut_to_dense(shape): + val = np.random.rand(*shape) + lut_params = WeightPalettizer.compress(val=val, nbits=4, mode="UNIFORM") + return mb.constexpr_lut_to_dense( + indices=lut_params.indices, + lut=lut_params.lut, + shape=np.uint32(lut_params.shape),) + +def _get_constexpr_affine_dequantize(shape): + val = np.random.rand(*shape) + quant_params = WeightAffineQuantizer.compress(val=val, axis=0, mode="LINEAR_SYMMETRIC") + return mb.constexpr_affine_dequantize( + quantized_data=quant_params.quantized_data, + zero_point=quant_params.zero_point, + scale=quant_params.scale, + axis=quant_params.axis,) + +CONSTEXPR_FUNCS = { + "constexpr_cast": _get_constexpr_cast, + "constexpr_sparse_to_dense": _get_constexpr_sparse_to_dense, + "constexpr_lut_to_dense": _get_constexpr_lut_to_dense, + "constexpr_affine_dequantize": _get_constexpr_affine_dequantize, +} + +CONSTEXPR_OPS = [ + "constexpr_cast", + "constexpr_sparse_to_dense", + "constexpr_lut_to_dense", + "constexpr_affine_dequantize" +] + +class TestSkipConstexprOps: + + @staticmethod + @pytest.mark.parametrize( + "constexpr_op", + CONSTEXPR_OPS, + ) + def test_skip_const_elimination(constexpr_op): + """ + constexpr_op + | + v + const -> linear + | + v + input --------------> add -> output + + We are testing that: + 1. constexpr_op can serve as a const input weight for linear op + 2. linear op shoudn't be removed by the const_elimination pass + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(4,))]) + def prog(x): + a = np.random.rand(2,) + constexpr = CONSTEXPR_FUNCS[constexpr_op]((4, 2)) + linear = mb.linear(x=a, weight=constexpr) + return mb.add(x=x, y=linear) + + PASS_REGISTRY["common::const_elimination"](prog) + assert get_op_types_in_program(prog) == [constexpr_op, "linear", "add"] + + @staticmethod + @pytest.mark.parametrize( + "constexpr_op, weight_constexpr, bias_constexpr", + itertools.product( + CONSTEXPR_OPS, + [True, False], + [True, False], + ) + ) + def test_skip_fuse_matmul_weight_bias(constexpr_op, weight_constexpr, bias_constexpr): + """ + const_1 const_2 + | | + v v + input -----> matmul -----> add ---> out + + In this case, if either const_1 or const_2 is constexpr op, they should be not fused into a single linear op + """ + + def get_matmul(x, weight_constexpr): + weight = CONSTEXPR_FUNCS[constexpr_op]((3, 2)) + if not weight_constexpr: + weight = weight.val + return mb.matmul(x=x, y=weight) + + def get_add(x, bias_constexpr): + bias = CONSTEXPR_FUNCS[constexpr_op]((2,)) + if not bias_constexpr: + bias = bias.val + return mb.add(x=x, y=bias) + + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 3))]) + def prog(x): + x = get_matmul(x, weight_constexpr) + x = get_add(x, bias_constexpr) + return x + + apply_pass_and_basic_check(prog, "common::fuse_matmul_weight_bias") + apply_pass_and_basic_check(prog, "common::const_elimination") + apply_pass_and_basic_check(prog, "common::dead_code_elimination") + + if not weight_constexpr and not bias_constexpr: + expected_ops = ["linear"] + else: + expected_ops = [] + if weight_constexpr: + expected_ops.append(constexpr_op) + expected_ops.append("matmul") + if bias_constexpr: + expected_ops.append(constexpr_op) + expected_ops.append("add") + + assert get_op_types_in_program(prog) == expected_ops + + @staticmethod + @pytest.mark.parametrize( + "constexpr_op, op, weight_constexpr, const_constexpr", + itertools.product( + CONSTEXPR_OPS, + ["mul", "add"], + [True, False], + [True, False], + ) + ) + def test_skip_fuse_conv(constexpr_op, op, weight_constexpr, const_constexpr): + + """ + const_1 const_2 + | | + v v + input -----> conv -----> mul/add ---> out + + This pattern shouldn't be fused into a single conv layer if one of const_1 or const_2 is a constexpr op. + """ + Cin, Cout = 3, 3 + input_shape = (2, Cin, 5, 5) + @mb.program(input_specs=[mb.TensorSpec(shape=input_shape)]) + def prog(x): + conv_weight = CONSTEXPR_FUNCS[constexpr_op]((Cout, Cin, 2, 2)) + if not weight_constexpr: + conv_weight = conv_weight.val + x = mb.conv(x=x, weight=conv_weight) + const = CONSTEXPR_FUNCS[constexpr_op]((Cout, 1, 1)) + if not const_constexpr: + const = const.val + return getattr(mb, op)(x=x, y=const) + + apply_pass_and_basic_check(prog, "common::fuse_conv_scale") + apply_pass_and_basic_check(prog, "common::fuse_conv_bias") + apply_pass_and_basic_check(prog, "common::const_elimination") + apply_pass_and_basic_check(prog, "common::dead_code_elimination") + + expected_ops = [] + if not weight_constexpr and not const_constexpr: + expected_ops = ["conv"] + else: + if weight_constexpr: + expected_ops.append(constexpr_op) + expected_ops.append("conv") + if const_constexpr: + expected_ops.append(constexpr_op) + if op != "add" or const_constexpr: + expected_ops.append(op) + + assert get_op_types_in_program(prog) == expected_ops + + @staticmethod + @pytest.mark.parametrize( + "constexpr_op, weight_constexpr, bias_constexpr", + itertools.product( + CONSTEXPR_OPS, + [True, False], + [True, False], + ) + ) + def test_skip_fuse_linear_bias(constexpr_op, weight_constexpr, bias_constexpr): + """ + const_1 const_2 + | | + v V + input -----> linear -----> add ---> out + + This pattern shouldn't be fused into a single linear layer if one of const_1 or const_2 is a constexpr op. + """ + @mb.program(input_specs=[mb.TensorSpec(shape=(2,))]) + def prog(x): + weight = CONSTEXPR_FUNCS[constexpr_op]((4, 2)) + if not weight_constexpr: + weight = weight.val + linear = mb.linear(x=x, weight=weight) + bias = CONSTEXPR_FUNCS[constexpr_op]((4,)) + if not bias_constexpr: + bias = bias.val + return mb.add(x=linear, y=bias) + + apply_pass_and_basic_check(prog, "common::fuse_linear_bias") + apply_pass_and_basic_check(prog, "common::const_elimination") + apply_pass_and_basic_check(prog, "common::dead_code_elimination") + + expected_ops = [] + if not weight_constexpr and not bias_constexpr: + expected_ops = ["linear"] + else: + if weight_constexpr: + expected_ops.append(constexpr_op) + expected_ops.append("linear") + if bias_constexpr: + expected_ops.append(constexpr_op) + expected_ops.append("add") + + assert get_op_types_in_program(prog) == expected_ops + + + @staticmethod + @pytest.mark.parametrize( + "constexpr_op, weight_constexpr, bias_constexpr", + itertools.product( + CONSTEXPR_OPS, + [True, False], + [True, False], + ) + ) + def test_skip_fuse_conv_batchnorm(constexpr_op, weight_constexpr, bias_constexpr): + """ + weight bias + | | + |_____ ____| + | | + v v + input -----> conv -----> batch_norm ---> out + + This pattern shouldn't be fused into a single conv layer if one of the weight / bias is a constexpr op. + """ + Cin, Cout = 2, 3 + input_shape = (2, Cin, 5, 5) + + @mb.program(input_specs=[mb.TensorSpec(shape=input_shape)]) + def prog(x): + # conv layer + weight = CONSTEXPR_FUNCS[constexpr_op]((Cout, Cin, 2, 2)) + if not weight_constexpr: + weight = weight.val + bias = CONSTEXPR_FUNCS[constexpr_op]((Cout, )) + if not bias_constexpr: + bias = bias.val + + x = mb.conv( + x=x, + weight=weight, + bias=bias, + ) + + # batch_norm layer + gamma = np.random.rand(Cout) + beta = np.random.rand(Cout) + mean = np.random.rand(Cout) + variance = np.random.rand(Cout) + epsilon = 1e-2 + return mb.batch_norm( + x=x, + mean=mean, + variance=variance, + gamma=gamma, + beta=beta, + epsilon=epsilon, + ) + + apply_pass_and_basic_check(prog, "common::fuse_conv_batchnorm") + apply_pass_and_basic_check(prog, "common::const_elimination") + apply_pass_and_basic_check(prog, "common::dead_code_elimination") + + expected_ops = [] + if not weight_constexpr and not bias_constexpr: + expected_ops = ["conv"] + else: + expected_ops = [constexpr_op] * sum([weight_constexpr, bias_constexpr]) + ["conv", "batch_norm"] + assert get_op_types_in_program(prog) == expected_ops diff --git a/coremltools/converters/mil/mil/passes/topological_reorder.py b/coremltools/converters/mil/mil/passes/topological_reorder.py index ce9e05044..6f0e96dc3 100644 --- a/coremltools/converters/mil/mil/passes/topological_reorder.py +++ b/coremltools/converters/mil/mil/passes/topological_reorder.py @@ -4,10 +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.pass_registry import register_pass from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass +from coremltools.converters.mil.mil.passes.helper import block_context_manager +from coremltools.converters.mil.mil.passes.pass_registry import register_pass - +@block_context_manager 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 @@ -40,20 +41,19 @@ def _move_operations_to_the_end_block(block, op_type_to_move): first_use_indices = [block.operations.index(first_use_op) for first_use_op in first_consumers] before_op = block.operations[min(first_use_indices)] - with block: - # Create new copy of current operation - new_var = getattr(mb, op.op_type)(**op.inputs, before_op=before_op) + # Create new copy of current operation + new_var = getattr(mb, op.op_type)(**op.inputs, before_op=before_op) - if not isinstance(new_var, (list, tuple)): - new_var = [new_var] + if not isinstance(new_var, (list, tuple)): + new_var = [new_var] - # Override current_op to be newly created op to ensure `first_use` - # points to newly created op instead of old one. - current_op = new_var[0].op + # Override current_op to be newly created op to ensure `first_use` + # points to newly created op instead of old one. + current_op = new_var[0].op - for old_output_var, new_output_var in zip(op.outputs, new_var): - block.replace_uses_of_var_after_op( - anchor_op=None, old_var=old_output_var, new_var=new_output_var) + for old_output_var, new_output_var in zip(op.outputs, new_var): + block.replace_uses_of_var_after_op( + anchor_op=None, old_var=old_output_var, new_var=new_output_var) # Collect input vars from sub-block if present relevant_inputs = set() diff --git a/coremltools/converters/mil/mil/passes/use_reflection_padding.py b/coremltools/converters/mil/mil/passes/use_reflection_padding.py index 29ca65665..7df9bad3e 100644 --- a/coremltools/converters/mil/mil/passes/use_reflection_padding.py +++ b/coremltools/converters/mil/mil/passes/use_reflection_padding.py @@ -5,6 +5,7 @@ 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.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -80,25 +81,23 @@ def _match_pattern(concat_op, block): def _replace_ops(block, concat_op, slice_ops, axis): - with block: - - pad_size = len(slice_ops) // 2 - if axis == -1: - pad = [pad_size, pad_size] - elif axis == -2: - pad = [pad_size, pad_size, 0, 0] - else: - return False - - x = mb.pad(x=slice_ops[0].inputs["x"], pad=pad, mode='reflect', before_op=concat_op) - concat_op.enclosing_block.replace_uses_of_var_after_op( - anchor_op=concat_op, old_var=concat_op.outputs[0], new_var=x - ) + pad_size = len(slice_ops) // 2 + if axis == -1: + pad = [pad_size, pad_size] + elif axis == -2: + pad = [pad_size, pad_size, 0, 0] + else: + return False - block.remove_ops([concat_op] + slice_ops) - return True + x = mb.pad(x=slice_ops[0].inputs["x"], pad=pad, mode='reflect', before_op=concat_op) + concat_op.enclosing_block.replace_uses_of_var_after_op( + anchor_op=concat_op, old_var=concat_op.outputs[0], new_var=x + ) + block.remove_ops([concat_op] + slice_ops) + return True +@block_context_manager def _reflection_padding_block(block): for op in list(block.operations): _match_pattern(op, block) diff --git a/coremltools/converters/mil/mil/program.py b/coremltools/converters/mil/mil/program.py index 68708f312..ba83c88e4 100644 --- a/coremltools/converters/mil/mil/program.py +++ b/coremltools/converters/mil/mil/program.py @@ -11,7 +11,9 @@ from .block import Function from .var import Var from .types.symbolic import k_used_symbols, k_num_internal_syms +from coremltools.converters.mil._deployment_compatibility import AvailableTarget as _target from coremltools.converters.mil.input_types import InputType +from coremltools.converters.mil.mil.ops.helper import _get_version_of_op class Program: @@ -22,10 +24,65 @@ def __init__(self): self.parameters = {} self.skip_all_passes = False + def _get_max_opset_version_and_op(self): + max_opset_version = _target.iOS13 + op_with_max_opset_version = None + def update_max_opset_version_block(block): + nonlocal max_opset_version + nonlocal op_with_max_opset_version + for op in list(block.operations): + for b in op.blocks: + update_max_opset_version_block(b) + if not hasattr(op, "_op_variants") or not isinstance(op._op_variants, dict): + continue + if op.opset_version > max_opset_version: + max_opset_version = op.opset_version + op_with_max_opset_version = op + for func in self.functions.values(): + update_max_opset_version_block(func) + return max_opset_version, op_with_max_opset_version + + def _check_ops_version_compatibility(self, max_opset_version): + def check_version_compatibility_block(block): + for op in list(block.operations): + for b in op.blocks: + check_version_compatibility_block(b) + if not hasattr(op, "_op_variants") or not isinstance(op._op_variants, dict): + continue + expected_op_cls = _get_version_of_op(op._op_variants, max_opset_version) + if type(op) is not expected_op_cls: + msg = ( + "Op {} with an out of date version {!s} is detected. Please use @mb.program(input_specs=..., " + "opset_version={!s})" + ).format(op.op_type, op.opset_version, max_opset_version) + raise ValueError(msg) + for func in self.functions.values(): + check_version_compatibility_block(func) + + def _check_or_set_functions_opset_version(self, max_opset_version): + funcs = list(self.functions.values()) + for func in funcs: + if func.opset_version is None: + func.opset_version = max_opset_version + else: + if func.opset_version < max_opset_version: + msg = "function should have at least opset_version {!s}. Got {!s}".format(max_opset_version, func.opset_version) + raise ValueError(msg) + for func in funcs: + if func.opset_version != funcs[0].opset_version: + msg = "all functions must have the same opset_version. Got {!s} and {!s}.".format(func.opset_version, funcs[0].opset_version) + raise ValueError(msg) + + def _check_program_opset_version(self): + max_opset_version, _ = self._get_max_opset_version_and_op() + self._check_ops_version_compatibility(max_opset_version) + self._check_or_set_functions_opset_version(max_opset_version) + def add_function(self, name, ssa_func): if not isinstance(ssa_func, Function): raise ValueError("Only Function can be added to Program.") self.functions[name] = ssa_func + self._check_program_opset_version() def add_parameters(self, name, ssa_val): raise NotImplementedError() diff --git a/coremltools/converters/mil/mil/tests/test_block.py b/coremltools/converters/mil/mil/tests/test_block.py index b246c6b3b..a6ca12bbf 100644 --- a/coremltools/converters/mil/mil/tests/test_block.py +++ b/coremltools/converters/mil/mil/tests/test_block.py @@ -4,8 +4,11 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import copy +import numpy as _np +import pytest from coremltools.converters.mil.mil import Builder as mb +from coremltools.converters.mil.mil.passes.test_passes import CONSTEXPR_FUNCS from coremltools.converters.mil.testing_utils import ( assert_same_output_names, assert_same_output_shapes, @@ -38,7 +41,6 @@ def prog(x0): assert len(block.inputs) == 1 assert len(block.outputs) == 1 assert block.inputs["x0"] == block.outputs[0] - print(prog) def test_add_op(): @@ -194,6 +196,59 @@ def remove_transpose(block): assert_same_output_shapes(prev_prog, prog) +def test_replace_nonreplaceable_vars(): + """ + The conversion should error out if an invalid replacement is invoked with nonreplaceable vars + """ + constexpr_op = "constexpr_sparse_to_dense" + @mb.program(input_specs=[mb.TensorSpec(shape=(4, 2))]) + def prog(x): + constexpr = CONSTEXPR_FUNCS[constexpr_op]((4, 2)) + return mb.add(x=x, y=constexpr) + + block = prog.functions["main"] + constexpr_op = block.find_ops(op_type=constexpr_op)[0] + + with block: + const = mb.const(val=_np.random.rand(4, 2), before_op=constexpr_op) + expected_err_str = "might potentially be removed during the replacement of those vars." + with pytest.raises(ValueError, match=expected_err_str): + block.replace_uses_of_var_after_op( + anchor_op=constexpr_op, + old_var=constexpr_op.outputs[0], + new_var=const + ) + +def test_replace_nonreplaceable_vars_force(): + """ + The conversion should not error out if the replace_uses_of_vars_after_op is executed with force_replace=True + Also we test that, the new nonreplaceable_vars_upstream is propagated after the code exist `with block`. + """ + constexpr_op = "constexpr_sparse_to_dense" + @mb.program(input_specs=[mb.TensorSpec(shape=(4, 2))]) + def prog(x): + constexpr = CONSTEXPR_FUNCS[constexpr_op]((4, 2)) + return mb.add(x=x, y=constexpr) + + block = prog.functions["main"] + constexpr_op = block.find_ops(op_type=constexpr_op)[0] + add_op = block.find_ops(op_type="add")[0] + + assert len(add_op.outputs[0].nonreplaceable_vars_upstream) == 1 + + with block: + const = mb.const(val=_np.random.rand(4, 2), before_op=constexpr_op) + block.replace_uses_of_var_after_op( + anchor_op=constexpr_op, + old_var=constexpr_op.outputs[0], + new_var=const, + force_replace=True, + ) + block.remove_ops([constexpr_op]) + + assert len(add_op.outputs[0].nonreplaceable_vars_upstream) == 0 + + def test_simple_substituion(): """Replace log(x+y) with log(x*y) """ diff --git a/coremltools/converters/mil/mil/tests/test_programs.py b/coremltools/converters/mil/mil/tests/test_programs.py index 5371beb8e..88ed72e3f 100644 --- a/coremltools/converters/mil/mil/tests/test_programs.py +++ b/coremltools/converters/mil/mil/tests/test_programs.py @@ -5,6 +5,7 @@ import logging import numpy as np +import pytest import coremltools as ct from coremltools.converters.mil.mil import Builder as mb @@ -138,3 +139,142 @@ def prog(a, b): if ct.utils._is_macos(): prediction = mlmodel.predict(feed_dict) assert len(prediction) == 2 + +def get_simple_topk_program(opset_version=None): + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=opset_version) + def prog(x): + x = mb.topk(x=x, k=1, axis=-1, ascending=True) + return x + return prog + +def get_simple_pixel_unshuffle_program(opset_version=None): + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=opset_version) + def prog(x): + x = mb.pixel_unshuffle(x=x, downscale_factor=np.uint32(2)) + return x + return prog + +def get_simple_topk_pixel_unshuffle_program(opset_version=None): + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=opset_version) + def prog(x): + x = mb.pixel_unshuffle(x=x, downscale_factor=np.uint32(2)) + x = mb.topk(x=x, k=1, axis=-1, ascending=True) + return x + return prog + +def get_simple_nested_block_program(opset_version=None): + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=opset_version) + def prog(x): + def true_fn(): + topk, _ = mb.topk(x=x, k=1, axis=-1, ascending=True) + return mb.add(x=topk, y=1) + + def false_fn(): + topk, _ = mb.topk(x=x, k=1, axis=-1, ascending=True) + return mb.add(x=topk, y=2) + + shape = mb.shape(x=x) + rank = mb.shape(x=shape) + pred = mb.squeeze(x=rank) + return mb.cond(pred=mb.cast(x=pred, dtype="bool"), _true_fn=true_fn, _false_fn=false_fn) + return prog + +class TestMLProgramVersionHandling: + + @staticmethod + def test_multi_versions_op_selection(): + ''' + Builder should pick up the right version of op based on opset_version + ''' + # pick up the oldest version (iOS13) topk by default + prog = get_simple_topk_program() + main_func = prog.functions["main"] + topk_op = main_func.find_ops(op_type="topk")[0] + assert topk_op.opset_version == ct.target.iOS13 + + # pick up iOS13 version topk + prog = get_simple_topk_program(opset_version=ct.target.iOS15) + main_func = prog.functions["main"] + topk_op = main_func.find_ops(op_type="topk")[0] + assert topk_op.opset_version == ct.target.iOS13 + + # pick up iOS16 version topk + prog = get_simple_topk_program(opset_version=ct.target.iOS16) + main_func = prog.functions["main"] + topk_op = main_func.find_ops(op_type="topk")[0] + assert topk_op.opset_version == ct.target.iOS16 + + @staticmethod + def test_pymil_front_end_conversion(): + prog = get_simple_topk_pixel_unshuffle_program(opset_version=ct.target.iOS16) + mlmodel = ct.convert(prog, minimum_deployment_target=ct.target.iOS16) + + @staticmethod + def test_nested_block_opset_version_selection(): + # pick up the oldest version (iOS13) topk by default + prog = get_simple_nested_block_program() + main_func = prog.functions["main"] + topk_ops = main_func.find_ops(op_type="topk") + assert all([topk.opset_version == ct.target.iOS13 for topk in topk_ops]) + + # pick up iOS16 version topk + prog = get_simple_nested_block_program(opset_version=ct.target.iOS16) + main_func = prog.functions["main"] + topk_ops = main_func.find_ops(op_type="topk") + assert all([topk.opset_version == ct.target.iOS16 for topk in topk_ops]) + + @staticmethod + def test_pymil_opset_version_inference(): + ''' + The program consist of pixel_unshuffle should be inferred as an iOS16 version program + ''' + prog = get_simple_pixel_unshuffle_program() + assert prog.functions["main"].opset_version == ct.target.iOS16 + + expected_err_str = ( + "Please update the minimum_deployment_target to coremltools.target.iOS16, " + "since op pixel_unshuffle is only available in opset coremltools.target.iOS16 or newer." + ) + with pytest.raises(ValueError, match=expected_err_str): + mlmodel = ct.convert(prog, convert_to="mlprogram") + + @staticmethod + def test_pymil_front_end_conversion_early_error_out(): + prog = get_simple_topk_pixel_unshuffle_program(opset_version=ct.target.iOS16) + expected_err_str = ( + "Please update the minimum_deployment_target to coremltools.target.iOS16, " + "since op pixel_unshuffle is only available in opset coremltools.target.iOS16 or newer." + ) + with pytest.raises(ValueError, match=expected_err_str): + mlmodel = ct.convert(prog, minimum_deployment_target=ct.target.iOS15) + + @staticmethod + def test_unsupported_op_early_error_out(): + ''' + We should error out at the point when Builder tries to add an op which is only supported in a newer spec version + ''' + expected_err_str = ( + "No available version for pixel_unshuffle in the coremltools.target.iOS15 opset. " + "Please update the minimum_deployment_target to at least coremltools.target.iOS16" + ) + with pytest.raises(ValueError, match=expected_err_str): + @mb.program(input_specs=[mb.TensorSpec(shape=(1, 1, 4, 4))], opset_version=ct.target.iOS15) + def prog(x): + x = mb.pixel_unshuffle(x=x, downscale_factor=np.uint32(2)) + return x + + @staticmethod + def test_bulid_non_compatible_program_early_error_out(): + ''' + `mb.program` API should detect potential non compatible ops in the program, and error out early + In this example, `pixel_unshuffle` is an iO16 op, and `topk` has iOS13 and iOS16 version. + If the builder version is not set, it is picking up the iOS13 version of topk, which would + potentially create an invalid program. + In this case, `mb.program` should error out, and tell the user to set `opset_version=target.iOS16` + ''' + expected_err_str = ( + "Op topk with an out of date version coremltools.target.iOS13 is detected. Please use @mb.program\(input_specs=..., opset_version=coremltools.target.iOS16\)" + ) + with pytest.raises(ValueError, match=expected_err_str): + get_simple_topk_pixel_unshuffle_program() + diff --git a/coremltools/converters/mil/mil/var.py b/coremltools/converters/mil/mil/var.py index 57260933f..17b545fa8 100644 --- a/coremltools/converters/mil/mil/var.py +++ b/coremltools/converters/mil/mil/var.py @@ -70,6 +70,9 @@ class Var: child_ops [_child_ops]: list[Operation] Ops that take this Var as an input. + + nonreplaceable_vars_upstream: set[Var] + Set that consists of nonreplaceable vars upstream """ __slots__ = [ @@ -80,6 +83,7 @@ class Var: "op_output_idx", "_child_ops", "consuming_blocks", + "_nonreplaceable_vars_upstream", ] def __init__(self, name, sym_type, sym_val=None, op=None, op_output_idx=None): @@ -102,6 +106,52 @@ def __init__(self, name, sym_type, sym_val=None, op=None, op_output_idx=None): # == 0) but is still used as block output. A var can be output of # multiple blocks (e.g., both current block and nested blocks) self.consuming_blocks = list() + + # replaceability + self._nonreplaceable_vars_upstream = set() + self._set_nonreplaceable_vars_upstream() + + @property + def nonreplaceable_vars_upstream(self): + return self._nonreplaceable_vars_upstream + + @nonreplaceable_vars_upstream.setter + def nonreplaceable_vars_upstream(self, val): + assert isinstance(val, set) + self._nonreplaceable_vars_upstream = val + + @staticmethod + def _is_nonreplaceable_var(var): + op = var.op + if op is None: + return False + return op.op_type.startswith("constexpr_") + + def _set_nonreplaceable_vars_upstream(self): + """ + A utility function to set the value of the "nonreplaceable_vars_upstream" property. + If the var is an output of the constexpr op, then "nonreplaceable_vars_upstream" is a single element set, containing this var. + Otherwise, its a union of the "nonreplaceable_vars_upstream" sets of all the input vars of its parent op. + """ + op = self.op + if op is None: + return + if Var._is_nonreplaceable_var(self): + self.nonreplaceable_vars_upstream = set([self]) + else: + flattened_inputs = op.get_flattened_inputs() + inputs_nonreplaceable_vars_upstream = [p.nonreplaceable_vars_upstream for p in flattened_inputs] + if len(inputs_nonreplaceable_vars_upstream) > 0: + self.nonreplaceable_vars_upstream = set.union(*inputs_nonreplaceable_vars_upstream) + + def _reset_nonreplaceable_vars_upstream(self): + self.nonreplaceable_vars_upstream = set() + + def can_be_replaced_by_var(self, new_var): + """ + A var can be replaced by a new var only if the new var's nonreplaceable_vars_upstream is the super set of the old one + """ + return self.nonreplaceable_vars_upstream.issubset(new_var.nonreplaceable_vars_upstream) @property def sym_type(self): diff --git a/coremltools/converters/mil/testing_utils.py b/coremltools/converters/mil/testing_utils.py index cf9e91784..3c5f63709 100644 --- a/coremltools/converters/mil/testing_utils.py +++ b/coremltools/converters/mil/testing_utils.py @@ -7,11 +7,10 @@ from functools import partial import os from pathlib import Path -from PIL import Image import re import numpy as np -import PIL.Image +from PIL import Image import coremltools as ct from coremltools._deps import _IS_MACOS @@ -178,7 +177,7 @@ def to_tuple(v): def run_core_ml_predict(mlmodel, input_key_values): for k, v in input_key_values.items(): - if isinstance(v, PIL.Image.Image): + if isinstance(v, Image.Image): continue elif not np.isscalar(v) and not v.shape == (): input_key_values[k] = v.astype(np.float32) @@ -440,10 +439,16 @@ def random_gen_input_feature_type(input_desc): shape = [3, input_desc.type.imageType.height, input_desc.type.imageType.width] x = np.random.randint(low=0, high=256, size=shape) return Image.fromarray(np.transpose(x, [1, 2, 0]).astype(np.uint8)) - else: + elif input_desc.type.imageType.colorSpace == ft.ImageFeatureType.GRAYSCALE: shape = [input_desc.type.imageType.height, input_desc.type.imageType.width] x = np.random.randint(low=0, high=256, size=shape) return Image.fromarray(x.astype(np.uint8), 'L') + elif input_desc.type.imageType.colorSpace == ft.ImageFeatureType.GRAYSCALE_FLOAT16: + shape = (input_desc.type.imageType.height, input_desc.type.imageType.width) + x = np.random.rand(*shape) + return Image.fromarray(x.astype(np.float32), 'F') + else: + raise ValueError("unrecognized image type") else: raise ValueError('unsupported type') diff --git a/coremltools/models/ml_program/compression_utils.py b/coremltools/models/ml_program/compression_utils.py index 36e8c8e80..ae6e40236 100644 --- a/coremltools/models/ml_program/compression_utils.py +++ b/coremltools/models/ml_program/compression_utils.py @@ -43,27 +43,22 @@ def _apply_graph_pass(mlmodel, graph_pass): raise TypeError("weight compression not applicable for model type {}".format(model_type)) assert isinstance(graph_pass, _AbstractQuantizationPass), "compression pass must be an AbstractQuantizationPass instance" - - program_spec = model_spec.mlProgram - model_specification_version = model_spec.specificationVersion - prog = _milproto_to_pymil( - program_spec=program_spec, - specification_version=model_specification_version, + specification_version = max(model_spec.specificationVersion, _DEFAULT_SPECIFICATION_VERSION_FOR_COMPRESSION) + prog = _milproto_to_pymil( + model_spec=model_spec, + specification_version=specification_version, file_weights_dir=mlmodel.weights_dir, ) - prog.skip_all_passes = True - # apply compression graph pass graph_pass.apply(prog) # convert the pymil program back to mlmodel - compress_model_specification_version = max(model_specification_version, _DEFAULT_SPECIFICATION_VERSION_FOR_COMPRESSION) compressed_mlmodel = _mil_convert( prog, convert_to="mlprogram", convert_from="milinternal", - specification_version=compress_model_specification_version, + specification_version=specification_version, compute_units=mlmodel.compute_unit, model_description=model_spec.description, ) diff --git a/coremltools/models/model.py b/coremltools/models/model.py index 16ead447a..88272da9c 100644 --- a/coremltools/models/model.py +++ b/coremltools/models/model.py @@ -12,8 +12,9 @@ import numpy as _numpy from ..proto import ( + FeatureTypes_pb2 as _ft, Model_pb2 as _Model_pb2, - MIL_pb2 as _MIL_pb2 + MIL_pb2 as _MIL_pb2, ) from .utils import ( _create_mlpackage, @@ -47,6 +48,12 @@ except: _ModelPackage = None +_HAS_PIL = True +try: + from PIL import Image as _PIL_IMAGE +except: + _HAS_PIL = False + _MLMODEL_FULL_PRECISION = "float32" _MLMODEL_HALF_PRECISION = "float16" @@ -308,6 +315,13 @@ def __init__(self, model, """ if not isinstance(compute_units, _ComputeUnit): raise TypeError('"compute_units" parameter must be of type: coremltools.ComputeUnit') + elif (compute_units == _ComputeUnit.CPU_AND_NE + and _is_macos() + and _macos_version() < (13, 0) + ): + raise ValueError( + 'coremltools.ComputeUnit.CPU_AND_NE is only available on macOS >= 13.0' + ) self.compute_unit = compute_units self.is_package = False @@ -500,11 +514,7 @@ def predict(self, data): ) if self.__proxy__: - # Check if the input name given by the user is valid. - # Although this is checked during prediction inside CoreML Framework, - # we still check it here to return early and - # return a more verbose error message - self._verify_input_name_exists(data) + self._verify_input_dict(data) self._convert_tensor_to_numpy(data) # TODO: remove the following call when this is fixed: rdar://92239209 self._update_float16_multiarray_input_to_float32(data) @@ -592,6 +602,37 @@ def _get_mil_internal(self): """ return _deepcopy(self._mil_program) + def _verify_input_dict(self, input_dict): + # Check if the input name given by the user is valid. + # Although this is checked during prediction inside CoreML Framework, + # we still check it here to return early and + # return a more verbose error message + self._verify_input_name_exists(input_dict) + + # verify that the pillow image modes are correct, for image inputs + self._verify_pil_image_modes(input_dict) + + def _verify_pil_image_modes(self, input_dict): + if not _HAS_PIL: + return + for input_desc in self._spec.description.input: + if input_desc.type.WhichOneof("Type") == "imageType": + input_val = input_dict.get(input_desc.name, None) + if not isinstance(input_val, _PIL_IMAGE.Image): + msg = "Image input, '{}' must be of type PIL.Image.Image in the input dict" + raise TypeError(msg.format(input_desc.name)) + if input_desc.type.imageType.colorSpace in (_ft.ImageFeatureType.BGR, _ft.ImageFeatureType.RGB): + if input_val.mode != 'RGB': + msg = "RGB/BGR image input, '{}', must be of type PIL.Image.Image with mode=='RGB'" + raise TypeError(msg.format(input_desc.name)) + elif input_desc.type.imageType.colorSpace == _ft.ImageFeatureType.GRAYSCALE: + if input_val.mode != 'L': + msg = "GRAYSCALE image input, '{}', must be of type PIL.Image.Image with mode=='L'" + raise TypeError(msg.format(input_desc.name)) + elif input_desc.type.imageType.colorSpace == _ft.ImageFeatureType.GRAYSCALE_FLOAT16: + if input_val.mode != 'F': + msg = "GRAYSCALE_FLOAT16 image input, '{}', must be of type PIL.Image.Image with mode=='F'" + raise TypeError(msg.format(input_desc.name)) def _verify_input_name_exists(self, input_dict): model_input_names = [inp.name for inp in self._spec.description.input] diff --git a/coremltools/models/neural_network/quantization_utils.py b/coremltools/models/neural_network/quantization_utils.py index 36351566d..03a8ab2e4 100644 --- a/coremltools/models/neural_network/quantization_utils.py +++ b/coremltools/models/neural_network/quantization_utils.py @@ -382,7 +382,7 @@ def _get_kmeans_lookup_table_and_weight( if _HAS_SKLEARN: from sklearn.cluster import KMeans else: - raise Exception("sklearn package required for k-means quantization") + raise Exception("scikit-learn package required for k-means quantization") units = _np.prod(w.shape) lut_len = 1 << nbits n_clusters = units if (units < lut_len) else lut_len diff --git a/coremltools/test/api/test_api_examples.py b/coremltools/test/api/test_api_examples.py index 9177ca981..e28bd3d72 100644 --- a/coremltools/test/api/test_api_examples.py +++ b/coremltools/test/api/test_api_examples.py @@ -732,10 +732,7 @@ def test_unsanitized_input_name_during_prediction(): assert "does not match any of the model input" in error_str @staticmethod - @pytest.mark.parametrize( - "to_tensor", [torch.tensor, tf.convert_to_tensor, lambda x: x.tolist()]) - @pytest.mark.skipif(not ct.utils._is_macos(), reason="Platform is not Mac OS") - def test_variant_input_type_prediction(to_tensor): + def _test_variant_input_type_prediction(to_tensor): prog = Program() func_inputs = {"x": mb.placeholder(shape=[2, 3]), "y": mb.placeholder(shape=[2, 3])} @@ -759,6 +756,24 @@ def test_variant_input_type_prediction(to_tensor): ) np.allclose(out_by_numpy["out"], out_by_tensor["out"]) + @staticmethod + @pytest.mark.skipif(not ct.utils._is_macos(), reason="test needs predictions") + def test_list_predict_input(): + TestInputs._test_variant_input_type_prediction(lambda x: x.tolist()) + + @staticmethod + @pytest.mark.skipif(not _HAS_TF_1, reason="requires TensorFlow") + @pytest.mark.skipif(not ct.utils._is_macos(), reason="test needs predictions") + def test_tf_predict_input(): + TestInputs._test_variant_input_type_prediction(tf.convert_to_tensor) + + @staticmethod + @pytest.mark.skipif(not _HAS_TORCH, reason="requires Torch") + @pytest.mark.skipif(not ct.utils._is_macos(), reason="test needs predictions") + def test_torch_predict_input(): + TestInputs._test_variant_input_type_prediction(torch.tensor) + + class TestFlexibleShape: @staticmethod @pytest.mark.parametrize( @@ -1210,6 +1225,29 @@ def test_torch_image_enumerated_shapes(): spec = model.get_spec() assert len(spec.description.input[0].type.imageType.enumeratedSizes.sizes) == 2 + @staticmethod + @pytest.mark.skipif(not _HAS_TORCH, reason=MSG_TORCH_NOT_FOUND) + def test_mean_op(): + # test for bug reported in https://github.com/apple/coremltools/issues/1420 + class Network(torch.nn.Module): + def forward(self, x): + return torch.mean(x, dim=(2, 3), keepdim=True) + + model = Network() + x = torch.rand(1, 3, 256, 256) + traced_model = torch.jit.trace(model, x) + input_x = ct.TensorType(shape=(1, 3, ct.RangeDim(default=256), ct.RangeDim(default=256)), + name="input") + cml = ct.convert(traced_model, + inputs=[input_x], + outputs=[ct.TensorType(name="out")], + convert_to="mlprogram", + compute_units=ct.ComputeUnit.CPU_ONLY) + input_dict = {"input": np.random.rand(1, 3, 112, 112)} + if ct.utils._is_macos(): + out = cml.predict(input_dict)["out"] + assert out.shape == (1, 3, 1, 1) + class TestOptionalInput: @staticmethod @@ -1330,6 +1368,9 @@ def forward(self, x, y): default_value=default_value) for compute_units in ct.ComputeUnit: + if compute_units == ct.ComputeUnit.CPU_AND_NE and ct.utils._macos_version() < (13, 0): + continue + mlmodel = ct.convert( traced_model, inputs=[required_input, optional_input], diff --git a/coremltools/test/api/test_api_visibilities.py b/coremltools/test/api/test_api_visibilities.py index bd614e4fc..2daf0cd1e 100644 --- a/coremltools/test/api/test_api_visibilities.py +++ b/coremltools/test/api/test_api_visibilities.py @@ -17,11 +17,7 @@ def _check_visible_modules(actual, expected): ) -class TestApiVisibilities: - """Test public coremltools API visibilities.""" - - def test_top_level(self): - expected = [ +EXPECTED_MODULES = [ "ClassifierConfig", "ComputeUnit", "EnumeratedShapes", @@ -46,9 +42,15 @@ def test_top_level(self): "libmodelpackage", "libmilstoragepython", ] + + +class TestApiVisibilities: + """Test public coremltools API visibilities.""" + + def test_top_level(self): if not ct.utils._is_macos(): - expected.remove("libcoremlpython") - _check_visible_modules(_get_visible_items(ct), expected) + EXPECTED_MODULES.remove("libcoremlpython") + _check_visible_modules(_get_visible_items(ct), EXPECTED_MODULES) def test_utils(self): expected = [ diff --git a/coremltools/test/ml_program/test_compression.py b/coremltools/test/ml_program/test_compression.py index 29bbaa924..98952a59c 100644 --- a/coremltools/test/ml_program/test_compression.py +++ b/coremltools/test/ml_program/test_compression.py @@ -4,14 +4,17 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import logging -import pytest +import unittest + import numpy as np +import pytest import torch import coremltools as ct - +from coremltools._deps import _HAS_SKLEARN from coremltools.converters.mil.testing_utils import get_op_types_in_program + def create_unique_weight(weight, nbits): shape = weight.detach().numpy().shape size = weight.detach().numpy().size @@ -92,6 +95,7 @@ def test_op_selector(): assert get_op_types_in_program(mlmodel_no_quantized._mil_program) == expected_ops @staticmethod + @unittest.skipIf(not _HAS_SKLEARN, "Missing scikit-learn. Skipping tests.") def test_weight_decompression(): """ This test is doing the following steps @@ -158,12 +162,12 @@ def test_compression_utils_error_handling(): mlmodel = ct.convert(torchmodel, inputs=inputs, convert_to="mlprogram") # Test invalid mode for affine quantization - expected_err_str = "supported for weight affine quantization. Got mode invalid_mode." + expected_err_str = "supported for weight affine quantization. Got mode" with pytest.raises(ValueError, match=expected_err_str): ct.compression_utils.affine_quantize_weights(mlmodel, mode="invalid_mode") # Test invalid mode for weight sparsification - expected_err_str = "supported for weight sparsification. Got mode invalid_mode." + expected_err_str = "supported for weight sparsification. Got mode" with pytest.raises(ValueError, match=expected_err_str): ct.compression_utils.sparsify_weights(mlmodel, mode="invalid_mode") @@ -178,7 +182,7 @@ def test_compression_utils_error_handling(): ct.compression_utils.sparsify_weights(mlmodel, mode="percentile_based", target_percentile=1.2) # Test invalid mode for weight palettization - expected_err_str = "supported for weight palettization. Got mode invalid_mode." + expected_err_str = "supported for weight palettization. Got mode" with pytest.raises(ValueError, match=expected_err_str): ct.compression_utils.palettize_weights(mlmodel, mode="invalid_mode") @@ -298,7 +302,7 @@ def test_weight_sparsify_percentile_based(percentile): @staticmethod @pytest.mark.parametrize( "mode", - ("uniform", "kmeans") + ("uniform", "kmeans") if _HAS_SKLEARN else ("uniform",) ) def test_weight_palettization(mode): model, inputs, torch_input_values, coreml_input_values = get_test_model_and_data() @@ -407,4 +411,4 @@ def lut_function(weight): assert lut_to_dense_op.shape.val.tolist() == list(model.weight.detach().numpy().shape) # validate the model - verify_model_outputs(mlmodel, mlmodel_palettized, coreml_input_values) \ No newline at end of file + verify_model_outputs(mlmodel, mlmodel_palettized, coreml_input_values) diff --git a/coremltools/test/modelpackage/test_modelpackage.py b/coremltools/test/modelpackage/test_modelpackage.py index 7b905ea66..003c04b6c 100644 --- a/coremltools/test/modelpackage/test_modelpackage.py +++ b/coremltools/test/modelpackage/test_modelpackage.py @@ -117,6 +117,10 @@ def test_predict_api(self): if utils._macos_version() >= (12, 0): for compute_units in coremltools.ComputeUnit: + if (compute_units == coremltools.ComputeUnit.CPU_AND_NE + and utils._macos_version() < (13, 0)): + continue + loaded_model = MLModel(package.name, compute_units=compute_units) preds = loaded_model.predict({"feature_1": 1.0, "feature_2": 1.0}) diff --git a/coremltools/test/neural_network/test_numpy_nn_layers.py b/coremltools/test/neural_network/test_numpy_nn_layers.py index 37428ef3c..fc7f21cda 100644 --- a/coremltools/test/neural_network/test_numpy_nn_layers.py +++ b/coremltools/test/neural_network/test_numpy_nn_layers.py @@ -917,52 +917,6 @@ def test_dynamic_weight_conv(self): self._test_model(builder.spec, feed_dict, expected, useCPUOnly=True) self._test_model(builder.spec, feed_dict, expected, useCPUOnly=False) - @pytest.mark.xfail - def test_dynamic_weight_deconv(self): - # Expect to fail in Core ML 3 - input_dim = (1, 1, 16, 16) - # weight layout: (output_channels, kernel_channels, height, width) - weight_dim = (1, 1, 3, 3) - output_dim = (1, 1, 18, 18) - output_channels, kernel_channels, height, width = weight_dim - - input_features = [ - ("data", datatypes.Array(*input_dim)), - ("weight", datatypes.Array(*weight_dim)), - ] - output_features = [("output", None)] - - builder = neural_network.NeuralNetworkBuilder( - input_features, output_features, disable_rank5_shape_mapping=True - ) - - builder.add_convolution( - name="deconv", - kernel_channels=kernel_channels, - output_channels=output_channels, - height=height, - width=width, - stride_height=1, - stride_width=1, - border_mode="valid", - groups=1, - W=None, - b=None, - has_bias=False, - is_deconv=True, - input_name=["data", "weight"], - output_name="output", - ) - - input_val = np.ones(input_dim) - weight_val = np.ones(weight_dim) - expected = np.ones(output_dim) * 27 - - feed_dict = {"data": input_val, "weight": weight_val} - expected = {"output": expected} - - self._test_model(builder.spec, feed_dict, expected) - def test_batched_mat_mul_cpu(self, cpu_only=True): a_shapes = [ (10,), @@ -1754,7 +1708,6 @@ def test_elementwise_binary_cpu(self, cpu_only=True): expected = {"output": func(a, b, dtype=np.float32)} self._test_model(builder.spec, input, expected, useCPUOnly=cpu_only) - @pytest.mark.xfail(reason="rdar://93912621") def test_elementwise_binary_gpu(self): self.test_elementwise_binary_cpu(cpu_only=False) @@ -4048,8 +4001,8 @@ def softmax(data): if isinstance(model, str): model = coremltools.models.MLModel(model) - model = coremltools.models.MLModel(model, useCPUOnly=True) - prediction = model.predict(inputs, useCPUOnly=True) + model = coremltools.models.MLModel(model) + prediction = model.predict(inputs) # validate each distribution separately logits = x.reshape(2, num_class) @@ -4324,7 +4277,6 @@ def test_reshape_like_cpu(self, cpu_only=True): self._test_model(builder.spec, inputs, expected, useCPUOnly=cpu_only) self.assertEqual(target_rank, builder._get_rank("output")) - @pytest.mark.xfail(reason="Fixed in https://github.com/apple/coremltools/pull/634") def test_reshape_like_gpu(self): self.test_reshape_like_cpu(cpu_only=False) @@ -6744,6 +6696,7 @@ def test_slice_by_size_cpu(self, cpu_only=True): slices.append(slice(None, None, None)) else: slices.append(slice(begin, begin + size, 1)) + slices = tuple(slices) expected = {"output": x[slices]} input_features = [ diff --git a/coremltools/test/neural_network/test_quantization.py b/coremltools/test/neural_network/test_quantization.py index c76c15dd3..c0ddebdcd 100644 --- a/coremltools/test/neural_network/test_quantization.py +++ b/coremltools/test/neural_network/test_quantization.py @@ -2,11 +2,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 -"""Module containing unit tests for verifying various quantization.""" + +""" +Module containing unit tests for verifying various quantizations. +""" + +import unittest import numpy as np import pytest -import unittest import coremltools from coremltools import ComputeUnit @@ -557,3 +561,9 @@ def test_embeddingND_quantize(compute_units): 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 + + @unittest.skipIf(coremltools.utils._macos_version() < (13, 0), + 'ComputeUnit.CPU_AND_NE is only available on macOS >= 13.0' + ) + def test_embeddingND_quantize_CPU_and_NE(self): + self.test_embeddingND_quantize(ComputeUnit.CPU_AND_NE) diff --git a/coremltools/test/sklearn_tests/test_random_forest_regression.py b/coremltools/test/sklearn_tests/test_random_forest_regression.py index f1c54bb44..e78258d99 100644 --- a/coremltools/test/sklearn_tests/test_random_forest_regression.py +++ b/coremltools/test/sklearn_tests/test_random_forest_regression.py @@ -4,7 +4,7 @@ # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause import unittest -from sklearn.ensemble import RandomForestRegressor + from coremltools._deps import _HAS_SKLEARN from coremltools.proto import Model_pb2 from coremltools.proto import FeatureTypes_pb2 diff --git a/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py b/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py index 73fd76059..0f3c81e74 100644 --- a/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py +++ b/coremltools/test/xgboost_tests/test_boosted_trees_classifier.py @@ -3,20 +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 -import unittest -import tempfile import json +import tempfile +import unittest -from sklearn.ensemble import GradientBoostingClassifier from coremltools.converters import sklearn as skl_converter from coremltools.models.utils import _macos_version from coremltools.proto import FeatureTypes_pb2, Model_pb2 from coremltools._deps import _HAS_SKLEARN, _HAS_XGBOOST +if _HAS_SKLEARN: + from sklearn.ensemble import GradientBoostingClassifier + if _HAS_XGBOOST: import xgboost from coremltools.converters import xgboost as xgb_converter + @unittest.skipIf(not _HAS_SKLEARN, "Missing sklearn. Skipping tests.") class GradientBoostingBinaryClassifierScikitTest(unittest.TestCase): """ diff --git a/coremltools/version.py b/coremltools/version.py index 55bb29f4c..997ca1738 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__ = "6.0b1" # VERSION_STRING +__version__ = "6.0b2" # VERSION_STRING diff --git a/reqs/build.pip b/reqs/build.pip index 74fe98c7d..48ba4e42c 100644 --- a/reqs/build.pip +++ b/reqs/build.pip @@ -1,3 +1,4 @@ +numpy==1.21.0; platform_machine == "arm64" numpy<1.20; platform_machine != "arm64" # rdar://93977023 diff --git a/reqs/test.pip b/reqs/test.pip index 425baa6fc..f595cb303 100644 --- a/reqs/test.pip +++ b/reqs/test.pip @@ -12,14 +12,14 @@ pytest; python_version < '3.7' pytest==5.3.4; python_version >= '3.7' pytest-cov pytest-sugar -scikit-learn==0.19.2 +scikit-learn==0.19.2; python_version < '3.8' scipy > 1.4 six sympy > 1.6 tensorflow==1.15.0; python_version < '3.8' torch==1.11.0 torchvision==0.12.0 -xgboost==1.4.2 +xgboost==1.4.2; platform_machine != "arm64" mock wrapt tqdm diff --git a/reqs/test_tf2.pip b/reqs/test_tf2.pip index 9f3811145..e279216b8 100644 --- a/reqs/test_tf2.pip +++ b/reqs/test_tf2.pip @@ -5,4 +5,4 @@ keras==2.8.0 tensorflow-addons==0.16.1 tensorflow-hub==0.12.0 -transformers==2.10.0; python_version > '3.6' +transformers==4.17.0; python_version > '3.6' diff --git a/scripts/env_create.sh b/scripts/env_create.sh index 29480a72d..ed723865b 100755 --- a/scripts/env_create.sh +++ b/scripts/env_create.sh @@ -91,11 +91,6 @@ fi echo "Installing basic build requirements." if [[ $include_build_deps == 1 ]]; then python -m pip install -r $COREMLTOOLS_HOME/reqs/build.pip - - if [[ $(uname -m) == "arm64" ]]; then - # Install the earliest version of numpy, we can find, that support arm64 - conda install -c apple numpy==1.19.5 -y - fi fi # Install test requirements (upgrades packages if required)