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

Add word conv embedding custom op (#229)

Browse source 
* run bw and fw sequentially for GRU if using MKLDNN

* word conv embedding custom op

* run bw and fw sequentially for GRU if using MKLDNN

* Add word conv embedding custom op

* fix build break in linux

* fix macos build break

* resolve the comments

* refine the comments

* remove unnessary comment

* rename the function to calculate the length of eache word in a sequence

* add license info and fix typo
This commit is contained in:
Yufeng Li 2018-12-21 10:48:51 -08:00 committed by GitHub
parent a37887cfa1
commit 4e74ffba91
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 455 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/contrib_ops/contrib_kernels.cc
View file

@ -17,6 +17,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, StringNormalizer);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WordConvEmbedding);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherND);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3);
@ -34,6 +35,7 @@ void RegisterContribKernels(std::function<void(KernelCreateInfo&&)> fn) {
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, string, StringNormalizer)>());
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, NonMaxSuppression)>());
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Range)>());
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, WordConvEmbedding)>());
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, GatherND)>());
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MurmurHash3)>());
}

219
onnxruntime/contrib_ops/cpu/word_conv_embedding.cc Normal file
View file

@ -0,0 +1,219 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "word_conv_embedding.h"
#include "core/util/math.h"
#include "core/util/math_cpuonly.h"
#include "core/mlas/inc/mlas.h"
namespace onnxruntime {
namespace contrib {
void WordConvEmbedding::CharEmbeddingLookup(
const int* seq_ptr,
const float* char_embedding_weight_p,
size_t seq_len,
size_t word_len,
size_t char_embedding_size,
const int* words_len_ptr,
float* dst) const {
for (size_t word_inx = 0; word_inx < seq_len; word_inx++) {
if (words_len_ptr[word_inx] > 0) {
const int* cur_seq_ptr = seq_ptr + word_inx * word_len;
float* cur_dst_ptr = dst + word_inx * word_len * char_embedding_size;
for (size_t char_inx = 0; char_inx < word_len; char_inx++) {
memcpy(cur_dst_ptr, char_embedding_weight_p + (*cur_seq_ptr) * char_embedding_size, sizeof(float) * char_embedding_size);
cur_dst_ptr += char_embedding_size;
cur_seq_ptr++;
}
}
}
}
//input : [sequence_length, word_length, char_embedding_size]
void WordConvEmbedding::ComputeConvMaxPoolWithActivation(
AllocatorPtr allocator,
const float* input,
const float* weights,
const float* bias,
const int* words_len_ptr,
int64_t seq_len,
int64_t word_len,
int64_t char_embedding_size,
int64_t filter_width,
int64_t num_filters,
float* output) const {
int64_t input_word_size = word_len * char_embedding_size;
int64_t unfolded_width = word_len - filter_width + 1;
int64_t unfolded_kernal_size = filter_width * char_embedding_size;
int64_t unfolded_segment_size = unfolded_width * unfolded_kernal_size;
int64_t conv_res_segment_size = unfolded_width * num_filters;
int64_t memcpy_size = unfolded_kernal_size * sizeof(float);
auto input_unfolded_buffer_p = IAllocator::MakeUniquePtr<float>(allocator, seq_len * unfolded_segment_size);
auto conv_result_p = IAllocator::MakeUniquePtr<float>(allocator, seq_len * conv_res_segment_size);
auto conv_activation_result_p = IAllocator::MakeUniquePtr<float>(allocator, seq_len * conv_res_segment_size);
for (int64_t word_inx = 0; word_inx < seq_len; word_inx++) {
if (words_len_ptr[word_inx] <= 0) continue;
const float* current_word_input = input + word_inx * input_word_size;
float* current_word_unfolded_buffer_p = input_unfolded_buffer_p.get() + word_inx * unfolded_segment_size;
float* conv_buf_p = conv_result_p.get() + word_inx * conv_res_segment_size;
float* pactivationbuf = conv_activation_result_p.get() + word_inx * conv_res_segment_size;
float* pres = output + word_inx * num_filters;
// Unfolding from pin to pufbuf.
float* tmp_unfolded_buffer_ptr = current_word_unfolded_buffer_p;
for (int64_t unfolded_inx = 0; unfolded_inx < unfolded_width; unfolded_inx++) {
memcpy(tmp_unfolded_buffer_ptr, current_word_input, memcpy_size);
current_word_input += char_embedding_size;
tmp_unfolded_buffer_ptr += unfolded_kernal_size;
}
math::GemmEx<float, CPUMathUtil>(
CblasNoTrans, CblasTrans,
static_cast<int>(unfolded_width), static_cast<int>(num_filters), static_cast<int>(unfolded_kernal_size), 1.0f,
current_word_unfolded_buffer_p, static_cast<int>(unfolded_kernal_size),
weights, static_cast<int>(unfolded_kernal_size), 0.0f,
conv_buf_p, static_cast<int>(num_filters), &CPUMathUtil::Instance());
for (int64_t unfolded_inx = 0; unfolded_inx < unfolded_width; unfolded_inx++)
for (int64_t filter_inx = 0; filter_inx < num_filters; filter_inx++) {
conv_buf_p[unfolded_inx * num_filters + filter_inx] += bias[filter_inx];
}
MlasComputeTanh(conv_buf_p, pactivationbuf, unfolded_width * num_filters);
// Max pooling.
for (int64_t filter_inx = 0; filter_inx < num_filters; filter_inx++) {
pres[filter_inx] = -1.0f * 1e12f;
}
for (int64_t unfolded_inx = 0; unfolded_inx < unfolded_width; unfolded_inx++) {
if (unfolded_inx > 0 && unfolded_inx > (words_len_ptr[word_inx] - filter_width)) break;
float* pcur = pactivationbuf + unfolded_inx * num_filters;
for (int64_t filter_inx = 0; filter_inx < num_filters; filter_inx++) {
pres[filter_inx] = std::max(pcur[filter_inx], pres[filter_inx]);
}
}
}
}
void WordConvEmbedding::CalculateLengthOfEachWordInSequence(
const int* seq_ptr,
int* words_len_ptr,
size_t seq_len,
size_t word_len) const {
for (size_t seq_inx = 0; seq_inx < seq_len; seq_inx++) {
size_t w_off = seq_inx * word_len;
int word_length = 0;
if (seq_ptr[w_off] > 0) {
for (size_t char_inx = 0; char_inx < word_len; char_inx++) {
if (seq_ptr[w_off + char_inx] > 0) word_length++;
}
}
words_len_ptr[seq_inx] = word_length;
}
}
Status WordConvEmbedding::ValidateInputShape(const TensorShape& w_conv_shape, const TensorShape& w_char_embedding_shape) const {
if (embedding_size_ != -1 && w_conv_shape[0] != embedding_size_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Conv filter size does not match embedding_size attribute.",
" embedding_size attribute: ", embedding_size_,
" conv filter size: ", w_conv_shape[0]);
}
if (conv_window_size_ != -1 && w_conv_shape[2] != conv_window_size_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Conv kernal size 1 does not match conv_window_size attribute .",
" conv_window_size attribute: ", conv_window_size_,
" conv kernal size 1: ", w_conv_shape[2]);
}
if (char_embedding_size_ != -1 && w_char_embedding_shape[1] != char_embedding_size_) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Char embedding size does not match char_embedding_size attribute.",
" char_embedding_size attribute: ", conv_window_size_,
" Char embedding size: ", w_conv_shape[1]);
}
if (w_char_embedding_shape[1] != w_conv_shape[3]) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Char embedding size does not match conv kernal size 2.",
" Char embedding size: ", conv_window_size_,
" Conv kernal size 2 : ", w_conv_shape[3]);
}
return Status::OK();
}
Status WordConvEmbedding::Compute(OpKernelContext* ctx) const {
// original lstm processing
const Tensor& sequence = *(ctx->Input<Tensor>(0)); // sequence: [sequence_length, word_length]
const Tensor& w_conv = *(ctx->Input<Tensor>(1)); // conv weight: [M, C/group, kH, kW]
const Tensor& b_conv = *(ctx->Input<Tensor>(2)); // conv bias: [M]
const Tensor& w_char_embedding = *(ctx->Input<Tensor>(3)); // conv weights. [index, char_embedding_size]
const TensorShape& sequence_shape = sequence.Shape();
const TensorShape& w_conv_shape = w_conv.Shape();
const TensorShape& w_char_embedding_shape = w_char_embedding.Shape();
ORT_RETURN_IF_ERROR(ValidateInputShape(w_conv_shape, w_char_embedding_shape));
int64_t seq_len = sequence_shape[0];
int64_t word_len = sequence_shape[1];
int64_t char_embedding_size = w_char_embedding_shape[1];
int64_t filter_size = w_conv_shape[0];
int64_t filter_width = w_conv_shape[2];
TensorShape Y_dims{seq_len, filter_size};
Tensor* Y = ctx->Output(/*index*/ 0, Y_dims);
const int* seq_ptr = sequence.Data<int>();
AllocatorPtr alloc;
ORT_RETURN_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc));
// allocate memory for char look up
// seq_len * word_len * char_embedding_size
size_t chars_embeddings_size = seq_len * word_len * char_embedding_size;
auto chars_embeddings_ptr = IAllocator::MakeUniquePtr<float>(alloc, chars_embeddings_size);
auto words_length_ptr = IAllocator::MakeUniquePtr<int>(alloc, seq_len);
std::memset(chars_embeddings_ptr.get(), 0, chars_embeddings_size * sizeof(float));
std::memset(words_length_ptr.get(), 0, seq_len * sizeof(int));
CalculateLengthOfEachWordInSequence(seq_ptr, words_length_ptr.get(), seq_len, word_len);
CharEmbeddingLookup(seq_ptr,
w_char_embedding.Data<float>(),
seq_len,
word_len,
char_embedding_size,
words_length_ptr.get(),
chars_embeddings_ptr.get());
ComputeConvMaxPoolWithActivation(
alloc,
chars_embeddings_ptr.get(),
w_conv.Data<float>(),
b_conv.Data<float>(),
words_length_ptr.get(),
seq_len,
word_len,
char_embedding_size,
filter_width,
filter_size,
Y->MutableData<float>());
return Status::OK();
}
/* Range operator */
ONNX_OPERATOR_KERNEL_EX(
WordConvEmbedding, //name
kMSDomain,
1,
kCpuExecutionProvider,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<int32_t>()).TypeConstraint("T1", DataTypeImpl::GetTensorType<float>()),
WordConvEmbedding);
} // namespace contrib
} // namespace onnxruntime

58
onnxruntime/contrib_ops/cpu/word_conv_embedding.h Normal file
View file

@ -0,0 +1,58 @@
// 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"
#include "core/framework/tensor.h"
namespace onnxruntime {
namespace contrib {
class WordConvEmbedding final : public OpKernel {
public:
explicit WordConvEmbedding(const OpKernelInfo& info) : OpKernel(info) {
}
Status Compute(OpKernelContext* context) const override;
private:
void CharEmbeddingLookup(
const int* seq_ptr,
const float* char_embedding_weight_p,
size_t seq_len,
size_t word_len,
size_t char_embedding_size,
const int* words_len_ptr,
float* dst) const;
void ComputeConvMaxPoolWithActivation(
AllocatorPtr allocator,
const float* input,
const float* weights,
const float* bias,
const int* words_len_ptr,
int64_t seq_len,
int64_t word_len,
int64_t char_embedding_size,
int64_t filter_width,
int64_t num_filters,
float* output) const;
void CalculateLengthOfEachWordInSequence(
const int* seq_ptr,
int* words_len_ptr,
size_t seq_len,
size_t word_len) const;
Status ValidateInputShape(
const TensorShape& w_conv_shape,
const TensorShape& w_char_embedding_shape) const;
private:
int64_t embedding_size_{Info().GetAttrOrDefault<int64_t>("embedding_size", -1)};
int64_t conv_window_size_{Info().GetAttrOrDefault<int64_t>("conv_window_size", -1)};
int64_t char_embedding_size_{Info().GetAttrOrDefault<int64_t>("char_embedding_size", -1)};
};
} // namespace contrib
} // namespace onnxruntime

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

@ -657,9 +657,46 @@ Example 4:
output = [[[2,3]],[[4,5]]]
)DOC");
ONNX_CONTRIB_OPERATOR_SCHEMA( WordConvEmbedding )
.SetDomain( kMSDomain )
.SinceVersion( 1 )
.Attr(
"embedding_size",
"Integer representing the embedding vector size for each word."
"If not provide, use the fileter size of conv weight",
AttributeProto::INT,
OPTIONAL)
.Attr(
"conv_window_size",
"This operator applies convolution to word from left to right with window equal to conv_window_size and stride to 1."
"Take word 'example' for example, with conv_window_size equal to 2, conv is applied to [ex],[xa], [am], [mp]..."
"If not provide, use the first dimension of conv kernal shape.",
AttributeProto::INT,
OPTIONAL)
.Attr(
"char_embedding_size",
"Integer representing the embedding vector size for each char."
"If not provide, use the char embedding size of embedding vector.",
AttributeProto::INT,
OPTIONAL)
.Input( 0, "Sequence", "Specify batchs of sequence words to embedding", "T" )
.Input( 1, "W", "Specify weights of conv", "T1" )
.Input( 2, "B", "Specify bias of conv", "T1" )
.Input( 3, "C", "Specify embedding vector of char", "T1" )
.Output( 0, "Y", "output", "T1" )
.TypeConstraint(
"T",
{ "tensor(int32)" },
"Constrain to tensor(int32)." )
.TypeConstraint(
"T1",
{ "tensor(float)" },
"Constrain to tensor(float).")
.SetDoc( R"DOC(The WordConvEmbedding takes in a batch of sequence words and embed each word to a vector.)DOC" );
#ifdef MICROSOFT_INTERNAL
// register internal ops
RegisterInternalSchemas();
// register internal ops
RegisterInternalSchemas();
#endif
}
} // namespace contrib

6
onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc
View file

@ -374,8 +374,10 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const {
gsl::span<T> hidden_output_2 = hidden_output.subspan(hidden_output_size_per_direction,
hidden_output_size_per_direction);
#ifndef USE_MKLDNN
std::packaged_task<void()> task_fw{
[&]() {
#endif // ! USE_MKLDNN
std::unique_ptr<detail::UniDirectionalGru<T>> fw = std::make_unique<detail::UniDirectionalGru<T>>(
alloc, logger,
seq_length, batch_size, input_size, hidden_size_, linear_before_reset_, Direction::kForward,
@ -384,9 +386,11 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const {
activation_funcs_.Entries()[1],
clip_, ttp_);
fw->Compute(input, sequence_lens_span, num_directions_, input_weights_1, recurrent_weights_1, output_1, hidden_output_1);
#ifndef USE_MKLDNN
}};
auto task_results_fw = task_fw.get_future();
ttp_.RunTask(std::move(task_fw));
#endif // ! USE_MKLDNN
std::unique_ptr<detail::UniDirectionalGru<T>> bw = std::make_unique<detail::UniDirectionalGru<T>>(
alloc, logger,
@ -397,7 +401,9 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const {
clip_, ttp_);
bw->Compute(input, sequence_lens_span, num_directions_, input_weights_2, recurrent_weights_2, output_2, hidden_output_2);
#ifndef USE_MKLDNN
task_results_fw.get();
#endif // ! USE_MKLDNN
} else {
std::unique_ptr<detail::UniDirectionalGru<T>> gru_p = std::make_unique<detail::UniDirectionalGru<T>>(
alloc, logger,

131
onnxruntime/test/contrib_ops/word_conv_embedding_test.cc Normal file
View file

@ -0,0 +1,131 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <codecvt>
#include <vector>
#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"
namespace onnxruntime {
namespace test {
void InitializeTestWithoutAttribute(OpTester& test) {
// sequence has 2 words and each words has 5 chars
std::vector<int64_t> seq_words_shape = {2, 5};
std::vector<int> seq_words{1, 2, 3, 4, 0,
4, 3, 2, 1, 0};
// Charset has 5 chars and each char is represented with a vector of 3
std::vector<int64_t> W_char_embedding_shape = {5, 3};
std::vector<float> W_char_embedding{0.1f, 0.2f, 0.3f,
0.2f, 0.3f, 0.1f,
0.3f, 0.1f, 0.2f,
0.4f, 0.5f, 0.6f,
0.7f, 0.8f, 0.9f};
std::vector<int64_t> W_conv_shape = {2, 1, 2, 3};
std::vector<float> W_conv{0.1f, 0.2f, 0.3f,
0.2f, 0.3f, 0.1f,
0.3f, 0.1f, 0.2f,
1.0f, 1.1f, 1.2f};
std::vector<int64_t> B_conv_shape = {2};
std::vector<float> B_conv{0.1f, 0.2f};
std::vector<int64_t> output_shape = {2, 2};
std::vector<float> output{0.711393774f, 0.996334076f, 0.711393774f, 0.981612563f};
test.AddInput<int>("Sequence", seq_words_shape, seq_words);
test.AddInput<float>("W", W_conv_shape, W_conv);
test.AddInput<float>("B", B_conv_shape, B_conv);
test.AddInput<float>("C", W_char_embedding_shape, W_char_embedding);
test.AddOutput<float>("Y", output_shape, output);
}
TEST(ContribOpTest, WordConvEmbedding) {
// Invalid input dimensions
OpTester test("WordConvEmbedding", 1, onnxruntime::kMSDomain);
InitializeTestWithoutAttribute(test);
test.Run();
}
TEST(ContribOpTest, WordConvEmbedding_valid_attribute) {
// Invalid input dimensions
OpTester test("WordConvEmbedding", 1, onnxruntime::kMSDomain);
InitializeTestWithoutAttribute(test);
test.AddAttribute<int64_t>("embedding_size", 2LL);
test.AddAttribute<int64_t>("conv_window_size", 2LL);
test.AddAttribute<int64_t>("char_embedding_size", 3LL);
test.Run();
}
TEST(ContribOpTest, WordConvEmbedding_embedding_size_mismatch) {
// Invalid input dimensions
OpTester test("WordConvEmbedding", 1, onnxruntime::kMSDomain);
InitializeTestWithoutAttribute(test);
test.AddAttribute<int64_t>("embedding_size", 3LL);
test.AddAttribute<int64_t>("conv_window_size", 2LL);
test.AddAttribute<int64_t>("char_embedding_size", 3LL);
test.Run(OpTester::ExpectResult::kExpectFailure);
}
TEST(ContribOpTest, WordConvEmbedding_conv_window_size_mismatch) {
// Invalid input dimensions
OpTester test("WordConvEmbedding", 1, onnxruntime::kMSDomain);
InitializeTestWithoutAttribute(test);
test.AddAttribute<int64_t>("embedding_size", 2LL);
test.AddAttribute<int64_t>("conv_window_size", 1LL);
test.AddAttribute<int64_t>("char_embedding_size", 3LL);
test.Run(OpTester::ExpectResult::kExpectFailure);
}
TEST(ContribOpTest, WordConvEmbedding_char_embedding_size_mismatch) {
// Invalid input dimensions
OpTester test("WordConvEmbedding", 1, onnxruntime::kMSDomain);
InitializeTestWithoutAttribute(test);
test.AddAttribute<int64_t>("embedding_size", 2LL);
test.AddAttribute<int64_t>("conv_window_size", 2LL);
test.AddAttribute<int64_t>("char_embedding_size", 4LL);
test.Run(OpTester::ExpectResult::kExpectFailure);
}
TEST(ContribOpTest, WordConvEmbedding_char_embedding_shape_conv_shape_not_match) {
// Invalid input dimensions
OpTester test("WordConvEmbedding", 1, onnxruntime::kMSDomain);
// sequence has 2 words and each words has 5 chars
std::vector<int64_t> seq_words_shape = {2, 5};
std::vector<int> seq_words{1, 2, 3, 4, 0,
4, 3, 2, 1, 0};
// Charset has 5 chars and each char is represented with a vector of 3
std::vector<int64_t> W_char_embedding_shape = {5, 3};
std::vector<float> W_char_embedding{0.1f, 0.2f, 0.3f,
0.2f, 0.3f, 0.1f,
0.3f, 0.1f, 0.2f,
0.4f, 0.5f, 0.6f,
0.7f, 0.8f, 0.9f};
std::vector<int64_t> W_conv_shape = {2, 1, 2, 2};
std::vector<float> W_conv{0.1f, 0.2f,
0.2f, 0.3f,
0.3f, 0.1f,
1.0f, 1.1f};
std::vector<int64_t> B_conv_shape = {2};
std::vector<float> B_conv{0.1f, 0.2f};
std::vector<int64_t> output_shape = {2, 2};
std::vector<float> output{0.711393774f, 0.996334076f, 0.711393774f, 0.981612563f};
test.AddInput<int>("Sequence", seq_words_shape, seq_words);
test.AddInput<float>("W", W_conv_shape, W_conv);
test.AddInput<float>("B", B_conv_shape, B_conv);
test.AddInput<float>("C", W_char_embedding_shape, W_char_embedding);
test.AddOutput<float>("Y", output_shape, output);
test.Run(OpTester::ExpectResult::kExpectFailure);
}
} // namespace test
} // namespace onnxruntime