saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-05 04:17:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

implement CPU contrib OP EmbedLayerNormalization (#2332)

Browse source
This commit is contained in:
Yulong Wang 2019-11-06 12:27:08 -08:00 committed by GitHub
parent 06a6d74a67
commit c0b8926863
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 244 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

114
onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.cc Normal file
View file

@ -0,0 +1,114 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "embed_layer_norm.h"
#include "embed_layer_norm_helper.h"
#include "core/util/math_cpuonly.h"
namespace onnxruntime {
namespace contrib {
// These ops are internal-only, so register outside of onnx
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
EmbedLayerNormalization, \
kMSDomain, \
1, \
T, \
kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
EmbedLayerNorm<T>);
REGISTER_KERNEL_TYPED(float)
template <typename T>
EmbedLayerNorm<T>::EmbedLayerNorm(const OpKernelInfo& info) : OpKernel(info) {}
template <typename T>
Status EmbedLayerNorm<T>::Compute(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context));
const Tensor* input_ids = context->Input<Tensor>(0);
const Tensor* segment_ids = context->Input<Tensor>(1);
const Tensor* mask = context->Input<Tensor>(2);
const Tensor* word_embedding = context->Input<Tensor>(3);
const Tensor* position_embedding = context->Input<Tensor>(4);
const Tensor* segment_embedding = context->Input<Tensor>(5);
const Tensor* gamma = context->Input<Tensor>(6);
const Tensor* beta = context->Input<Tensor>(7);
const auto input_dims = input_ids->Shape().GetDims();
int64_t hidden_size = word_embedding->Shape()[1];
std::vector<int64_t> out_dims;
out_dims.reserve(3);
out_dims.push_back(input_dims[0]);
out_dims.push_back(input_dims[1]);
out_dims.push_back(hidden_size);
TensorShape output_shape(out_dims);
Tensor* output = context->Output(0, output_shape);
std::vector<int64_t> mask_index_dims;
mask_index_dims.push_back(input_dims[0]);
TensorShape mask_index_shape(mask_index_dims);
Tensor* mask_index = context->Output(1, mask_index_shape);
int batch_size = static_cast<int>(input_dims[0]);
int sequence_length = static_cast<int>(input_dims[1]);
int word_embedding_length = static_cast<int>(word_embedding->Shape()[0]);
int position_embedding_length = static_cast<int>(position_embedding->Shape()[0]);
int segment_embedding_length = static_cast<int>(segment_embedding->Shape()[0]);
ConstEigenArrayMap<T> word_embedding_arr(word_embedding->template Data<T>(), hidden_size, word_embedding_length);
ConstEigenArrayMap<T> position_embedding_arr(position_embedding->template Data<T>(), hidden_size, position_embedding_length);
ConstEigenArrayMap<T> segment_embedding_arr(segment_embedding->template Data<T>(), hidden_size, segment_embedding_length);
ConstEigenVectorMap<T> gamma_vector(gamma->template Data<T>(), hidden_size);
ConstEigenVectorMap<T> beta_vector(beta->template Data<T>(), hidden_size);
EigenArrayMap<T> output_arr(output->template MutableData<T>(), hidden_size, batch_size * sequence_length);
// Calculate output
{
size_t index = 0;
for (int b = 0; b < batch_size; b++) {
for (int s = 0; s < sequence_length; s++) {
int word_col_index = input_ids->template Data<int>()[index];
if (word_col_index < 0 || word_col_index >= word_embedding_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "word_col_index out of range");
}
int position_col_index = s;
if (position_col_index >= position_embedding_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "position_col_index out of range");
}
int segment_col_index = segment_ids->template Data<int>()[index];
if (segment_col_index < 0 || segment_col_index >= segment_embedding_length) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "segment_col_index out of range");
}
output_arr.col(index) = word_embedding_arr.col(word_col_index) +
position_embedding_arr.col(position_col_index) +
segment_embedding_arr.col(segment_col_index);
output_arr.col(index) -= output_arr.col(index).mean();
output_arr.col(index) /= static_cast<T>(sqrt(output_arr.col(index).pow(2).mean() + 1.0e-13));
output_arr.col(index) *= gamma_vector.array();
output_arr.col(index) += beta_vector.array();
index++;
}
}
}
// Calculate mask
{
const int* mask_data = mask->template Data<int>();
for (int b = 0; b < batch_size; b++) {
mask_index->template MutableData<int>()[b] = static_cast<int>(std::count_if(mask_data + (b * sequence_length),
mask_data + (b * sequence_length) + sequence_length,
[](int v) { return v == 1; }));
}
}
return Status::OK();
}
} // namespace contrib
} // namespace onnxruntime

18
onnxruntime/contrib_ops/cpu/bert/embed_layer_norm.h Normal file
View file

@ -0,0 +1,18 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
template <typename T>
class EmbedLayerNorm : public OpKernel {
public:
explicit EmbedLayerNorm(const OpKernelInfo& info);
Status Compute(OpKernelContext* context) const override;
};
} // namespace contrib
} // namespace onnxruntime

83
onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.cc Normal file
View file

@ -0,0 +1,83 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "embed_layer_norm_helper.h"
#include "core/framework/tensorprotoutils.h"
#include "onnx/defs/tensor_proto_util.h"
namespace onnxruntime {
namespace contrib {
namespace embed_layer_norm {
Status CheckInputs(const OpKernelContext* context) {
const Tensor* input_ids = context->Input<Tensor>(0);
const Tensor* segment_ids = context->Input<Tensor>(1);
const Tensor* mask = context->Input<Tensor>(2);
const Tensor* word_embedding = context->Input<Tensor>(3);
const Tensor* position_embedding = context->Input<Tensor>(4);
const Tensor* segment_embedding = context->Input<Tensor>(5);
const Tensor* gamma = context->Input<Tensor>(6);
const Tensor* beta = context->Input<Tensor>(7);
if (input_ids->Shape() != segment_ids->Shape() || input_ids->Shape() != mask->Shape()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 0, 1 and 2 shall have same shape");
}
const auto input_dims = input_ids->Shape().GetDims();
if (input_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"input_ids is expected to have 2 dimensions, got ", input_dims.size());
}
const auto word_embedding_dims = word_embedding->Shape().GetDims();
if (word_embedding_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"word_embedding is expected to have 2 dimensions, got ", word_embedding_dims.size());
}
const auto position_embedding_dims = position_embedding->Shape().GetDims();
if (position_embedding_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size());
}
const auto segment_embedding_dims = segment_embedding->Shape().GetDims();
if (segment_embedding_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"segment_embedding is expected to have 2 dimensions, got ", segment_embedding_dims.size());
}
if (word_embedding_dims[1] != position_embedding_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"word_embedding and position_embedding shall have same dimension 1");
}
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 dimensions, got ", beta_dims.size());
}
if (beta_dims[0] != word_embedding_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"beta is expected to have size of ", word_embedding_dims[1], ", got ", beta_dims[0]);
}
const auto gamma_dims = gamma->Shape().GetDims();
if (gamma_dims.size() != 1) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"gamma is expected to have 1 dimensions, got ", gamma_dims.size());
}
if (gamma_dims[0] != word_embedding_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"gamma is expected to have size of ", word_embedding_dims[1], ", got ", gamma_dims[0]);
}
return Status::OK();
}
} // namespace embed_layer_norm
} // namespace contrib
} // namespace onnxruntime

17
onnxruntime/contrib_ops/cpu/bert/embed_layer_norm_helper.h Normal file
View file

@ -0,0 +1,17 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
namespace embed_layer_norm {
Status CheckInputs(const OpKernelContext* context);
} // namespace embed_layer_norm
} // namespace contrib
} // namespace onnxruntime

2
onnxruntime/contrib_ops/cpu_contrib_kernels.cc
View file

@ -9,6 +9,7 @@ namespace onnxruntime {
namespace contrib {
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedGemm);
@ -87,6 +88,7 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, SampleOp)>,
// add more kernels here
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, EmbedLayerNormalization)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ExpandDims)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedConv)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, FusedGemm)>,

37
onnxruntime/contrib_ops/cuda/bert/embed_layer_norm.cc
View file

@ -5,6 +5,7 @@
#include "core/providers/cuda/cudnn_common.h"
#include "core/framework/tensorprotoutils.h"
#include "onnx/defs/tensor_proto_util.h"
#include "contrib_ops/cpu/bert/embed_layer_norm_helper.h"
#include "embed_layer_norm.h"
#include "embed_layer_norm_impl.h"
@ -34,6 +35,8 @@ EmbedLayerNorm<T>::EmbedLayerNorm(const OpKernelInfo& op_kernel_info) : CudaKern
template <typename T>
Status EmbedLayerNorm<T>::ComputeInternal(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(embed_layer_norm::CheckInputs(context));
const Tensor* input_ids = context->Input<Tensor>(0);
const Tensor* segment_ids = context->Input<Tensor>(1);
const Tensor* mask = context->Input<Tensor>(2);
@ -43,40 +46,8 @@ Status EmbedLayerNorm<T>::ComputeInternal(OpKernelContext* context) const {
const Tensor* gamma = context->Input<Tensor>(6);
const Tensor* beta = context->Input<Tensor>(7);
if (input_ids->Shape() != segment_ids->Shape() || input_ids->Shape() != mask->Shape()) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"Input 0, 1 and 2 shall have same shape");
}
const auto input_dims = input_ids->Shape().GetDims();
if (input_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"input_ids is expected to have 2 dimensions, got ", input_dims.size());
}
const auto word_embedding_dims = word_embedding->Shape().GetDims();
if (word_embedding_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"word_embedding is expected to have 2 dimensions, got ", word_embedding_dims.size());
}
const auto position_embedding_dims = position_embedding->Shape().GetDims();
if (position_embedding_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"position_embedding is expected to have 2 dimensions, got ", position_embedding_dims.size());
}
const auto segment_embedding_dims = segment_embedding->Shape().GetDims();
if (segment_embedding_dims.size() != 2) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"segment_embedding is expected to have 2 dimensions, got ", segment_embedding_dims.size());
}
if (word_embedding_dims[1] != position_embedding_dims[1]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"word_embedding and position_embedding shall have same dimension 1");
}
int64_t hidden_size = word_embedding_dims[1];
int64_t hidden_size = word_embedding->Shape()[1];
std::vector<int64_t> out_dims;
out_dims.reserve(3);

10
onnxruntime/test/contrib_ops/embedlayernorm_op_test.cc
View file

@ -24,7 +24,11 @@ static void RunTest(
int hidden_size,
bool use_float16 = false) {
int min_cuda_architecture = use_float16 ? 530 : 0;
if (HasCudaEnvironment(min_cuda_architecture)) {
bool enable_cuda = HasCudaEnvironment(min_cuda_architecture);
bool enable_cpu = !use_float16;
if (enable_cpu || enable_cuda) {
// Input and output shapes
// Input 0 - input_ids : (batch_size, sequence_length)
// Input 1 - segment_ids : (batch_size, sequence_length)
@ -76,9 +80,7 @@ static void RunTest(
}
tester.AddOutput<int32_t>("mask_index", mask_index_dims, mask_index_data);
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCudaExecutionProvider());
tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
tester.Run();
}
}