From 4e74ffba9122695555754e5ec14a3faf094ce24f Mon Sep 17 00:00:00 2001 From: Yufeng Li Date: Fri, 21 Dec 2018 10:48:51 -0800 Subject: [PATCH] Add word conv embedding custom op (#229) * 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 --- onnxruntime/contrib_ops/contrib_kernels.cc | 2 + .../contrib_ops/cpu/word_conv_embedding.cc | 219 ++++++++++++++++++ .../contrib_ops/cpu/word_conv_embedding.h | 58 +++++ .../core/graph/contrib_ops/contrib_defs.cc | 41 +++- .../core/providers/cpu/rnn/deep_cpu_gru.cc | 6 + .../contrib_ops/word_conv_embedding_test.cc | 131 +++++++++++ 6 files changed, 455 insertions(+), 2 deletions(-) create mode 100644 onnxruntime/contrib_ops/cpu/word_conv_embedding.cc create mode 100644 onnxruntime/contrib_ops/cpu/word_conv_embedding.h create mode 100644 onnxruntime/test/contrib_ops/word_conv_embedding_test.cc diff --git a/onnxruntime/contrib_ops/contrib_kernels.cc b/onnxruntime/contrib_ops/contrib_kernels.cc index a8f8876e77..9788bd13af 100644 --- a/onnxruntime/contrib_ops/contrib_kernels.cc +++ b/onnxruntime/contrib_ops/contrib_kernels.cc @@ -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 fn) { fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); + fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); } diff --git a/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc b/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc new file mode 100644 index 0000000000..4868fa8779 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/word_conv_embedding.cc @@ -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(allocator, seq_len * unfolded_segment_size); + auto conv_result_p = IAllocator::MakeUniquePtr(allocator, seq_len * conv_res_segment_size); + auto conv_activation_result_p = IAllocator::MakeUniquePtr(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( + CblasNoTrans, CblasTrans, + static_cast(unfolded_width), static_cast(num_filters), static_cast(unfolded_kernal_size), 1.0f, + current_word_unfolded_buffer_p, static_cast(unfolded_kernal_size), + weights, static_cast(unfolded_kernal_size), 0.0f, + conv_buf_p, static_cast(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(0)); // sequence: [sequence_length, word_length] + const Tensor& w_conv = *(ctx->Input(1)); // conv weight: [M, C/group, kH, kW] + const Tensor& b_conv = *(ctx->Input(2)); // conv bias: [M] + const Tensor& w_char_embedding = *(ctx->Input(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(); + + 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(alloc, chars_embeddings_size); + auto words_length_ptr = IAllocator::MakeUniquePtr(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(), + seq_len, + word_len, + char_embedding_size, + words_length_ptr.get(), + chars_embeddings_ptr.get()); + + ComputeConvMaxPoolWithActivation( + alloc, + chars_embeddings_ptr.get(), + w_conv.Data(), + b_conv.Data(), + words_length_ptr.get(), + seq_len, + word_len, + char_embedding_size, + filter_width, + filter_size, + Y->MutableData()); + + return Status::OK(); +} + +/* Range operator */ +ONNX_OPERATOR_KERNEL_EX( + WordConvEmbedding, //name + kMSDomain, + 1, + kCpuExecutionProvider, + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()).TypeConstraint("T1", DataTypeImpl::GetTensorType()), + WordConvEmbedding); + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/word_conv_embedding.h b/onnxruntime/contrib_ops/cpu/word_conv_embedding.h new file mode 100644 index 0000000000..b402f3ff72 --- /dev/null +++ b/onnxruntime/contrib_ops/cpu/word_conv_embedding.h @@ -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("embedding_size", -1)}; + int64_t conv_window_size_{Info().GetAttrOrDefault("conv_window_size", -1)}; + int64_t char_embedding_size_{Info().GetAttrOrDefault("char_embedding_size", -1)}; +}; + +} // namespace contrib +} // namespace onnxruntime diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc index 9a2e58057c..bc347e320c 100644 --- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc @@ -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 diff --git a/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc b/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc index 16a7158634..311844b89e 100644 --- a/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc +++ b/onnxruntime/core/providers/cpu/rnn/deep_cpu_gru.cc @@ -374,8 +374,10 @@ Status DeepCpuGruOp::ComputeImpl(OpKernelContext& context) const { gsl::span hidden_output_2 = hidden_output.subspan(hidden_output_size_per_direction, hidden_output_size_per_direction); +#ifndef USE_MKLDNN std::packaged_task task_fw{ [&]() { +#endif // ! USE_MKLDNN std::unique_ptr> fw = std::make_unique>( 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> bw = std::make_unique>( 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> gru_p = std::make_unique>( alloc, logger, diff --git a/onnxruntime/test/contrib_ops/word_conv_embedding_test.cc b/onnxruntime/test/contrib_ops/word_conv_embedding_test.cc new file mode 100644 index 0000000000..c2e12569f5 --- /dev/null +++ b/onnxruntime/test/contrib_ops/word_conv_embedding_test.cc @@ -0,0 +1,131 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#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 seq_words_shape = {2, 5}; + std::vector 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 W_char_embedding_shape = {5, 3}; + std::vector 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 W_conv_shape = {2, 1, 2, 3}; + std::vector 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 B_conv_shape = {2}; + std::vector B_conv{0.1f, 0.2f}; + + std::vector output_shape = {2, 2}; + std::vector output{0.711393774f, 0.996334076f, 0.711393774f, 0.981612563f}; + + test.AddInput("Sequence", seq_words_shape, seq_words); + test.AddInput("W", W_conv_shape, W_conv); + test.AddInput("B", B_conv_shape, B_conv); + test.AddInput("C", W_char_embedding_shape, W_char_embedding); + test.AddOutput("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("embedding_size", 2LL); + test.AddAttribute("conv_window_size", 2LL); + test.AddAttribute("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("embedding_size", 3LL); + test.AddAttribute("conv_window_size", 2LL); + test.AddAttribute("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("embedding_size", 2LL); + test.AddAttribute("conv_window_size", 1LL); + test.AddAttribute("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("embedding_size", 2LL); + test.AddAttribute("conv_window_size", 2LL); + test.AddAttribute("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 seq_words_shape = {2, 5}; + std::vector 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 W_char_embedding_shape = {5, 3}; + std::vector 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 W_conv_shape = {2, 1, 2, 2}; + std::vector W_conv{0.1f, 0.2f, + 0.2f, 0.3f, + 0.3f, 0.1f, + 1.0f, 1.1f}; + + std::vector B_conv_shape = {2}; + std::vector B_conv{0.1f, 0.2f}; + + std::vector output_shape = {2, 2}; + std::vector output{0.711393774f, 0.996334076f, 0.711393774f, 0.981612563f}; + + test.AddInput("Sequence", seq_words_shape, seq_words); + test.AddInput("W", W_conv_shape, W_conv); + test.AddInput("B", B_conv_shape, B_conv); + test.AddInput("C", W_char_embedding_shape, W_char_embedding); + test.AddOutput("Y", output_shape, output); + + test.Run(OpTester::ExpectResult::kExpectFailure); +} + +} // namespace test +} // namespace onnxruntime