Adding custom op ConvTransposeWithDynamicPads. (#638)

* Adding custom op ConvTransposeWithDynamicPads.

* Adding custom op ConvTransposeWithDynamicPads.

* adding cuda kernels

* fix a bug

* fix build issue.

* Integrate PR comments.
This commit is contained in:
Du Li 2019-05-31 11:48:43 -07:00 committed by GitHub
parent 6c408c3a75
commit 05110a6558
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
16 changed files with 248 additions and 31 deletions

View file

@ -8,9 +8,11 @@ file(GLOB_RECURSE onnxruntime_providers_srcs CONFIGURE_DEPENDS
file(GLOB_RECURSE onnxruntime_contrib_ops_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/contrib_ops/*.h"
"${ONNXRUNTIME_ROOT}/contrib_ops/*.cc"
"${ONNXRUNTIME_ROOT}/contrib_ops/contrib_kernels.cc"
"${ONNXRUNTIME_ROOT}/contrib_ops/cpu/*.h"
"${ONNXRUNTIME_ROOT}/contrib_ops/cpu/*.cc"
"${ONNXRUNTIME_ROOT}/contrib_ops/cpu/attnlstm/*.h"
"${ONNXRUNTIME_ROOT}/contrib_ops/cpu/attnlstm/*.cc"
)
file(GLOB onnxruntime_providers_common_srcs CONFIGURE_DEPENDS
@ -58,12 +60,14 @@ if (onnxruntime_USE_CUDA)
file(GLOB_RECURSE onnxruntime_providers_cuda_cc_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/core/providers/cuda/*.h"
"${ONNXRUNTIME_ROOT}/core/providers/cuda/*.cc"
"${ONNXRUNTIME_ROOT}/contrib_ops/cuda/*.h"
"${ONNXRUNTIME_ROOT}/contrib_ops/cuda/*.cc"
)
file(GLOB_RECURSE onnxruntime_providers_cuda_cu_srcs CONFIGURE_DEPENDS
"${ONNXRUNTIME_ROOT}/core/providers/cuda/*.cu"
"${ONNXRUNTIME_ROOT}/core/providers/cuda/*.cuh"
)
source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_cuda_cc_srcs} ${onnxruntime_providers_cuda_cu_srcs})
source_group(TREE ${ONNXRUNTIME_ROOT} FILES ${onnxruntime_providers_cuda_cc_srcs} ${onnxruntime_providers_cuda_cu_srcs})
add_library(onnxruntime_providers_cuda ${onnxruntime_providers_cuda_cc_srcs} ${onnxruntime_providers_cuda_cu_srcs})
if (UNIX)
target_compile_options(onnxruntime_providers_cuda PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xcompiler -Wno-reorder>"

View file

@ -221,6 +221,13 @@ template <typename T>
KernelCreateInfo BuildKernelCreateInfo();
} // namespace contrib
namespace contrib {
namespace cuda {
template <typename T>
KernelCreateInfo BuildKernelCreateInfo();
} // namespace cuda
} // namespace contrib
using BuildKernelCreateInfoFn = KernelCreateInfo (*)();
// Naming convention for operator kernel classes

View file

@ -20,6 +20,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Murmu
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MaxpoolWithMask);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Pad);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Unique);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads);
// This section includes all opkernel declarations for former experimental ops which have now been removed from onnx.
// To maintain backward compatibility these are added as contrib ops.
@ -64,6 +65,8 @@ void RegisterContribKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, MaxpoolWithMask)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Pad)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Unique)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads)>,
// These ops were experimental ops in onnx domain which have been removed now. We add them here as
// contrib ops to main backward compatibility
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 1, Affine)>,

View file

@ -9,5 +9,8 @@
namespace onnxruntime {
namespace contrib {
void RegisterContribKernels(KernelRegistry& kernel_registry);
} // namespace contrib
namespace cuda {
void RegisterCudaContribKernels(KernelRegistry& kernel_registry);
}
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,16 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "conv_transpose_with_dynamic_pads.h"
namespace onnxruntime {
namespace contrib {
ONNX_CPU_OPERATOR_TYPED_MS_KERNEL(
ConvTransposeWithDynamicPads,
1,
float,
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
ConvTransposeWithDynamicPads<float>);
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,20 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/cpu/nn/conv_transpose.h"
namespace onnxruntime {
namespace contrib {
template <typename T>
class ConvTransposeWithDynamicPads : public ConvTranspose<T> {
public:
ConvTransposeWithDynamicPads(const OpKernelInfo& info) : ConvTranspose<T>(info) {}
Status Compute(OpKernelContext* context) const override {
return ConvTranspose<T>::DoConvTranspose(context, true);
}
};
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,19 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/cuda/conv_transpose_with_dynamic_pads.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
ONNX_OPERATOR_TYPED_KERNEL_EX(
ConvTransposeWithDynamicPads,
kMSDomain,
1,
float,
kCudaExecutionProvider,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()).InputMemoryType<OrtMemTypeCPUInput>(2),
ConvTransposeWithDynamicPads<float>);
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

View file

@ -0,0 +1,23 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/cuda/nn/conv_transpose.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
template <typename T>
class ConvTransposeWithDynamicPads : public ::onnxruntime::cuda::ConvTranspose<T> {
public:
ConvTransposeWithDynamicPads(const OpKernelInfo& info) : ::onnxruntime::cuda::ConvTranspose<T>(info) {}
Status ComputeInternal(OpKernelContext* context) const override {
return ::onnxruntime::cuda::ConvTranspose<T>::DoConvTranspose(context, true);
}
};
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

View file

@ -530,6 +530,61 @@ Sample echo operator.)DOC");
ONNX_NAMESPACE::convPoolShapeInference(ctx, false, true, 0, 1);
});
ONNX_CONTRIB_OPERATOR_SCHEMA(ConvTransposeWithDynamicPads)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetDoc(R"DOC()DOC")
.Attr(
"kernel_shape",
"",
AttributeProto::INTS,
OPTIONAL)
.Attr("output_padding",
"",
AttributeProto::INTS,
OPTIONAL)
.Attr(
"dilations",
"",
AttributeProto::INTS,
OPTIONAL)
.Attr(
"strides",
"",
AttributeProto::INTS,
OPTIONAL)
.Attr(
"auto_pad",
"",
AttributeProto::STRING,
std::string("NOTSET"))
.Attr(
"group",
"",
AttributeProto::INT,
static_cast<int64_t>(1))
.Input(
0,
"X",
"",
"T")
.Input(
1,
"W",
"",
"T")
.Input(2, "Pads", "", "tensor(int64)", OpSchema::Optional)
.Input(3, "B", "", "T", OpSchema::Optional)
.Output(
0,
"Y",
"",
"T")
.TypeConstraint("T", {"tensor(float16)", "tensor(float)", "tensor(double)"}, "Constrain input and output types to float tensors")
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
});
ONNX_CONTRIB_OPERATOR_SCHEMA(FusedConv)
.SetDomain(kMSDomain)
.SinceVersion(1)
@ -1192,6 +1247,6 @@ Example 4:
// register internal ops
RegisterInternalSchemas();
#endif
}
} // namespace contrib
} // namespace contrib
} // namespace onnxruntime

View file

@ -53,30 +53,31 @@ inline void ComputeTransposePadAndOutputShape(
}
return;
}
if (pad_type != AutoPadType::NOTSET) {
switch (pad_type) {
// We handle cases of AutoPadType::VALID and AutoPadType::SAME_UPPER/LOWER,
// the same way
case AutoPadType::VALID:
case AutoPadType::SAME_UPPER:
case AutoPadType::SAME_LOWER:
*pad_head = 0;
*pad_tail = 0;
*out_size = (in_size - 1) * stride + kernel + dilation - 1 + adj;
break;
default:
throw NotImplementedException("pad type not supported");
}
} else {
*out_size =
(in_size - 1) * stride + kernel + dilation - 1 + adj - *pad_head - *pad_tail;
if (pad_type != AutoPadType::NOTSET) {
switch (pad_type) {
// We handle cases of AutoPadType::VALID and AutoPadType::SAME_UPPER/LOWER,
// the same way
case AutoPadType::VALID:
case AutoPadType::SAME_UPPER:
case AutoPadType::SAME_LOWER:
*pad_head = 0;
*pad_tail = 0;
*out_size = (in_size - 1) * stride + kernel + dilation - 1 + adj;
break;
default:
throw NotImplementedException("pad type not supported");
}
} else {
*out_size =
(in_size - 1) * stride + kernel + dilation - 1 + adj - *pad_head - *pad_tail;
}
}
Status ConvTransposeBase::PrepareForCompute(OpKernelContext* context, bool has_bias, ConvTransposeBase::Prepare& p) const {
const auto* X = context->Input<Tensor>(0);
const auto* F = context->Input<Tensor>(1);
const Tensor* B = has_bias ? context->Input<Tensor>(2) : nullptr;
Status ConvTransposeBase::PrepareForCompute(OpKernelContext* context, bool has_bias, ConvTransposeBase::Prepare& p, bool dynamic_padding) const {
const Tensor* X = context->Input<Tensor>(0);
const Tensor* F = context->Input<Tensor>(1);
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();
// input validations
@ -129,7 +130,15 @@ Status ConvTransposeBase::PrepareForCompute(OpKernelContext* context, bool has_b
if (output_padding.empty()) {
output_padding.resize(kernel_shape.size(), 0);
}
std::vector<int64_t> pads(pads_);
std::vector<int64_t> pads;
pads.reserve(2 * (input_shape.NumDimensions() - 2));
if (dynamic_padding) {
for (int64_t i = 0; i < Pads->Shape().SizeFromDimension(0); ++i) {
pads.push_back(Pads->Data<int64_t>()[i]);
}
} else {
pads.assign(pads_.begin(), pads_.end());
}
if (pads.empty()) {
pads.resize(kernel_shape.size() * 2, 0);
}
@ -214,9 +223,15 @@ void ConvTransposeBase::ComputePadsAndOutputShape(
template <typename T>
Status ConvTranspose<T>::Compute(OpKernelContext* context) const {
return ConvTranspose<T>::DoConvTranspose(context, false);
}
template <typename T>
Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const {
size_t num_inputs = OpKernel::Node().InputDefs().size();
Prepare p;
ORT_RETURN_IF_ERROR(PrepareForCompute(context, num_inputs == 3, p));
bool has_bias = dynamic_padding ? num_inputs == 4 : num_inputs == 3;
ORT_RETURN_IF_ERROR(PrepareForCompute(context, has_bias, p, dynamic_padding));
const int64_t input_image_size = p.H * p.W;
const int64_t X_offset = p.num_input_channels / group_ * input_image_size;

View file

@ -45,7 +45,7 @@ class ConvTransposeBase : public ConvBase {
std::vector<int64_t> strides;
};
Status PrepareForCompute(OpKernelContext* context, bool has_bias, Prepare& p) const;
Status PrepareForCompute(OpKernelContext* context, bool has_bias, Prepare& p, bool dynamic_padding = false) const;
void ComputePadsAndOutputShape(TensorShape input_shape, int64_t output_channel,
const std::vector<int64_t>& kernel_shape, const std::vector<int64_t>& strides,
@ -62,6 +62,9 @@ class ConvTranspose : public OpKernel, public ConvTransposeBase {
ConvTranspose(const OpKernelInfo& info) : OpKernel(info), ConvTransposeBase(info) {}
Status Compute(OpKernelContext* context) const override;
protected:
Status DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const;
};
} // namespace onnxruntime

View file

@ -0,0 +1,17 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/contrib_kernels.h"
#include "core/graph/constants.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads);
void RegisterCudaContribKernels(KernelRegistry& kernel_registry) {
kernel_registry.Register(BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kMSDomain, 1, float, ConvTransposeWithDynamicPads)>());
}
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

View file

@ -8,6 +8,7 @@
#include "cuda_allocator.h"
#include "core/framework/kernel_registry.h"
#include "core/framework/compute_capability.h"
#include "contrib_ops/contrib_kernels.h"
using namespace onnxruntime::common;
@ -881,6 +882,8 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
std::shared_ptr<KernelRegistry> GetCudaKernelRegistry() {
std::shared_ptr<KernelRegistry> kernel_registry = std::make_shared<KernelRegistry>();
RegisterCudaKernels(*kernel_registry);
::onnxruntime::contrib::cuda::RegisterCudaContribKernels(*kernel_registry);
return kernel_registry;
}

View file

@ -22,6 +22,11 @@ REGISTER_KERNEL_TYPED(MLFloat16)
template <typename T>
Status ConvTranspose<T>::ComputeInternal(OpKernelContext* context) const {
return DoConvTranspose(context, false);
}
template <typename T>
Status ConvTranspose<T>::DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const {
typedef typename ToCudaType<T>::MappedType CudaT;
const Tensor* X = context->Input<Tensor>(0);
@ -35,7 +40,7 @@ Status ConvTranspose<T>::ComputeInternal(OpKernelContext* context) const {
auto w_data = reinterpret_cast<const CudaT*>(W->template Data<T>());
size_t num_inputs = OpKernel::Node().InputDefs().size();
bool has_bias = (num_inputs == 3);
bool has_bias = dynamic_padding ? num_inputs == 4 : num_inputs == 3;
CudaT* y_data = nullptr;
@ -54,7 +59,7 @@ Status ConvTranspose<T>::ComputeInternal(OpKernelContext* context) const {
}
Prepare p;
ORT_RETURN_IF_ERROR(PrepareForCompute(context, has_bias, p));
ORT_RETURN_IF_ERROR(PrepareForCompute(context, has_bias, p, dynamic_padding));
const auto& y_dims = p.Y->Shape().GetDims();
s_.y_dims = y_dims;
@ -143,7 +148,7 @@ Status ConvTranspose<T>::ComputeInternal(OpKernelContext* context) const {
y_data));
if (has_bias) {
const Tensor* B = context->Input<Tensor>(2);
const Tensor* B = dynamic_padding ? context->Input<Tensor>(3) : context->Input<Tensor>(2);
auto b_data = reinterpret_cast<const CudaT*>(B->template Data<T>());
CUDNN_RETURN_IF_ERROR(cudnnAddTensor(CudnnHandle(), &alpha, s_.b_tensor, b_data, &alpha, s_.y_tensor, y_data));
}

View file

@ -15,6 +15,7 @@ class ConvTranspose : public CudaKernel, public ConvTransposeBase {
public:
ConvTranspose(const OpKernelInfo& info) : CudaKernel(info), ConvTransposeBase(info){};
Status ComputeInternal(OpKernelContext* context) const override;
Status DoConvTranspose(OpKernelContext* context, bool dynamic_padding) const;
private:
mutable CudnnConvState<cudnnConvolutionBwdDataAlgoPerf_t> s_;

View file

@ -0,0 +1,23 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "gtest/gtest.h"
#include "test/providers/provider_test_utils.h"
namespace onnxruntime {
namespace test {
TEST(ContribOpTest, ConvTransposeWithDynamicPads) {
OpTester test("ConvTransposeWithDynamicPads", 1, onnxruntime::kMSDomain);
test.AddAttribute("kernel_shape", std::vector<int64_t>{3, 3});
test.AddAttribute("output_padding", std::vector<int64_t>{1, 1});
test.AddAttribute("strides", std::vector<int64_t>{2, 2});
test.AddAttribute("dilations", std::vector<int64_t>{1, 1});
test.AddInput<float>("X", {1, 1, 3, 3}, std::vector<float>{0.16857791f, -0.15161794f, 0.08540368f, 0.1820628f, -0.21746576f, 0.08245695f, 0.1431433f, -0.43156421f, 0.30591947f});
test.AddInput<float>("W", {1, 1, 3, 3}, std::vector<float>{-0.06230065f, 0.37932432f, -0.25388849f, 0.33878803f, 0.43709868f, -0.22477469f, 0.04118127f, -0.44696793f, 0.06373066f});
test.AddInput<int64_t>("Pads", {4}, std::vector<int64_t>{1, 1, 1, 1});
test.AddOutput<float>("Y", {1, 1, 6, 6}, std::vector<float>{0.07368518f, -0.08925839f, -0.06627201f, 0.06301362f, 0.03732984f, -0.01919658f, -0.00628807f, -0.02817563f, -0.01472169f, 0.04392925f, -0.00689478f, -0.01549204f, 0.07957941f, -0.11459791f, -0.09505399f, 0.07681622f, 0.03604182f, -0.01853423f, -0.0270785f, -0.00680824f, -0.06650258f, 0.08004665f, 0.07918708f, -0.0724144f, 0.06256775f, -0.17838378f, -0.18863615f, 0.20064656f, 0.133717f, -0.06876295f, -0.06398046f, -0.00864975f, 0.19289537f, -0.01490572f, -0.13673618f, 0.01949645f});
test.Run();
}
} // namespace test
} // namespace onnxruntime