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.
This commit is contained in:
Zhang Lei 2021-04-27 16:49:03 -07:00 committed by GitHub
parent d21304ceb0
commit e64e30ee0d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 264 additions and 74 deletions

View file

@ -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<float>()),
ConvTranspose<float>);
template <typename T>
Status ConvTranspose<T>::PrePack(const Tensor& /* tensor */, int /* input_idx */, bool& is_packed) {
is_packed = false;
return Status::OK();
}
template <>
Status ConvTranspose<float>::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<size_t>(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<float>() + (N * K * group_id),
((float*)packed_filter_data) + (group_id * packed_elements_per_group),
K, N);
}
is_packed = true;
}
return Status::OK();
}
template <typename T>
Status ConvTranspose<T>::Compute(OpKernelContext* context) const {
return ConvTranspose<T>::DoConvTranspose(context, false);
@ -73,25 +113,25 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
const T* Xdata = p.X->template Data<T>();
const T* filter_data = p.F->template Data<T>();
T* Ydata = p.Y->template MutableData<T>();
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<T>(
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<T>(
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<T, CPUMathUtil, StorageOrder::NCHW>(
col_buffer_data,
p.num_output_channels / conv_transpose_attrs_.group,
@ -109,37 +149,7 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
p.strides[1],
Ydata + group_id * Y_offset,
&CPUMathUtil::Instance());
}
if (p.B != nullptr) {
auto Ymatrix = EigenMatrixMap<T>(Ydata, output_size, p.num_output_channels);
auto Bvec = ConstEigenVectorMap<T>(p.B->template Data<T>(), 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<T>(
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<T, CPUMathUtil, StorageOrder::NCHW>(
col_buffer_data,
output_shape.GetDims().data(),
@ -154,16 +164,116 @@ Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_
Ydata + group_id * Y_offset,
&CPUMathUtil::Instance());
}
if (p.B != nullptr) {
auto Ymatrix = EigenMatrixMap<T>(Ydata, output_size, p.num_output_channels);
auto Bvec = ConstEigenVectorMap<T>(p.B->template Data<T>(), 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<T>(Ydata, output_size, p.num_output_channels);
auto Bvec = ConstEigenVectorMap<T>(p.B->template Data<T>(), 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<float>::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<size_t>(sizeof(float)) * col_buffer_size);
BufferUniquePtr col_buffer(col_data, BufferDeleter(alloc));
float* col_buffer_data = static_cast<float*>(col_buffer.get());
const float* Xdata = p.X->template Data<float>();
const float* filter_data = p.F ? p.F->template Data<float>() : static_cast<float*>(transposed_filter_.get());
float* Ydata = p.Y->template MutableData<float>();
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<float>(
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<float, CPUMathUtil, StorageOrder::NCHW>(
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<float, CPUMathUtil, StorageOrder::NCHW>(
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<int>(p.kernel_shape.size()),
Ydata + group_id * Y_offset,
&CPUMathUtil::Instance());
}
}
if (p.B != nullptr) {
auto Ymatrix = EigenMatrixMap<float>(Ydata, output_size, p.num_output_channels);
auto Bvec = ConstEigenVectorMap<float>(p.B->template Data<float>(), 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();

View file

@ -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

View file

@ -43,16 +43,18 @@ struct ConvTransposeAttributes : public ConvAttributes {
std::vector<int64_t> 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<Tensor>(0);
const Tensor* F = context->Input<Tensor>(1);
const Tensor* F = (filter_shape != nullptr) ? nullptr : context->Input<Tensor>(1);
const TensorShape& F_Shape = (filter_shape != nullptr) ? *filter_shape : F->Shape();
const Tensor* Pads = dynamic_padding ? context->Input<Tensor>(2) : nullptr;
const Tensor* B = has_bias ? (dynamic_padding ? context->Input<Tensor>(3) : context->Input<Tensor>(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<int64_t> kernel_shape;
ORT_RETURN_IF_ERROR(ComputeKernelShape(F->Shape(), kernel_shape));
ORT_RETURN_IF_ERROR(ComputeKernelShape(F_Shape, kernel_shape));
std::vector<int64_t> local_output_padding(output_padding);
if (local_output_padding.empty()) {

View file

@ -20,14 +20,15 @@ struct ConvTransposeOpAttributes {
string auto_pad;
};
void TestConvTransposeOp(const ConvTransposeOpAttributes& attributes,
const vector<vector<float>>& inputs,
const vector<vector<int64_t>>& input_shapes,
const std::initializer_list<float>& expected_output,
const vector<int64_t>& expected_output_shape,
OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess,
const std::string& err_str = "",
const std::unordered_set<std::string>& excluded_provider_types = {kTensorrtExecutionProvider}) {
void TestConvTransposeOpInitializer(const ConvTransposeOpAttributes& attributes,
const vector<vector<float>>& inputs,
const vector<vector<int64_t>>& input_shapes,
const std::initializer_list<float>& expected_output,
const vector<int64_t>& expected_output_shape,
bool is_filter_initializer = false,
OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess,
const std::string& err_str = "",
const std::unordered_set<std::string>& 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<float>(szNames[i], input_shapes[i], inputs[i]);
test.AddInput<float>(szNames[i], input_shapes[i], inputs[i], isInitializers[i]);
}
test.AddOutput<float>("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<vector<float>>& inputs,
const vector<vector<int64_t>>& input_shapes,
const std::initializer_list<float>& expected_output,
const vector<int64_t>& expected_output_shape,
OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess,
const std::string& err_str = "",
const std::unordered_set<std::string>& excluded_provider_types = {kTensorrtExecutionProvider}) {
std::unordered_set<std::string> 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<int64_t>{3, 3}, // kernel_shape
{}, // output_padding
vector<int64_t>{1, 6, 4, 4}, // output_shape
vector<int64_t>{0, 0, 0, 0}, // pads
vector<int64_t>{1, 1}, // strides
vector<int64_t>{1, 1}, // dilations
2, // group
"NOTSET" // auto_pad
};
int image_size = 4 * 4;
int input_channels = 3 * 2;
int output_channels = 3;
std::vector<float> X;
for (int i = 0; i < input_channels * image_size; i++)
X.push_back(1.0f);
std::vector<float> W;
int kernel_size = output_channels * input_channels * 3 * 3;
for (int i = 0; i < kernel_size; i++)
W.push_back(1.0f);
vector<int64_t> X_shape = {1, 6, 4, 4};
vector<int64_t> W_shape = {6, 3, 3, 3};
vector<int64_t> 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<int64_t>{1, 5}, // kernel_shape