From e64e30ee0d584bae66d77958326b35cd4de66805 Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Tue, 27 Apr 2021 16:49:03 -0700 Subject: [PATCH] Improve ConvTranspose by transposing const filter during prepacking. (#7388) * Improve ConvTranspose by transposing const filter during prepacking. * Fix CI build break for openvino which can not load such onnx model now. --- .../core/providers/cpu/nn/conv_transpose.cc | 224 +++++++++++++----- .../core/providers/cpu/nn/conv_transpose.h | 6 + .../cpu/nn/conv_transpose_attributes.h | 18 +- .../cpu/nn/conv_transpose_op_test.cc | 90 ++++++- 4 files changed, 264 insertions(+), 74 deletions(-) diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc index 67f3cca4cb..d28a63268b 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose.cc +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose.cc @@ -17,6 +17,7 @@ #include "core/providers/cpu/nn/conv_transpose.h" +#include "core/mlas/inc/mlas.h" #include "core/common/safeint.h" #include "core/util/math.h" #include "core/util/math_cpuonly.h" @@ -35,6 +36,45 @@ ONNX_CPU_OPERATOR_KERNEL( KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), ConvTranspose); +template +Status ConvTranspose::PrePack(const Tensor& /* tensor */, int /* input_idx */, bool& is_packed) { + is_packed = false; + return Status::OK(); +} + +template <> +Status ConvTranspose::PrePack(const Tensor& tensor, int input_idx, bool& is_packed) { + is_packed = false; + + // only pack filter tensor + if (input_idx == 1) { + if (tensor.Shape().NumDimensions() <= 2) { + return Status::OK(); + } + filter_shape_ = tensor.Shape(); + + const size_t K = static_cast(filter_shape_[0]) / conv_transpose_attrs_.group; + const size_t N = filter_shape_.SizeFromDimension(1); + auto packed_elements_per_group = N * K; + if (packed_elements_per_group == 0 || N == 1 || K == 1) { // No need for single row or single col case + return Status::OK(); + } + + auto alloc = Info().GetAllocator(0, OrtMemTypeDefault); + auto* packed_filter_data = alloc->Alloc(packed_elements_per_group * sizeof(float) * conv_transpose_attrs_.group); + transposed_filter_ = BufferUniquePtr(packed_filter_data, BufferDeleter(alloc)); + + for (int64_t group_id = 0; group_id < conv_transpose_attrs_.group; ++group_id) { + MlasTranspose(tensor.Data() + (N * K * group_id), + ((float*)packed_filter_data) + (group_id * packed_elements_per_group), + K, N); + } + + is_packed = true; + } + return Status::OK(); +} + template Status ConvTranspose::Compute(OpKernelContext* context) const { return ConvTranspose::DoConvTranspose(context, false); @@ -73,25 +113,25 @@ Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_ const T* Xdata = p.X->template Data(); const T* filter_data = p.F->template Data(); T* Ydata = p.Y->template MutableData(); + TensorShape output_shape = p.Y->Shape().Slice(2); - if (p.X->Shape().NumDimensions() == 4) { - for (auto image_id = 0; image_id < p.N; ++image_id) { - for (int group_id = 0; group_id < conv_transpose_attrs_.group; ++group_id) { - // Weight term - math::Gemm( - CblasTrans, - CblasNoTrans, - kernel_dim, - input_image_size, - p.num_input_channels / conv_transpose_attrs_.group, - 1, - filter_data + group_id * W_offset, - Xdata + group_id * X_offset, - 0, - col_buffer_data, - thread_pool); + for (auto image_id = 0; image_id < p.N; ++image_id) { + for (int group_id = 0; group_id < conv_transpose_attrs_.group; ++group_id) { + // Weight term + math::Gemm( + CblasTrans, + CblasNoTrans, + kernel_dim, + input_image_size, + p.num_input_channels / conv_transpose_attrs_.group, + 1, + filter_data + group_id * W_offset, + Xdata + group_id * X_offset, + 0, + col_buffer_data, + thread_pool); - // Col2im + if (p.X->Shape().NumDimensions() == 4) { math::Col2im( col_buffer_data, p.num_output_channels / conv_transpose_attrs_.group, @@ -109,37 +149,7 @@ Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_ p.strides[1], Ydata + group_id * Y_offset, &CPUMathUtil::Instance()); - } - - if (p.B != nullptr) { - auto Ymatrix = EigenMatrixMap(Ydata, output_size, p.num_output_channels); - auto Bvec = ConstEigenVectorMap(p.B->template Data(), p.num_output_channels); - Ymatrix.rowwise() += Bvec.transpose(); - } - - Xdata += X_offset * conv_transpose_attrs_.group; - Ydata += Y_offset * conv_transpose_attrs_.group; - } - } else { - TensorShape output_shape = p.Y->Shape().Slice(2); - - for (auto image_id = 0; image_id < p.N; ++image_id) { - for (int group_id = 0; group_id < conv_transpose_attrs_.group; ++group_id) { - // Weight term - math::Gemm( - CblasTrans, - CblasNoTrans, - kernel_dim, - input_image_size, - p.num_input_channels / conv_transpose_attrs_.group, - 1, - filter_data + group_id * W_offset, - Xdata + group_id * X_offset, - 0, - col_buffer_data, - thread_pool); - - // Col2im + } else { math::Col2imNd( col_buffer_data, output_shape.GetDims().data(), @@ -154,16 +164,116 @@ Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_ Ydata + group_id * Y_offset, &CPUMathUtil::Instance()); } - - if (p.B != nullptr) { - auto Ymatrix = EigenMatrixMap(Ydata, output_size, p.num_output_channels); - auto Bvec = ConstEigenVectorMap(p.B->template Data(), p.num_output_channels); - Ymatrix.rowwise() += Bvec.transpose(); - } - - Xdata += X_offset * conv_transpose_attrs_.group; - Ydata += Y_offset * conv_transpose_attrs_.group; } + + if (p.B != nullptr) { + auto Ymatrix = EigenMatrixMap(Ydata, output_size, p.num_output_channels); + auto Bvec = ConstEigenVectorMap(p.B->template Data(), p.num_output_channels); + Ymatrix.rowwise() += Bvec.transpose(); + } + + Xdata += X_offset * conv_transpose_attrs_.group; + Ydata += Y_offset * conv_transpose_attrs_.group; + } + + return Status::OK(); +} + +template <> +Status ConvTranspose::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const { + concurrency::ThreadPool* thread_pool = context->GetOperatorThreadPool(); + + size_t num_inputs = OpKernel::Node().InputDefs().size(); + ConvTransposeAttributes::Prepare p; + bool has_bias = dynamic_padding ? num_inputs == 4 : num_inputs == 3; + ORT_RETURN_IF_ERROR(conv_transpose_attrs_.PrepareForCompute( + context, has_bias, p, dynamic_padding, transposed_filter_ ? &filter_shape_ : nullptr)); + + // Bail out early if one of the dimensions is zero. + if (p.Y->Shape().Size() == 0) { + return Status::OK(); + } + + const int64_t input_image_size = p.input_shape.Size(); + const int64_t X_offset = p.num_input_channels / conv_transpose_attrs_.group * input_image_size; + const int64_t Y_offset = p.Y->Shape().Size() / p.Y->Shape()[0] / conv_transpose_attrs_.group; + const int64_t W_offset = (p.F ? p.F->Shape().Size() : filter_shape_.Size()) / conv_transpose_attrs_.group; + const int64_t kernel_size = TensorShape(p.kernel_shape).Size(); + const int64_t kernel_dim = p.num_output_channels / conv_transpose_attrs_.group * kernel_size; + const int64_t output_size = (p.Y->Shape().Slice(2)).Size(); + + AllocatorPtr alloc; + ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&alloc)); + + const int64_t col_buffer_size = kernel_dim * p.input_shape.Size(); + auto col_data = alloc->Alloc(SafeInt(sizeof(float)) * col_buffer_size); + BufferUniquePtr col_buffer(col_data, BufferDeleter(alloc)); + float* col_buffer_data = static_cast(col_buffer.get()); + + const float* Xdata = p.X->template Data(); + const float* filter_data = p.F ? p.F->template Data() : static_cast(transposed_filter_.get()); + float* Ydata = p.Y->template MutableData(); + TensorShape output_shape = p.Y->Shape().Slice(2); + + for (auto image_id = 0; image_id < p.N; ++image_id) { + for (int group_id = 0; group_id < conv_transpose_attrs_.group; ++group_id) { + // Weight term + math::Gemm( + p.F ? CblasTrans : CblasNoTrans, + CblasNoTrans, + kernel_dim, + input_image_size, + p.num_input_channels / conv_transpose_attrs_.group, + 1, + filter_data + group_id * W_offset, + Xdata + group_id * X_offset, + 0, + col_buffer_data, + thread_pool); + + if (p.X->Shape().NumDimensions() == 4) { + math::Col2im( + col_buffer_data, + p.num_output_channels / conv_transpose_attrs_.group, + p.Y->Shape()[2], + p.Y->Shape()[3], + p.kernel_shape[0], + p.kernel_shape[1], + p.dilations[0], + p.dilations[1], + p.pads[0], + p.pads[1], + p.pads[2], + p.pads[3], + p.strides[0], + p.strides[1], + Ydata + group_id * Y_offset, + &CPUMathUtil::Instance()); + } else { + math::Col2imNd( + col_buffer_data, + output_shape.GetDims().data(), + p.input_shape.GetDims().data(), + kernel_dim, + Y_offset, + p.kernel_shape.data(), + p.strides.data(), + p.dilations.data(), + p.pads.data(), + static_cast(p.kernel_shape.size()), + Ydata + group_id * Y_offset, + &CPUMathUtil::Instance()); + } + } + + if (p.B != nullptr) { + auto Ymatrix = EigenMatrixMap(Ydata, output_size, p.num_output_channels); + auto Bvec = ConstEigenVectorMap(p.B->template Data(), p.num_output_channels); + Ymatrix.rowwise() += Bvec.transpose(); + } + + Xdata += X_offset * conv_transpose_attrs_.group; + Ydata += Y_offset * conv_transpose_attrs_.group; } return Status::OK(); diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose.h b/onnxruntime/core/providers/cpu/nn/conv_transpose.h index af2b26bf73..6de3499e62 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose.h +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose.h @@ -27,6 +27,8 @@ class ConvTranspose : public OpKernel { public: ConvTranspose(const OpKernelInfo& info) : OpKernel(info), conv_transpose_attrs_(info) {} + Status PrePack(const Tensor& tensor, int input_idx, bool& is_packed) override; + Status Compute(OpKernelContext* context) const override; protected: @@ -34,6 +36,10 @@ class ConvTranspose : public OpKernel { private: ConvTransposeAttributes conv_transpose_attrs_; + + // for pre-packing usage + TensorShape filter_shape_; + BufferUniquePtr transposed_filter_; }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h b/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h index 07a36ccccc..23d34d4675 100644 --- a/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h +++ b/onnxruntime/core/providers/cpu/nn/conv_transpose_attributes.h @@ -43,16 +43,18 @@ struct ConvTransposeAttributes : public ConvAttributes { std::vector strides; }; - Status PrepareForCompute(OpKernelContext* context, bool has_bias, Prepare& p, bool dynamic_padding = false) const { + Status PrepareForCompute(OpKernelContext* context, bool has_bias, Prepare& p, + bool dynamic_padding = false, const TensorShape* filter_shape = nullptr) const { const Tensor* X = context->Input(0); - const Tensor* F = context->Input(1); + const Tensor* F = (filter_shape != nullptr) ? nullptr : context->Input(1); + const TensorShape& F_Shape = (filter_shape != nullptr) ? *filter_shape : F->Shape(); const Tensor* Pads = dynamic_padding ? context->Input(2) : nullptr; const Tensor* B = has_bias ? (dynamic_padding ? context->Input(3) : context->Input(2)) : nullptr; const TensorShape& input_shape = X->Shape().Slice(2); const int64_t num_input_channels = X->Shape()[1]; const int64_t N = X->Shape()[0]; - const int64_t num_output_channels_multiplier = F->Shape()[1]; + const int64_t num_output_channels_multiplier = F_Shape[1]; const int64_t num_output_channels = num_output_channels_multiplier * group; // input validations @@ -61,15 +63,15 @@ struct ConvTransposeAttributes : public ConvAttributes { " group: ", group); } - if (X->Shape().NumDimensions() != F->Shape().NumDimensions()) { + if (X->Shape().NumDimensions() != F_Shape.NumDimensions()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "X num_dims does not match W num_dims.", " X: ", X->Shape().ToString().c_str(), - " W: ", F->Shape().ToString().c_str()); + " W: ", F_Shape.ToString().c_str()); } - if (F->Shape()[0] != num_input_channels) { + if (F_Shape[0] != num_input_channels) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "filter number not equal to input channel number.", - " filter_number: ", F->Shape()[0], + " filter_number: ", F_Shape[0], " num_input_channels: ", num_input_channels); } @@ -83,7 +85,7 @@ struct ConvTransposeAttributes : public ConvAttributes { } std::vector kernel_shape; - ORT_RETURN_IF_ERROR(ComputeKernelShape(F->Shape(), kernel_shape)); + ORT_RETURN_IF_ERROR(ComputeKernelShape(F_Shape, kernel_shape)); std::vector local_output_padding(output_padding); if (local_output_padding.empty()) { diff --git a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc index 42da30ab4c..324d3f7392 100644 --- a/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc +++ b/onnxruntime/test/providers/cpu/nn/conv_transpose_op_test.cc @@ -20,14 +20,15 @@ struct ConvTransposeOpAttributes { string auto_pad; }; -void TestConvTransposeOp(const ConvTransposeOpAttributes& attributes, - const vector>& inputs, - const vector>& input_shapes, - const std::initializer_list& expected_output, - const vector& expected_output_shape, - OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess, - const std::string& err_str = "", - const std::unordered_set& excluded_provider_types = {kTensorrtExecutionProvider}) { +void TestConvTransposeOpInitializer(const ConvTransposeOpAttributes& attributes, + const vector>& inputs, + const vector>& input_shapes, + const std::initializer_list& expected_output, + const vector& expected_output_shape, + bool is_filter_initializer = false, + OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess, + const std::string& err_str = "", + const std::unordered_set& excluded_provider_types = {kTensorrtExecutionProvider}) { OpTester test("ConvTranspose"); test.AddAttribute("kernel_shape", attributes.kernel_shape); test.AddAttribute("group", attributes.group); @@ -56,13 +57,31 @@ void TestConvTransposeOp(const ConvTransposeOpAttributes& attributes, ORT_ENFORCE(inputs.size() <= 3, "Our name array is only setup to handle 3 inputs"); const char* szNames[] = {"X", "W", "B"}; + bool isInitializers[] = {false, is_filter_initializer, false}; for (size_t i = 0; i < inputs.size(); i++) { - test.AddInput(szNames[i], input_shapes[i], inputs[i]); + test.AddInput(szNames[i], input_shapes[i], inputs[i], isInitializers[i]); } test.AddOutput("Y", expected_output_shape, expected_output); test.Run(expect_result, err_str, excluded_provider_types); // Disable TensorRT because weight as input is not supported } + +void TestConvTransposeOp(const ConvTransposeOpAttributes& attributes, + const vector>& inputs, + const vector>& input_shapes, + const std::initializer_list& expected_output, + const vector& expected_output_shape, + OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess, + const std::string& err_str = "", + const std::unordered_set& excluded_provider_types = {kTensorrtExecutionProvider}) { + std::unordered_set extra_exclude_openvino_for_initializer_filter = excluded_provider_types; + extra_exclude_openvino_for_initializer_filter.insert(kOpenVINOExecutionProvider); + TestConvTransposeOpInitializer(attributes, inputs, input_shapes, expected_output, expected_output_shape, + true, expect_result, err_str, extra_exclude_openvino_for_initializer_filter); + TestConvTransposeOpInitializer(attributes, inputs, input_shapes, expected_output, expected_output_shape, + false, expect_result, err_str, excluded_provider_types); +} + } // namespace TEST(ConvTransposeTest, ConvTranspose_1D) { @@ -237,6 +256,59 @@ TEST(ConvTransposeTest, ConvTranspose_2D_OutputShape_1) { TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape); } +TEST(ConvTransposeTest, ConvTranspose_2D_OutputShape_1_group_2_for_tranpose_path) { + ConvTransposeOpAttributes attrs = { + vector{3, 3}, // kernel_shape + {}, // output_padding + vector{1, 6, 4, 4}, // output_shape + vector{0, 0, 0, 0}, // pads + vector{1, 1}, // strides + vector{1, 1}, // dilations + 2, // group + "NOTSET" // auto_pad + }; + int image_size = 4 * 4; + int input_channels = 3 * 2; + int output_channels = 3; + std::vector X; + for (int i = 0; i < input_channels * image_size; i++) + X.push_back(1.0f); + std::vector W; + int kernel_size = output_channels * input_channels * 3 * 3; + for (int i = 0; i < kernel_size; i++) + W.push_back(1.0f); + + vector X_shape = {1, 6, 4, 4}; + vector W_shape = {6, 3, 3, 3}; + + vector Y_shape = {1, 6, 4, 4}; + auto expected_vals = {12.0f, 18.0f, 18.0f, 12.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 12.0f, 18.0f, 18.0f, 12.0f, // duplicate below + 12.0f, 18.0f, 18.0f, 12.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 12.0f, 18.0f, 18.0f, 12.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 18.0f, 27.0f, 27.0f, 18.0f, + 12.0f, 18.0f, 18.0f, 12.0f,}; + TestConvTransposeOp(attrs, {X, W}, {X_shape, W_shape}, expected_vals, Y_shape); +} + TEST(ConvTransposeTest, ConvTranspose_2D_OutputShape_2) { ConvTransposeOpAttributes attrs = { vector{1, 5}, // kernel_shape