mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Rework CPU MeanVarianceNormalization kernel to support arbitrary axes. (#16420)
This commit is contained in:
Edward Chen
GitHub
parent
e2c214d81f
commit
4a331ef667
9 changed files with 527 additions and 218 deletions
|
|
@ -13,6 +13,6 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
|
|||
1,
|
||||
8,
|
||||
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
|
||||
MeanVarianceNormalization_0<float>);
|
||||
MeanVarianceNormalization);
|
||||
} // namespace contrib
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -3,17 +3,206 @@
|
|||
|
||||
#include "core/providers/cpu/tensor/mean_variance_normalization.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <optional>
|
||||
|
||||
#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<int64_t> GetAxesFromAttribute(const OpKernelInfo& info) {
|
||||
// legacy attribute that affects default axes value
|
||||
const bool across_channels = info.GetAttrOrDefault<int64_t>("across_channels", int64_t{0}) == int64_t{1};
|
||||
|
||||
const auto default_axes = across_channels
|
||||
? std::vector<int64_t>{0, 1, 2, 3}
|
||||
: std::vector<int64_t>{0, 2, 3};
|
||||
|
||||
const auto axes = info.GetAttrsOrDefault<int64_t>("axes", default_axes);
|
||||
|
||||
return InlinedVector<int64_t>(axes.begin(), axes.end());
|
||||
}
|
||||
|
||||
// Drop out of range, make non-negative, sort, and make unique.
|
||||
InlinedVector<size_t> NormalizeAxes(gsl::span<const int64_t> axes, size_t rank) {
|
||||
InlinedVector<size_t> normalized_axes{};
|
||||
normalized_axes.reserve(axes.size());
|
||||
for (int64_t axis : axes) {
|
||||
if (IsAxisInRange(axis, static_cast<int64_t>(rank))) {
|
||||
normalized_axes.push_back(
|
||||
static_cast<size_t>(HandleNegativeAxis(axis, static_cast<int64_t>(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<InlinedVector<size_t>> GetTransposePermutationIfNeeded(gsl::span<const size_t> 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<size_t> 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<const float> X, gsl::span<float> 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<Eigen::Index>(M), idx_N = narrow<Eigen::Index>(N);
|
||||
// Note: Eigen arrays have column-major storage by default, so we specify N rows x M columns.
|
||||
ConstEigenArrayMap<float> X_array(X.data(), idx_N, idx_M);
|
||||
EigenArrayMap<float> 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<size_t> InvertPerm(gsl::span<const size_t> perm) {
|
||||
InlinedVector<size_t> 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<Tensor>(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<int64_t> 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<size_t>(compute_shape.SizeToDimension(num_unspecified_axes)),
|
||||
N = narrow<size_t>(compute_shape.SizeFromDimension(num_unspecified_axes));
|
||||
|
||||
const gsl::span<const float> X = is_transpose_required
|
||||
? transposed_input.DataAsSpan<float>()
|
||||
: input.DataAsSpan<float>();
|
||||
const gsl::span<float> Y = is_transpose_required
|
||||
? transposed_result.MutableDataAsSpan<float>()
|
||||
: output.MutableDataAsSpan<float>();
|
||||
|
||||
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<float>()),
|
||||
MeanVarianceNormalization_1<float>);
|
||||
MeanVarianceNormalization);
|
||||
|
||||
ONNX_CPU_OPERATOR_KERNEL(
|
||||
MeanVarianceNormalization,
|
||||
13,
|
||||
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
|
||||
MeanVarianceNormalization_1<float>);
|
||||
MeanVarianceNormalization);
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
|
|
@ -2,121 +2,18 @@
|
|||
// Licensed under the MIT License.
|
||||
|
||||
#pragma once
|
||||
#include <core/common/safeint.h>
|
||||
#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 <typename T>
|
||||
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<int64_t>("across_channels", &across_channels_).IsOK());
|
||||
ORT_ENFORCE(info.GetAttr<int64_t>("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<Tensor>(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>();
|
||||
T* Ydata = Y->MutableData<T>();
|
||||
|
||||
const int64_t sample_size = H * W;
|
||||
Eigen::Array<float, Eigen::Dynamic, 1> mean(C, 1);
|
||||
Eigen::Array<float, Eigen::Dynamic, 1> var(C, 1);
|
||||
mean.setZero();
|
||||
var.setZero();
|
||||
|
||||
ConstEigenArrayMap<T> X_arr(Xdata, onnxruntime::narrow<std::ptrdiff_t>(sample_size), SafeInt<ptrdiff_t>(N) * C);
|
||||
for (int nc = 0; nc < N * C; ++nc) {
|
||||
mean(nc % C) += X_arr.col(nc).sum();
|
||||
}
|
||||
mean /= gsl::narrow_cast<T>(N * sample_size);
|
||||
for (int64_t nc = 0; nc < N * C; ++nc) {
|
||||
var(onnxruntime::narrow<std::ptrdiff_t>(nc % C)) += (X_arr.col(onnxruntime::narrow<std::ptrdiff_t>(nc)) - mean(onnxruntime::narrow<std::ptrdiff_t>(nc % C))).matrix().squaredNorm();
|
||||
}
|
||||
var /= gsl::narrow_cast<T>(N * sample_size);
|
||||
|
||||
Eigen::Array<T, Eigen::Dynamic, 1> inv_std;
|
||||
EigenArrayMap<T> Y_arr(Ydata, onnxruntime::narrow<std::ptrdiff_t>(sample_size), SafeInt<ptrdiff_t>(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<std::ptrdiff_t>(nc)) = (X_arr.col(onnxruntime::narrow<std::ptrdiff_t>(nc)) - mean(onnxruntime::narrow<std::ptrdiff_t>(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<std::ptrdiff_t>(nc)) = (X_arr.col(onnxruntime::narrow<std::ptrdiff_t>(nc)) - mean(onnxruntime::narrow<std::ptrdiff_t>(nc % C))) * inv_std(onnxruntime::narrow<std::ptrdiff_t>(nc % C));
|
||||
}
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
protected:
|
||||
int64_t across_channels_;
|
||||
int64_t normalize_variance_;
|
||||
private:
|
||||
const bool normalize_variance_;
|
||||
const InlinedVector<int64_t> axes_;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class MeanVarianceNormalization_1 final : public MeanVarianceNormalization_0<T> {
|
||||
public:
|
||||
MeanVarianceNormalization_1(const OpKernelInfo& info) : MeanVarianceNormalization_0<T>(info, false) {
|
||||
std::vector<int64_t> 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
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ namespace test {
|
|||
TEST(TensorGenerator, DiscreteFloat) {
|
||||
FixedPatternValueGenerator random{};
|
||||
const std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<float> data = random.Discrete<float>(shape, {-1.f, 0.f, 1.f});
|
||||
std::vector<float> data = random.Discrete<float>(shape, AsSpan({-1.f, 0.f, 1.f}));
|
||||
|
||||
ASSERT_EQ(data.size(), static_cast<size_t>(6));
|
||||
for (float value : data) {
|
||||
|
|
@ -22,7 +22,7 @@ TEST(TensorGenerator, DiscreteFloat) {
|
|||
TEST(TensorGenerator, DiscreteInt) {
|
||||
FixedPatternValueGenerator random{};
|
||||
const std::vector<int64_t> shape = {2, 3};
|
||||
std::vector<int> data = random.Discrete<int>(shape, {-1, 0, 1});
|
||||
std::vector<int> data = random.Discrete<int>(shape, AsSpan({-1, 0, 1}));
|
||||
|
||||
ASSERT_EQ(data.size(), static_cast<size_t>(6));
|
||||
for (int value : data) {
|
||||
|
|
@ -34,7 +34,7 @@ TEST(TensorGenerator, DiscreteInt) {
|
|||
TEST(TensorGenerator, CircularFloat) {
|
||||
FixedPatternValueGenerator random{};
|
||||
const std::vector<int64_t> shape = {3, 2};
|
||||
std::vector<float> data = random.Circular<float>(shape, {-1.f, 0.f, 1.f});
|
||||
std::vector<float> data = random.Circular<float>(shape, AsSpan({-1.f, 0.f, 1.f}));
|
||||
|
||||
ASSERT_EQ(data.size(), static_cast<size_t>(6));
|
||||
EXPECT_EQ(data[0], -1.f);
|
||||
|
|
@ -48,7 +48,7 @@ TEST(TensorGenerator, CircularFloat) {
|
|||
TEST(TensorGenerator, CircularInt) {
|
||||
FixedPatternValueGenerator random{};
|
||||
const std::vector<int64_t> shape = {3, 2};
|
||||
std::vector<int> data = random.Circular<int>(shape, {-1, 0, 1});
|
||||
std::vector<int> data = random.Circular<int>(shape, AsSpan({-1, 0, 1}));
|
||||
|
||||
ASSERT_EQ(data.size(), static_cast<size_t>(6));
|
||||
EXPECT_EQ(data[0], -1);
|
||||
|
|
@ -62,7 +62,7 @@ TEST(TensorGenerator, CircularInt) {
|
|||
TEST(TensorGenerator, CircularBool) {
|
||||
FixedPatternValueGenerator random{};
|
||||
const std::vector<int64_t> shape = {3, 2};
|
||||
std::vector<bool> data = random.Circular<bool>(shape, {false, true});
|
||||
std::vector<bool> data = random.Circular<bool>(shape, AsSpan({false, true}));
|
||||
|
||||
ASSERT_EQ(data.size(), static_cast<size_t>(6));
|
||||
EXPECT_EQ(data[0], false);
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ class FixedPatternValueGenerator {
|
|||
|
||||
template <typename TValue>
|
||||
std::vector<TValue>
|
||||
Discrete(gsl::span<const int64_t> dims, const std::vector<TValue>& value_candidates) {
|
||||
Discrete(gsl::span<const int64_t> dims, gsl::span<const TValue> value_candidates) {
|
||||
std::vector<TValue> values(detail::SizeFromDims(dims));
|
||||
std::uniform_int_distribution<size_t> 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 <typename TValue>
|
||||
std::vector<TValue>
|
||||
Circular(gsl::span<const int64_t> dims, const std::vector<TValue>& value_candidates) {
|
||||
Circular(gsl::span<const int64_t> dims, gsl::span<const TValue> 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<MLFloat16> ValueRange<MLFloat16>(size_t count, MLFloat16 star
|
|||
return result;
|
||||
}
|
||||
|
||||
inline std::pair<float, float> MeanStdev(std::vector<float>& v) {
|
||||
inline std::pair<float, float> MeanStdev(gsl::span<const float> v) {
|
||||
float sum = std::accumulate(v.begin(), v.end(), 0.0f);
|
||||
float mean = sum / v.size();
|
||||
|
||||
|
|
@ -132,7 +132,7 @@ inline std::pair<float, float> MeanStdev(std::vector<float>& v) {
|
|||
}
|
||||
|
||||
inline void Normalize(std::vector<float>& v,
|
||||
std::pair<float, float>& mean_stdev, bool normalize_variance) {
|
||||
const std::pair<float, float>& mean_stdev, bool normalize_variance) {
|
||||
float mean = mean_stdev.first;
|
||||
float stdev = mean_stdev.second;
|
||||
|
||||
|
|
|
|||
37
onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py
Normal file
37
onnxruntime/test/providers/cpu/tensor/gen_mvn_test_data.py
Normal file
|
|
@ -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)}")
|
||||
|
|
@ -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<float> N1C1 = {3.0f, -3.0f, -1.0f,
|
||||
1.0f, 2.0f, -1.0f};
|
||||
std::vector<float> N1C2 = {-2.0f, -2.0f, -2.0f,
|
||||
4.0f, 1.0f, 4.0f};
|
||||
std::vector<float> N2C1 = {
|
||||
0.0f,
|
||||
-2.0f,
|
||||
-2.0f,
|
||||
-4.0f,
|
||||
5.0f,
|
||||
7.0f,
|
||||
};
|
||||
std::vector<float> N2C2 = {
|
||||
5.0f,
|
||||
-5.0f,
|
||||
-5.0f,
|
||||
3.0f,
|
||||
4.0f,
|
||||
4.0f,
|
||||
};
|
||||
|
||||
std::vector<float> 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<float> 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<float> 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<float> 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<float> 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<float>("input", {N, C, H, W}, X);
|
||||
test.AddOutput<float>("output", {N, C, H, W}, result);
|
||||
test.Run();
|
||||
}
|
||||
|
||||
static void TestMeanVarianceNormalizationOverAllAxes(const std::vector<int64_t>& shape) {
|
||||
SCOPED_TRACE(MakeString("shape: ", TensorShape(shape)));
|
||||
|
||||
FixedPatternValueGenerator generator{};
|
||||
|
||||
const auto X_value_candidates = ValueRange<float>(11, -5.0f);
|
||||
const auto X = generator.Discrete<float>(shape, X_value_candidates);
|
||||
const auto mean_stdev = MeanStdev(X);
|
||||
|
||||
std::vector<float> Y(X);
|
||||
Normalize(Y, mean_stdev, true);
|
||||
|
||||
OpTester test("MeanVarianceNormalization", 9);
|
||||
const auto all_axes = ValueRange<int64_t>(shape.size());
|
||||
test.AddAttribute("axes", all_axes);
|
||||
test.AddInput<float>("input", shape, X);
|
||||
test.AddOutput<float>("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 <axes values>
|
||||
|
||||
auto axes_to_str = [](gsl::span<const int64_t> axes) {
|
||||
std::ostringstream s;
|
||||
s << "{ ";
|
||||
std::copy(axes.begin(), axes.end(), std::ostream_iterator<int64_t>(s, " "));
|
||||
s << "}";
|
||||
return s.str();
|
||||
};
|
||||
|
||||
auto test_with_axes = [&](gsl::span<const int64_t> axes, gsl::span<const float> 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<int64_t> shape{2, 2, 2, 2, 2};
|
||||
|
||||
OpTester test("MeanVarianceNormalization", 9);
|
||||
test.AddAttribute("axes", axes);
|
||||
test.AddInput<float>("input", shape, X.data(), X.size());
|
||||
test.AddOutput<float>("output", shape, Y.data(), Y.size());
|
||||
|
||||
test.Run();
|
||||
};
|
||||
|
||||
test_with_axes(
|
||||
AsSpan<int64_t>({0, 2, 4}),
|
||||
AsSpan<float>({
|
||||
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<int64_t>({1, 2, 3}),
|
||||
AsSpan<float>({
|
||||
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<int64_t>({0, 1, 4}),
|
||||
AsSpan<float>({
|
||||
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
|
||||
|
|
@ -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<float> N1C1 = {3.0f, -3.0f, -1.0f,
|
||||
1.0f, 2.0f, -1.0f};
|
||||
std::vector<float> N1C2 = {-2.0f, -2.0f, -2.0f,
|
||||
4.0f, 1.0f, 4.0f};
|
||||
std::vector<float> N2C1 = {
|
||||
0.0f,
|
||||
-2.0f,
|
||||
-2.0f,
|
||||
-4.0f,
|
||||
5.0f,
|
||||
7.0f,
|
||||
};
|
||||
std::vector<float> N2C2 = {
|
||||
5.0f,
|
||||
-5.0f,
|
||||
-5.0f,
|
||||
3.0f,
|
||||
4.0f,
|
||||
4.0f,
|
||||
};
|
||||
|
||||
std::vector<float> 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<float> 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<float> 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<float> 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<float> 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<float>("input", {N, C, H, W}, X);
|
||||
test.AddOutput<float>("output", {N, C, H, W}, result);
|
||||
test.Run();
|
||||
}
|
||||
|
||||
void MeanVarianceNormalizationFunctionAcrossChannels(std::vector<int64_t> axes) {
|
||||
constexpr int64_t N = 2, C = 2, H = 2, W = 3;
|
||||
|
||||
std::vector<float> 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<float> result(X);
|
||||
Normalize(result, mean_stdev, 1);
|
||||
|
||||
OpTester test("MeanVarianceNormalization", 9);
|
||||
test.AddAttribute("axes", axes);
|
||||
test.AddInput<float>("input", {N, C, H, W}, X);
|
||||
test.AddOutput<float>("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
|
||||
|
|
|
|||
Loading…
Reference in a new issue