saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-20 19:12:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support no bias in layernorm and skiplayernorm op (#6554)

Browse source 
* add noBias attribute in layernorm

* skip bias in skiplayernorm

* fix

* fix cuda tets

* add tests

* fix windows build

* fix win build issue

* review comments
This commit is contained in:
Ye Wang 2021-02-05 16:48:22 -08:00 committed by GitHub
parent 299ace0759
commit 82229c8e61
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 122 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/contrib_ops/cpu/layer_norm.cc
View file

@ -31,7 +31,7 @@ namespace contrib {
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
LayerNorm<T, true>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED(double)
@ -50,7 +50,7 @@ Status LayerNorm<T, simplified>::Compute(OpKernelContext* p_ctx) const {
const Tensor* bias = p_ctx->Input<Tensor>(2);
auto X_data = X->template Data<T>();
auto scale_data = scale->template Data<T>();
auto bias_data = simplified ? nullptr : bias->template Data<T>();
auto bias_data = (simplified || nullptr == bias) ? nullptr : bias->template Data<T>();
const TensorShape& x_shape = X->Shape();
const int64_t axis = HandleNegativeAxis(axis_, x_shape.NumDimensions());
@ -124,6 +124,8 @@ Status LayerNorm<T, simplified>::Compute(OpKernelContext* p_ctx) const {
for (int64_t h = 0; h < norm_size; h++) {
if (simplified) {
p_output[h] = p_input[h] / mean_square * scale_data[h];
} else if (nullptr == bias){
p_output[h] = (p_input[h] - mean) / mean_square * scale_data[h];
} else {
p_output[h] = (p_input[h] - mean) / mean_square * scale_data[h] + bias_data[h];
}

26
onnxruntime/contrib_ops/cpu/skip_layer_norm.cc
View file

@ -61,14 +61,16 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
"Last dimension of gamma and input does not match");
}
const auto& beta_dims = beta->Shape().GetDims();
if (beta_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have 1 dimension, got ", beta_dims.size());
}
if (beta_dims[0] != input_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Last dimension of beta and input does not match");
if (nullptr != beta) {
const auto& beta_dims = beta->Shape().GetDims();
if (beta_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have 1 dimension, got ", beta_dims.size());
}
if (beta_dims[0] != input_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Last dimension of beta and input does not match");
}
}
if (nullptr != bias) {
@ -91,7 +93,7 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
const T* input_data = input->Data<T>();
const T* skip_data = skip->Data<T>();
const T* gamma_data = gamma->Data<T>();
const T* beta_data = beta->Data<T>();
const T* beta_data = beta == nullptr ? nullptr : beta->Data<T>();
const T* bias_data = bias == nullptr ? nullptr : bias->Data<T>();
T* output_data = output->MutableData<T>();
@ -119,7 +121,11 @@ Status SkipLayerNorm<T>::Compute(OpKernelContext* p_ctx) const {
mean_square = sqrt(mean_square / hidden_size - mean * mean + epsilon_);
for (int64_t h = 0; h < hidden_size; h++) {
p_output[h] = (p_output[h] - mean) / mean_square * gamma_data[h] + beta_data[h];
if (nullptr == beta_data) {
p_output[h] = (p_output[h] - mean) / mean_square * gamma_data[h];
} else {
p_output[h] = (p_output[h] - mean) / mean_square * gamma_data[h] + beta_data[h];
}
}
}, 0);

8
onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh
View file

@ -1,7 +1,7 @@
/*
The implementation of this file is based on bert plugins in TensorRT demo:
https://github.com/NVIDIA/TensorRT/tree/release/5.1/demo/BERT/
Copyright 2019 NVIDIA Corporation
Licensed under the Apache License, Version 2.0 (the "License");
@ -79,7 +79,7 @@ struct KeyValuePairSum {
template <typename T, int TPB>
__device__ inline void LayerNorm(
const cub::KeyValuePair<T, T>& thread_data, const int ld, const int offset, const T* beta,
const cub::KeyValuePair<T, T>& thread_data, const int ld, const int offset, const T* beta,
const T* gamma, const T epsilon, T* output) {
// Assuming thread_data is already divided by ld
@ -101,7 +101,7 @@ __device__ inline void LayerNorm(
const int idx = offset + i;
const T val = output[idx];
const T g(gamma[i]);
const T b(beta[i]);
const T b = (nullptr == beta) ? (T)0 : beta[i];
output[idx] = g * (val - mu) * rsigma + b;
}
}
@ -129,7 +129,7 @@ __device__ inline void LayerNormSmall(const T val, const cub::KeyValuePair<T, T>
if (threadIdx.x < ld) {
const T g(gamma[threadIdx.x]);
const T b(beta[threadIdx.x]);
const T b = (nullptr == beta) ? (T)0 : beta[threadIdx.x];
output[idx] = g * (val - mu) * rsigma + b;
}
}

20
onnxruntime/contrib_ops/cuda/bert/skip_layer_norm.cc
View file

@ -65,14 +65,16 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
"Last dimension of gamma and input does not match");
}
const auto& beta_dims = beta->Shape().GetDims();
if (beta_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have 1 dimension, got ", beta_dims.size());
}
if (beta_dims[0] != input_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Last dimension of beta and input does not match");
if (nullptr != beta) {
const auto& beta_dims = beta->Shape().GetDims();
if (beta_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have 1 dimension, got ", beta_dims.size());
}
if (beta_dims[0] != input_dims[2]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Last dimension of beta and input does not match");
}
}
if (nullptr != bias) {
@ -98,7 +100,7 @@ Status SkipLayerNorm<T>::ComputeInternal(OpKernelContext* ctx) const {
input->template Data<T>(),
skip->template Data<T>(),
gamma->template Data<T>(),
beta->template Data<T>(),
beta != nullptr ? beta->template Data<T>() : nullptr,
bias != nullptr ? bias->template Data<T>() : nullptr,
epsilon_,
hidden_size,

4
onnxruntime/contrib_ops/cuda/layer_norm.cc
View file

@ -59,7 +59,7 @@ Status LayerNorm<T, U, simplified>::ComputeInternal(OpKernelContext* ctx) const
auto X_data = reinterpret_cast<const CudaT*>(X->template Data<T>());
auto scale_data = reinterpret_cast<const CudaT*>(scale->template Data<T>());
auto bias_data = simplified ? nullptr: reinterpret_cast<const CudaT*>(bias->template Data<T>());
auto bias_data = (simplified || (nullptr == bias)) ? nullptr: reinterpret_cast<const CudaT*>(bias->template Data<T>());
const TensorShape& x_shape = X->Shape();
const int64_t axis = HandleNegativeAxis(axis_, x_shape.NumDimensions());
@ -91,7 +91,7 @@ Status LayerNorm<T, U, simplified>::ComputeInternal(OpKernelContext* ctx) const
mean_data = reinterpret_cast<CudaU*>(mean->template MutableData<U>());
}
}
Tensor* var = ctx->Output(output_index, TensorShape(mean_inv_std_var_dim));
CudaU* inv_var_data = nullptr;
if (var != nullptr) {

4
onnxruntime/core/graph/contrib_ops/contrib_defs.cc
View file

@ -546,7 +546,7 @@ GELU (Gaussian Error Linear Unit) approximation: Y=0.5*X*(1+tanh(0.797885*X+0.03
.Input(0, "input", "3D input tensor with shape (batch_size, sequence_length, hidden_size)", "T")
.Input(1, "skip", "3D skip tensor with shape (batch_size, sequence_length, hidden_size)", "T")
.Input(2, "gamma", "1D input tensor with shape (hidden_size)", "T")
.Input(3, "beta", "1D skip tensor with shape (hidden_size", "T")
.Input(3, "beta", "1D skip tensor with shape (hidden_size", "T", OpSchema::Optional)
.Input(4, "bias", "1D bias tensor with shape (hidden_size", "T", OpSchema::Optional)
.Output(0, "output", "3D output tensor with shape (batch_size, sequence_length, hidden_size)", "T")
.Output(1, "mean", "Saved mean used during training to speed up gradient computation", "U", OpSchema::Optional)
@ -2085,7 +2085,7 @@ Example 4:
.AllowUncheckedAttributes()
.Input(0, "X", "Input data tensor from the previous layer.", "T")
.Input(1, "scale", "Scale tensor.", "T")
.Input(2, "B", "Bias tensor.", "T")
.Input(2, "B", "Bias tensor.", "T", OpSchema::Optional)
.Output(0, "Y", "Output data tensor.", "T")
.Output(1, "mean", "Saved mean used during training to speed up gradient computation", "U", OpSchema::Optional)
.Output(2, "inv_std_var", "Saved inverse standard variance used during training to speed up gradient computation.", "U", OpSchema::Optional)

23
onnxruntime/test/contrib_ops/layer_norm_op_test.cc
View file

@ -43,5 +43,28 @@ TEST(LayerNormTest, BERTLayerNorm) {
tester.Run();
}
TEST(LayerNormTest, BERTLayerNorm_NoBias) {
OpTester tester("LayerNormalization", 1 /*opset_version*/);
tester.AddAttribute<int64_t>("axis", -1);
tester.AddAttribute<float>("epsilon", 1e-12f);
// create rand inputs
RandomValueGenerator random{};
std::vector<int64_t> X_dims{4, 128};
std::vector<float> X_data = random.Uniform<float>(X_dims, 0.0f, 1.0f);
tester.AddInput<float>("X", X_dims, X_data);
std::vector<int64_t> scale_dims{128};
std::vector<float> scale_data = random.Uniform<float>(scale_dims, 0.0f, 1.0f);
tester.AddInput<float>("Scale", scale_dims, scale_data);
tester.AddMissingOptionalInput<float>();
tester.AddReferenceOutputs("testdata/layernorm_no_bias.onnx");
tester.Run();
}
} // namespace test
} // namespace onnxruntime

15
onnxruntime/test/contrib_ops/layer_norm_test.cc
View file

@ -29,7 +29,8 @@ static void TestLayerNorm(const std::vector<int64_t>& x_dims,
const std::string& op,
optional<float> epsilon,
int64_t axis = -1,
int64_t keep_dims = 1) {
int64_t keep_dims = 1,
bool no_bias = false) {
const std::vector<int64_t>& n_x_m_dims = x_dims;
std::vector<int64_t> n_dims, m_dims;
ASSERT_TRUE(SplitDims(n_x_m_dims, axis, n_dims, m_dims).IsOK());
@ -41,7 +42,7 @@ static void TestLayerNorm(const std::vector<int64_t>& x_dims,
ASSERT_NE(keep_dims, 0);
const std::vector<int64_t>& stats_dims = keep_dims ? n_and_ones_dims : n_dims;
CompareOpTester test(op.c_str());
test.AddAttribute("axis", axis);
test.AddAttribute("keep_dims", keep_dims);
@ -57,7 +58,7 @@ static void TestLayerNorm(const std::vector<int64_t>& x_dims,
test.AddInput<float>("X", n_x_m_dims, X_data);
test.AddInput<float>("scale", m_dims, scale_data, true);
if (op.compare(SIMPLIFIED_LAYER_NORM_OP) != 0) {
if (op.compare(SIMPLIFIED_LAYER_NORM_OP) != 0 && no_bias == false) {
test.AddInput<float>("B", m_dims, B_data, true);
}
@ -99,6 +100,14 @@ TEST(CudaKernelTest, LayerNorm_LargeSizeTensor) {
TestLayerNorm(X_dims, LAYER_NORM_OP, k_epsilon_default);
}
TEST(CudaKernelTest, LayerNorm_MidSizeTensor_NoBias) {
std::vector<int64_t> X_dims{8, 80, 768};
const int64_t axis = -1;
const int64_t keep_dims = 1;
const bool no_bias = true;
TestLayerNorm(X_dims, LAYER_NORM_OP, k_epsilon_default, axis, keep_dims, no_bias);
}
TEST(CudaKernelTest, SimplifiedLayerNorm_SmallSizeTensor) {
const std::vector<int64_t> X_dims{4, 20, 128};
TestLayerNorm(X_dims, SIMPLIFIED_LAYER_NORM_OP, k_epsilon_default);

51
onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc
View file

@ -21,7 +21,8 @@ static void RunTest(
int batch_size,
int sequence_length,
int hidden_size,
bool use_float16 = false) {
bool use_float16 = false,
bool no_beta = false) {
// Input and output shapes
// Input 0 - input: (batch_size, sequence_length, hidden_size)
// Input 1 - skip : (batch_size, sequence_length, hidden_size)
@ -40,7 +41,11 @@ static void RunTest(
test.AddInput<float>("input", input_dims, input_data);
test.AddInput<float>("skip", skip_dims, skip_data);
test.AddInput<float>("gamma", gamma_dims, gamma_data);
test.AddInput<float>("beta", beta_dims, beta_data);
if (!no_beta) {
test.AddInput<float>("beta", beta_dims, beta_data);
} else {
test.AddMissingOptionalInput<float>();
}
test.AddAttribute("epsilon", epsilon);
if (!bias_data.empty()) {
test.AddInput<float>("bias", bias_dims, bias_data);
@ -53,7 +58,11 @@ static void RunTest(
test.AddInput<MLFloat16>("input", input_dims, ToFloat16(input_data));
test.AddInput<MLFloat16>("skip", skip_dims, ToFloat16(skip_data));
test.AddInput<MLFloat16>("gamma", gamma_dims, ToFloat16(gamma_data));
test.AddInput<MLFloat16>("beta", beta_dims, ToFloat16(beta_data));
if (!no_beta) {
test.AddInput<MLFloat16>("beta", beta_dims, ToFloat16(beta_data));
} else {
test.AddMissingOptionalInput<float>();
}
test.AddAttribute("epsilon", epsilon);
if (!bias_data.empty()) {
test.AddInput<MLFloat16>("bias", bias_dims, ToFloat16(bias_data));
@ -138,6 +147,42 @@ TEST(SkipLayerNormTest, SkipLayerNormBatch1_Float16) {
true);
}
TEST(SkipLayerNormTest, SkipLayerNormBatch1_NoBeta) {
int batch_size = 1;
int sequence_length = 2;
int hidden_size = 4;
std::vector<float> input_data = {
0.8f, -0.5f, 0.0f, 1.f,
0.5f, 0.2f, 0.3f, -0.6f};
std::vector<float> skip_data = {
0.1f, -0.2f, 0.3f, 1.0f,
0.5f, 0.1f, 0.4f, 1.6f};
std::vector<float> gamma_data = {
0.3f, 0.2f, 4.0f, 2.2f};
std::vector<float> beta_data = {};
std::vector<float> output_data = {
0.08433859348297119f, -0.27090578377246857f, -1.32897164821624756f, 3.0924152374267578f,
0.26111652255058289f, -0.31333980560302734f, -0.69631003737449646f, 1.9148544311523438f};
RunTest(input_data,
skip_data,
gamma_data,
beta_data,
std::vector<float>(),
output_data,
epsilon_,
batch_size,
sequence_length,
hidden_size,
false,
true);
}
TEST(SkipLayerNormTest, SkipLayerNormBatch2) {
int batch_size = 2;
int sequence_length = 2;

BIN
onnxruntime/test/testdata/layernorm_no_bias.onnx vendored Normal file
View file

Binary file not shown.