From 4a331ef66726dbf27221b4ac436a7919704b15d2 Mon Sep 17 00:00:00 2001 From: Edward Chen <18449977+edgchen1@users.noreply.github.com> Date: Mon, 26 Jun 2023 15:29:50 -0700 Subject: [PATCH] Rework CPU MeanVarianceNormalization kernel to support arbitrary axes. (#16420) --- .../cpu/mean_variance_normalization_exp.cc | 2 +- onnxruntime/core/providers/common.h | 9 +- .../cpu/tensor/mean_variance_normalization.cc | 193 ++++++++++++- .../cpu/tensor/mean_variance_normalization.h | 119 +------- .../test/common/random_generator_test.cc | 10 +- .../test/common/tensor_op_test_utils.h | 8 +- .../providers/cpu/tensor/gen_mvn_test_data.py | 37 +++ .../mean_variance_normalization_test.cc | 273 ++++++++++++++++++ .../providers/cpu/tensor/tensor_op_test.cc | 94 ------ 9 files changed, 527 insertions(+), 218 deletions(-) create mode 100644 onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py create mode 100644 onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc diff --git a/onnxruntime/contrib_ops/cpu/mean_variance_normalization_exp.cc b/onnxruntime/contrib_ops/cpu/mean_variance_normalization_exp.cc index 564a49d33c..b10a622be8 100644 --- a/onnxruntime/contrib_ops/cpu/mean_variance_normalization_exp.cc +++ b/onnxruntime/contrib_ops/cpu/mean_variance_normalization_exp.cc @@ -13,6 +13,6 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( 1, 8, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - MeanVarianceNormalization_0); + MeanVarianceNormalization); } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/providers/common.h b/onnxruntime/core/providers/common.h index c821ff0650..6990bb9f63 100644 --- a/onnxruntime/core/providers/common.h +++ b/onnxruntime/core/providers/common.h @@ -13,6 +13,13 @@ namespace onnxruntime { +/** +Returns whether `axis` is in range [-`tensor_rank`, `tensor_rank`). +**/ +constexpr inline bool IsAxisInRange(int64_t axis, int64_t tensor_rank) { + return axis >= -tensor_rank && axis <= tensor_rank - 1; +} + /** Handle a potentially negative axis. Enforces negative axis is valid. @param axis Axis to convert from negative to positive if needed. @@ -20,7 +27,7 @@ Handle a potentially negative axis. Enforces negative axis is valid. @returns non-negative axis. */ inline int64_t HandleNegativeAxis(int64_t axis, int64_t tensor_rank) { - ORT_ENFORCE(axis >= -tensor_rank && axis <= tensor_rank - 1, "axis ", axis, + ORT_ENFORCE(IsAxisInRange(axis, tensor_rank), "axis ", axis, " is not in valid range [-", tensor_rank, ",", tensor_rank - 1, "]"); // Handle negative axis return axis < 0 ? axis + tensor_rank : axis; diff --git a/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.cc b/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.cc index 6f8827508d..62c3dbfc87 100644 --- a/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.cc +++ b/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.cc @@ -3,17 +3,206 @@ #include "core/providers/cpu/tensor/mean_variance_normalization.h" +#include +#include + +#include "core/common/gsl.h" +#include "core/providers/common.h" +#include "core/providers/cpu/tensor/transpose.h" +#include "core/util/math_cpuonly.h" + namespace onnxruntime { + +namespace { +InlinedVector GetAxesFromAttribute(const OpKernelInfo& info) { + // legacy attribute that affects default axes value + const bool across_channels = info.GetAttrOrDefault("across_channels", int64_t{0}) == int64_t{1}; + + const auto default_axes = across_channels + ? std::vector{0, 1, 2, 3} + : std::vector{0, 2, 3}; + + const auto axes = info.GetAttrsOrDefault("axes", default_axes); + + return InlinedVector(axes.begin(), axes.end()); +} + +// Drop out of range, make non-negative, sort, and make unique. +InlinedVector NormalizeAxes(gsl::span axes, size_t rank) { + InlinedVector normalized_axes{}; + normalized_axes.reserve(axes.size()); + for (int64_t axis : axes) { + if (IsAxisInRange(axis, static_cast(rank))) { + normalized_axes.push_back( + static_cast(HandleNegativeAxis(axis, static_cast(rank)))); + } + } + std::sort(normalized_axes.begin(), normalized_axes.end()); + normalized_axes.erase(std::unique(normalized_axes.begin(), normalized_axes.end()), normalized_axes.end()); + return normalized_axes; +} + +std::optional> GetTransposePermutationIfNeeded(gsl::span normalized_axes, size_t rank) { + // We need to transpose if anything other than the trailing axes are specified. + // Assume `normalized_axes` is sorted with unique values. + const bool is_transpose_needed = [&]() { + for (size_t i = 0, num_axes = normalized_axes.size(); i < num_axes; ++i) { + if (normalized_axes[i] != rank - num_axes + i) { + return true; + } + } + return false; + }(); + + if (!is_transpose_needed) { + return std::nullopt; + } + + // permutation of [ { unspecified axes }, { specified axes } ] + InlinedVector permutation{}; + permutation.reserve(rank); + + auto specified_axis_it = normalized_axes.begin(); + for (size_t axis = 0; axis < rank; ++axis) { + if (specified_axis_it != normalized_axes.end() && + axis == *specified_axis_it) { + // skip specified axis for now, add them all to the end later + ++specified_axis_it; + } else { + // add unspecified axis + permutation.push_back(axis); + } + } + + // add all specified axes + permutation.insert(permutation.end(), normalized_axes.begin(), normalized_axes.end()); + + return permutation; +} + +// Given an M x N array where N is the inner dimension, compute normalized quantities for M sets of N values. +// X is the input and Y is the output. +Status ComputeMeanVarianceNormalization2D(size_t M, size_t N, + gsl::span X, gsl::span Y, + bool normalize_variance) { + ORT_RETURN_IF_NOT(X.size() == M * N && X.size() == Y.size(), "X and Y must both have M * N elements."); + + const auto idx_M = narrow(M), idx_N = narrow(N); + // Note: Eigen arrays have column-major storage by default, so we specify N rows x M columns. + ConstEigenArrayMap X_array(X.data(), idx_N, idx_M); + EigenArrayMap Y_array(Y.data(), idx_N, idx_M); + + // for each column, compute Y = X - E[X] + Y_array = X_array.rowwise() - X_array.colwise().mean(); + + if (normalize_variance) { + // for each column, compute Y' = (X - E[X]) / ( E[ (X - E[X])^2 ] )^(1/2) + // we start with Y = X - E[X], + // so Y' = Y / ( E[ Y^2 ] )^(1/2) + Y_array = (Y_array.rowwise() / (Y_array.square().colwise().mean().sqrt())).eval(); + } + + return Status::OK(); +} + +InlinedVector InvertPerm(gsl::span perm) { + InlinedVector inverted_perm{}; + inverted_perm.resize(perm.size()); + for (size_t i = 0; i < perm.size(); ++i) { + inverted_perm[perm[i]] = i; + } + return inverted_perm; +} +} // namespace + +MeanVarianceNormalization::MeanVarianceNormalization(const OpKernelInfo& info) + : OpKernel{info}, + normalize_variance_{ + // legacy attribute + info.GetAttrOrDefault("normalize_variance", int64_t{1}) == int64_t{1}}, + axes_{GetAxesFromAttribute(info)} { +} + +Status MeanVarianceNormalization::Compute(OpKernelContext* context) const { + const auto& input = context->RequiredInput(0); + const auto& input_shape = input.Shape(); + + Tensor& output = context->RequiredOutput(0, input_shape); + + // approach for normalizing values across arbitrary dimensions: + // - transpose to [unspecified axes, specified axes] + // - do normalization of inner dimension values + // - transpose back + + const auto rank = input_shape.GetDims().size(); + + const auto normalized_axes = NormalizeAxes(axes_, rank); + + // The ONNX spec doesn't specify what to do if no axes are specified so we won't try to do anything. + ORT_RETURN_IF(normalized_axes.empty(), "No valid axes are specified. This is not handled now."); + + if (input_shape.Size() == 0) { + return Status::OK(); + } + + const auto transpose_permutation = GetTransposePermutationIfNeeded(normalized_axes, rank); + const bool is_transpose_required = transpose_permutation.has_value(); + + // intermediate tensors if transposing is necessary + Tensor transposed_input; + Tensor transposed_result; + + TensorShape compute_shape = input_shape; + + if (is_transpose_required) { + AllocatorPtr alloc; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); + + InlinedVector transposed_dims{}; + transposed_dims.reserve(rank); + std::transform(transpose_permutation->begin(), transpose_permutation->end(), std::back_inserter(transposed_dims), + [input_dims = input.Shape().GetDims()](size_t axis) { return input_dims[axis]; }); + compute_shape = TensorShape(transposed_dims); + + transposed_input = Tensor{input.DataType(), compute_shape, alloc}; + + ORT_RETURN_IF_ERROR(TransposeBase::DoTranspose(*transpose_permutation, input, transposed_input)); + + transposed_result = Tensor{input.DataType(), compute_shape, alloc}; + } + + const size_t num_unspecified_axes = rank - normalized_axes.size(); + const size_t M = narrow(compute_shape.SizeToDimension(num_unspecified_axes)), + N = narrow(compute_shape.SizeFromDimension(num_unspecified_axes)); + + const gsl::span X = is_transpose_required + ? transposed_input.DataAsSpan() + : input.DataAsSpan(); + const gsl::span Y = is_transpose_required + ? transposed_result.MutableDataAsSpan() + : output.MutableDataAsSpan(); + + ORT_RETURN_IF_ERROR(ComputeMeanVarianceNormalization2D(M, N, X, Y, normalize_variance_)); + + if (is_transpose_required) { + const auto inverted_permutation = InvertPerm(*transpose_permutation); + ORT_RETURN_IF_ERROR(TransposeBase::DoTranspose(inverted_permutation, transposed_result, output)); + } + + return Status::OK(); +} + ONNX_CPU_OPERATOR_VERSIONED_KERNEL( MeanVarianceNormalization, 9, 12, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - MeanVarianceNormalization_1); + MeanVarianceNormalization); ONNX_CPU_OPERATOR_KERNEL( MeanVarianceNormalization, 13, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), - MeanVarianceNormalization_1); + MeanVarianceNormalization); + } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h b/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h index 89c9bccb07..2d0f3d9688 100644 --- a/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h +++ b/onnxruntime/core/providers/cpu/tensor/mean_variance_normalization.h @@ -2,121 +2,18 @@ // Licensed under the MIT License. #pragma once -#include -#include "core/common/common.h" + +#include "core/common/inlined_containers.h" #include "core/framework/op_kernel.h" -#include "core/util/math_cpuonly.h" -#include "core/common/gsl.h" namespace onnxruntime { -template -class MeanVarianceNormalization_0 : public OpKernel { +class MeanVarianceNormalization : public OpKernel { public: - MeanVarianceNormalization_0(const OpKernelInfo& info, bool old_attr = true) : OpKernel(info) { - if (old_attr) { - ORT_ENFORCE(info.GetAttr("across_channels", &across_channels_).IsOK()); - ORT_ENFORCE(info.GetAttr("normalize_variance", &normalize_variance_).IsOK()); - } - } + MeanVarianceNormalization(const OpKernelInfo& info); + Status Compute(OpKernelContext* context) const override; - Status Compute(OpKernelContext* context) const override { - const auto* X = context->Input(0); - if (X == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch"); - - const auto dims = X->Shape().GetDims(); - - if (dims.size() < 4) { - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, - "Input is expected to have four dimensions corresponding to [N,C,H,W]"); - } - - const int64_t N = dims[0]; - const int64_t C = dims[1]; - const int64_t H = dims[2]; - const int64_t W = dims[3]; - - Tensor* Y = context->Output(0, {N, C, H, W}); - const T* Xdata = X->Data(); - T* Ydata = Y->MutableData(); - - const int64_t sample_size = H * W; - Eigen::Array mean(C, 1); - Eigen::Array var(C, 1); - mean.setZero(); - var.setZero(); - - ConstEigenArrayMap X_arr(Xdata, onnxruntime::narrow(sample_size), SafeInt(N) * C); - for (int nc = 0; nc < N * C; ++nc) { - mean(nc % C) += X_arr.col(nc).sum(); - } - mean /= gsl::narrow_cast(N * sample_size); - for (int64_t nc = 0; nc < N * C; ++nc) { - var(onnxruntime::narrow(nc % C)) += (X_arr.col(onnxruntime::narrow(nc)) - mean(onnxruntime::narrow(nc % C))).matrix().squaredNorm(); - } - var /= gsl::narrow_cast(N * sample_size); - - Eigen::Array inv_std; - EigenArrayMap Y_arr(Ydata, onnxruntime::narrow(sample_size), SafeInt(N) * C); - - if (across_channels_) { - // m_c = sum(m_i) / n - float global_mean = mean.mean(); - - // var_c = [(var_1 + (m_1 - m_c)^2) + ... + (var_n + (m_n - m_c)^2)] / n - // = [sum(var_i) + squared_norm(m_i - m_c)] / n - float global_var = ((mean - global_mean).matrix().squaredNorm() + var.sum()) / C; - - // For across channels we can directly use eigen because global_mean and global_var - // are just scalars. - if (!normalize_variance_) { - Y_arr = X_arr - global_mean; - } else { - float inv_std_scalar = 1 / std::sqrt(global_var); - Y_arr = (X_arr - global_mean) * inv_std_scalar; - } - } else { - if (!normalize_variance_) { - // inv_std = 1 - for (int64_t nc = 0; nc < N * C; ++nc) { - // y = (x - mean) - Y_arr.col(onnxruntime::narrow(nc)) = (X_arr.col(onnxruntime::narrow(nc)) - mean(onnxruntime::narrow(nc % C))); - } - } else { - inv_std = var.sqrt().inverse(); - for (int64_t nc = 0; nc < N * C; ++nc) { - // y = (x - mean) * (inv_std) - Y_arr.col(onnxruntime::narrow(nc)) = (X_arr.col(onnxruntime::narrow(nc)) - mean(onnxruntime::narrow(nc % C))) * inv_std(onnxruntime::narrow(nc % C)); - } - } - } - return Status::OK(); - } - - protected: - int64_t across_channels_; - int64_t normalize_variance_; + private: + const bool normalize_variance_; + const InlinedVector axes_; }; - -template -class MeanVarianceNormalization_1 final : public MeanVarianceNormalization_0 { - public: - MeanVarianceNormalization_1(const OpKernelInfo& info) : MeanVarianceNormalization_0(info, false) { - std::vector axes; - if (!info.GetAttrs("axes", axes).IsOK()) { - axes = {0, 2, 3}; - } - constexpr int64_t cross_channel_axes[] = {0, 1, 2, 3}; - constexpr int64_t batch_spatial_axes[] = {0, 2, 3}; - - if (std::equal(std::begin(axes), std::end(axes), std::begin(cross_channel_axes), std::end(cross_channel_axes))) { - this->across_channels_ = true; - } else if (std::equal(std::begin(axes), std::end(axes), std::begin(batch_spatial_axes), std::end(batch_spatial_axes))) { - this->across_channels_ = false; - } else { - ORT_THROW("MeanVarianceNormalization CPU EP only supports NHW and NCHW reduction for axes attribute."); - } - this->normalize_variance_ = 1; - } -}; - } // namespace onnxruntime diff --git a/onnxruntime/test/common/random_generator_test.cc b/onnxruntime/test/common/random_generator_test.cc index f9cc66ba69..98f19188d6 100644 --- a/onnxruntime/test/common/random_generator_test.cc +++ b/onnxruntime/test/common/random_generator_test.cc @@ -11,7 +11,7 @@ namespace test { TEST(TensorGenerator, DiscreteFloat) { FixedPatternValueGenerator random{}; const std::vector shape = {2, 3}; - std::vector data = random.Discrete(shape, {-1.f, 0.f, 1.f}); + std::vector data = random.Discrete(shape, AsSpan({-1.f, 0.f, 1.f})); ASSERT_EQ(data.size(), static_cast(6)); for (float value : data) { @@ -22,7 +22,7 @@ TEST(TensorGenerator, DiscreteFloat) { TEST(TensorGenerator, DiscreteInt) { FixedPatternValueGenerator random{}; const std::vector shape = {2, 3}; - std::vector data = random.Discrete(shape, {-1, 0, 1}); + std::vector data = random.Discrete(shape, AsSpan({-1, 0, 1})); ASSERT_EQ(data.size(), static_cast(6)); for (int value : data) { @@ -34,7 +34,7 @@ TEST(TensorGenerator, DiscreteInt) { TEST(TensorGenerator, CircularFloat) { FixedPatternValueGenerator random{}; const std::vector shape = {3, 2}; - std::vector data = random.Circular(shape, {-1.f, 0.f, 1.f}); + std::vector data = random.Circular(shape, AsSpan({-1.f, 0.f, 1.f})); ASSERT_EQ(data.size(), static_cast(6)); EXPECT_EQ(data[0], -1.f); @@ -48,7 +48,7 @@ TEST(TensorGenerator, CircularFloat) { TEST(TensorGenerator, CircularInt) { FixedPatternValueGenerator random{}; const std::vector shape = {3, 2}; - std::vector data = random.Circular(shape, {-1, 0, 1}); + std::vector data = random.Circular(shape, AsSpan({-1, 0, 1})); ASSERT_EQ(data.size(), static_cast(6)); EXPECT_EQ(data[0], -1); @@ -62,7 +62,7 @@ TEST(TensorGenerator, CircularInt) { TEST(TensorGenerator, CircularBool) { FixedPatternValueGenerator random{}; const std::vector shape = {3, 2}; - std::vector data = random.Circular(shape, {false, true}); + std::vector data = random.Circular(shape, AsSpan({false, true})); ASSERT_EQ(data.size(), static_cast(6)); EXPECT_EQ(data[0], false); diff --git a/onnxruntime/test/common/tensor_op_test_utils.h b/onnxruntime/test/common/tensor_op_test_utils.h index 3a509b02c0..a6369ed0e6 100644 --- a/onnxruntime/test/common/tensor_op_test_utils.h +++ b/onnxruntime/test/common/tensor_op_test_utils.h @@ -29,7 +29,7 @@ class FixedPatternValueGenerator { template std::vector - Discrete(gsl::span dims, const std::vector& value_candidates) { + Discrete(gsl::span dims, gsl::span value_candidates) { std::vector values(detail::SizeFromDims(dims)); std::uniform_int_distribution distribution(0, value_candidates.size() - 1); // Tier 2 RNG. Use it if `RandomValueGenerator::Uniform` method causes large numerical errors @@ -61,7 +61,7 @@ class FixedPatternValueGenerator { template std::vector - Circular(gsl::span dims, const std::vector& value_candidates) { + Circular(gsl::span dims, gsl::span value_candidates) { // Tier 3 RNG. Use it if `Discrete` method causes large numerical errors // (e.g., when elementwise relative error > 1e-3). // Suggested value_candidates to alleviate numerical error (listed @@ -118,7 +118,7 @@ inline std::vector ValueRange(size_t count, MLFloat16 star return result; } -inline std::pair MeanStdev(std::vector& v) { +inline std::pair MeanStdev(gsl::span v) { float sum = std::accumulate(v.begin(), v.end(), 0.0f); float mean = sum / v.size(); @@ -132,7 +132,7 @@ inline std::pair MeanStdev(std::vector& v) { } inline void Normalize(std::vector& v, - std::pair& mean_stdev, bool normalize_variance) { + const std::pair& mean_stdev, bool normalize_variance) { float mean = mean_stdev.first; float stdev = mean_stdev.second; diff --git a/onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py b/onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py new file mode 100644 index 0000000000..a7c1be435d --- /dev/null +++ b/onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py @@ -0,0 +1,37 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import argparse + +import numpy as np + +parser = argparse.ArgumentParser() +parser.add_argument("--shape", type=int, nargs="+", required=True) +parser.add_argument("--axes", type=int, nargs="+", required=True) +args = parser.parse_args() + +shape = tuple(args.shape) +axes = tuple(args.axes) + +random_seed = 0 +rng = np.random.default_rng(random_seed) + +X = rng.random(size=shape, dtype=float) + +# Calculate expected output data +X_mean = np.mean(X, axis=axes, keepdims=True) +X_std = np.std(X, axis=axes, keepdims=True) +Y = (X - X_mean) / X_std + + +def to_c_float_literals(arr): + literals_per_line = 8 + literals = [f"{literal:.7f}f" for literal in arr.flatten().tolist()] + result = "" + for i, literal in enumerate(literals): + result += "{},{}".format(literal, "\n" if (i + 1) % literals_per_line == 0 else " ") + return result + + +print(f"input:\n{to_c_float_literals(X)}") +print(f"expected output:\n{to_c_float_literals(Y)}") diff --git a/onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc b/onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc new file mode 100644 index 0000000000..b6720ae2a9 --- /dev/null +++ b/onnxruntime/test/providers/cpu/tensor/mean_variance_normalization_test.cc @@ -0,0 +1,273 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" + +#include "test/common/tensor_op_test_utils.h" +#include "test/providers/provider_test_utils.h" + +namespace onnxruntime::test { + +TEST(MeanVarianceNormalizationTest, DefaultAxes) { + constexpr int64_t N = 2, C = 2, H = 2, W = 3; + + std::vector N1C1 = {3.0f, -3.0f, -1.0f, + 1.0f, 2.0f, -1.0f}; + std::vector N1C2 = {-2.0f, -2.0f, -2.0f, + 4.0f, 1.0f, 4.0f}; + std::vector N2C1 = { + 0.0f, + -2.0f, + -2.0f, + -4.0f, + 5.0f, + 7.0f, + }; + std::vector N2C2 = { + 5.0f, + -5.0f, + -5.0f, + 3.0f, + 4.0f, + 4.0f, + }; + + std::vector X; + X.reserve(N * C * H * W); + X.insert(X.end(), N1C1.begin(), N1C1.end()); + X.insert(X.end(), N1C2.begin(), N1C2.end()); + X.insert(X.end(), N2C1.begin(), N2C1.end()); + X.insert(X.end(), N2C2.begin(), N2C2.end()); + + std::vector C1; + C1.reserve(N * H * W); + C1.insert(C1.end(), N1C1.begin(), N1C1.end()); + C1.insert(C1.end(), N2C1.begin(), N2C1.end()); + auto C1_meam_stdev = MeanStdev(C1); + + std::vector C2; + C2.reserve(N * H * W); + C2.insert(C2.end(), N1C2.begin(), N1C2.end()); + C2.insert(C2.end(), N2C2.begin(), N2C2.end()); + auto C2_meam_stdev = MeanStdev(C2); + + std::vector N1C1_result(N1C1), N1C2_result(N1C2), + N2C1_result(N2C1), N2C2_result(N2C2); + Normalize(N1C1_result, C1_meam_stdev, true); + Normalize(N2C1_result, C1_meam_stdev, true); + Normalize(N1C2_result, C2_meam_stdev, true); + Normalize(N2C2_result, C2_meam_stdev, true); + + std::vector result; + result.reserve(N * C * H * W); + result.insert(result.end(), N1C1_result.begin(), N1C1_result.end()); + result.insert(result.end(), N1C2_result.begin(), N1C2_result.end()); + result.insert(result.end(), N2C1_result.begin(), N2C1_result.end()); + result.insert(result.end(), N2C2_result.begin(), N2C2_result.end()); + + OpTester test("MeanVarianceNormalization", 9); + test.AddInput("input", {N, C, H, W}, X); + test.AddOutput("output", {N, C, H, W}, result); + test.Run(); +} + +static void TestMeanVarianceNormalizationOverAllAxes(const std::vector& shape) { + SCOPED_TRACE(MakeString("shape: ", TensorShape(shape))); + + FixedPatternValueGenerator generator{}; + + const auto X_value_candidates = ValueRange(11, -5.0f); + const auto X = generator.Discrete(shape, X_value_candidates); + const auto mean_stdev = MeanStdev(X); + + std::vector Y(X); + Normalize(Y, mean_stdev, true); + + OpTester test("MeanVarianceNormalization", 9); + const auto all_axes = ValueRange(shape.size()); + test.AddAttribute("axes", all_axes); + test.AddInput("input", shape, X); + test.AddOutput("output", shape, Y); + + test.Run(); +} + +TEST(MeanVarianceNormalizationTest, AllAxes) { + TestMeanVarianceNormalizationOverAllAxes({2, 2, 4}); + TestMeanVarianceNormalizationOverAllAxes({2, 2, 2, 3}); + TestMeanVarianceNormalizationOverAllAxes({2, 2, 2, 2, 2}); +} + +TEST(MeanVarianceNormalizationTest, AxesSubsets5D) { + // test data was generated with: + // python onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py --shape 2 2 2 2 2 --axes + + auto axes_to_str = [](gsl::span axes) { + std::ostringstream s; + s << "{ "; + std::copy(axes.begin(), axes.end(), std::ostream_iterator(s, " ")); + s << "}"; + return s.str(); + }; + + auto test_with_axes = [&](gsl::span axes, gsl::span Y) { + SCOPED_TRACE(axes_to_str(axes)); + + constexpr std::array X = { + 0.6369617f, + 0.2697867f, + 0.0409735f, + 0.0165276f, + 0.8132702f, + 0.9127556f, + 0.6066358f, + 0.7294966f, + 0.5436250f, + 0.9350724f, + 0.8158536f, + 0.0027385f, + 0.8574043f, + 0.0335856f, + 0.7296554f, + 0.1756556f, + 0.8631789f, + 0.5414612f, + 0.2997119f, + 0.4226872f, + 0.0283197f, + 0.1242833f, + 0.6706244f, + 0.6471895f, + 0.6153851f, + 0.3836776f, + 0.9972099f, + 0.9808353f, + 0.6855420f, + 0.6504593f, + 0.6884467f, + 0.3889214f, + }; + + const std::vector shape{2, 2, 2, 2, 2}; + + OpTester test("MeanVarianceNormalization", 9); + test.AddAttribute("axes", axes); + test.AddInput("input", shape, X.data(), X.size()); + test.AddOutput("output", shape, Y.data(), Y.size()); + + test.Run(); + }; + + test_with_axes( + AsSpan({0, 2, 4}), + AsSpan({ + 0.3508345f, + -0.7870349f, + -1.4605863f, + -1.5525494f, + 0.8972119f, + 1.2055154f, + 0.6673803f, + 1.1295706f, + -0.1683330f, + 1.3134559f, + 0.6321192f, + -1.7208749f, + 1.0194501f, + -2.0990413f, + 0.3826789f, + -1.2204870f, + 1.0518781f, + 0.0548801f, + -0.4872377f, + -0.0246164f, + -1.5353374f, + -1.2379477f, + 0.9080993f, + 0.8199395f, + 0.1033084f, + -0.7737996f, + 1.1569287f, + 1.1095439f, + 0.3688809f, + 0.2360785f, + 0.2634291f, + -0.6033379f, + })); + + test_with_axes( + AsSpan({1, 2, 3}), + AsSpan({ + 0.0260567f, + -0.3008327f, + -2.3950341f, + -0.9652744f, + 0.7422773f, + 1.3860379f, + -0.0971367f, + 0.9052460f, + -0.3531062f, + 1.4445876f, + 0.7527716f, + -1.0014510f, + 0.9215636f, + -0.9205217f, + 0.4026078f, + -0.5477917f, + 0.8924309f, + 0.1015335f, + -1.0632416f, + -0.4004898f, + -2.0051854f, + -1.6617567f, + 0.2241158f, + 0.5484163f, + 0.0323921f, + -0.5653723f, + 1.3576236f, + 1.9586404f, + 0.2758914f, + 0.5622366f, + 0.2859732f, + -0.5432080f, + })); + + test_with_axes( + AsSpan({0, 1, 4}), + AsSpan({ + 0.1843672f, + -1.5822912f, + -1.0098907f, + -1.0706838f, + 0.8350127f, + 1.1118552f, + 0.1472972f, + 0.8161319f, + -0.2647213f, + 1.6187237f, + 0.9171133f, + -1.1049752f, + 0.9578265f, + -1.3346525f, + 0.8169968f, + -2.1988905f, + 1.2728089f, + -0.2751323f, + -0.3664493f, + -0.0606291f, + -1.3493062f, + -1.0822638f, + 0.4956413f, + 0.3680653f, + 0.0805517f, + -1.0343067f, + 1.3681179f, + 1.3273969f, + 0.4795772f, + 0.3819509f, + 0.5926631f, + -1.0379052f, + })); +} + +} // namespace onnxruntime::test diff --git a/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc b/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc index 987ded80d1..cd83a3dbd2 100644 --- a/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/tensor_op_test.cc @@ -426,99 +426,5 @@ TEST(TensorOpTest, ShapeTest3D) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); } -void MeanVarianceNormalizationFunctionDefaultPerChannel() { - constexpr int64_t N = 2, C = 2, H = 2, W = 3; - - std::vector N1C1 = {3.0f, -3.0f, -1.0f, - 1.0f, 2.0f, -1.0f}; - std::vector N1C2 = {-2.0f, -2.0f, -2.0f, - 4.0f, 1.0f, 4.0f}; - std::vector N2C1 = { - 0.0f, - -2.0f, - -2.0f, - -4.0f, - 5.0f, - 7.0f, - }; - std::vector N2C2 = { - 5.0f, - -5.0f, - -5.0f, - 3.0f, - 4.0f, - 4.0f, - }; - - std::vector X; - X.reserve(N * C * H * W); - X.insert(X.end(), N1C1.begin(), N1C1.end()); - X.insert(X.end(), N1C2.begin(), N1C2.end()); - X.insert(X.end(), N2C1.begin(), N2C1.end()); - X.insert(X.end(), N2C2.begin(), N2C2.end()); - - std::vector C1; - C1.reserve(N * H * W); - C1.insert(C1.end(), N1C1.begin(), N1C1.end()); - C1.insert(C1.end(), N2C1.begin(), N2C1.end()); - auto C1_meam_stdev = MeanStdev(C1); - - std::vector C2; - C2.reserve(N * H * W); - C2.insert(C2.end(), N1C2.begin(), N1C2.end()); - C2.insert(C2.end(), N2C2.begin(), N2C2.end()); - auto C2_meam_stdev = MeanStdev(C2); - - std::vector N1C1_result(N1C1), N1C2_result(N1C2), - N2C1_result(N2C1), N2C2_result(N2C2); - Normalize(N1C1_result, C1_meam_stdev, 1); - Normalize(N2C1_result, C1_meam_stdev, 1); - Normalize(N1C2_result, C2_meam_stdev, 1); - Normalize(N2C2_result, C2_meam_stdev, 1); - - std::vector result; - result.reserve(N * C * H * W); - result.insert(result.end(), N1C1_result.begin(), N1C1_result.end()); - result.insert(result.end(), N1C2_result.begin(), N1C2_result.end()); - result.insert(result.end(), N2C1_result.begin(), N2C1_result.end()); - result.insert(result.end(), N2C2_result.begin(), N2C2_result.end()); - - OpTester test("MeanVarianceNormalization", 9); - test.AddInput("input", {N, C, H, W}, X); - test.AddOutput("output", {N, C, H, W}, result); - test.Run(); -} - -void MeanVarianceNormalizationFunctionAcrossChannels(std::vector axes) { - constexpr int64_t N = 2, C = 2, H = 2, W = 3; - - std::vector X = {3.0f, -3.0f, -1.0f, - 1.0f, 2.0f, -1.0f, - -2.0f, -2.0f, -2.0f, - 4.0f, 1.0f, 4.0f, - 0.0f, -2.0f, -2.0f, - -4.0f, 5.0f, 7.0f, - 5.0f, -5.0f, -5.0f, - 3.0f, 4.0f, 4.0f}; - auto mean_stdev = MeanStdev(X); - - std::vector result(X); - Normalize(result, mean_stdev, 1); - - OpTester test("MeanVarianceNormalization", 9); - test.AddAttribute("axes", axes); - test.AddInput("input", {N, C, H, W}, X); - test.AddOutput("output", {N, C, H, W}, result); - test.Run(); -} - -TEST(TensorOpTest, MeanVarianceNormalizationCPUTest) { - // axes: {0, 1, 2, 3} for across_channels - MeanVarianceNormalizationFunctionAcrossChannels({0, 1, 2, 3}); - - // Default (axes: {0, 2, 3}) for non across_channels - MeanVarianceNormalizationFunctionDefaultPerChannel(); -} - } // namespace test } // namespace onnxruntime