Upsample opset 9 cuda implementation (#330)

This commit is contained in:
Hector Li 2019-01-14 23:04:46 -08:00 committed by GitHub
parent 0efc48a11a
commit 779123cf55
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 57 additions and 40 deletions

View file

@ -17,19 +17,32 @@ enum UpsampleMode {
class UpsampleBase {
protected:
UpsampleBase(OpKernelInfo info) {
UpsampleBase(OpKernelInfo info): scales_cached_(false) {
std::string mode;
ORT_ENFORCE(info.GetAttr<std::string>("mode", &mode).IsOK());
mode_ = StringToUpsampleMode(mode);
if (info.GetInputCount() == 1) {
auto input_count = info.GetInputCount();
if (input_count == 1) {
ORT_ENFORCE(info.GetAttrs<float>("scales", scales_).IsOK());
ScalesValidation(scales_, mode_);
}
// opset 9
if (input_count > 1) {
const Tensor* scale;
bool get_scale = info.TryGetConstantInput(1, &scale);
if (get_scale) {
ParseScalesData(scale, scales_);
scales_cached_ = true;
}
}
}
UpsampleMode mode_;
std::vector<float> scales_;
bool scales_cached_;
UpsampleMode StringToUpsampleMode(const std::string& mode) {
if (strcmp(mode.c_str(), UpsampleModeNN) == 0) {
@ -53,28 +66,7 @@ class UpsampleBase {
"Upsample: linear mode upsample only support bilinear, the first 2 scales should be 1.");
}
}
};
template <typename T>
class Upsample : public UpsampleBase, public OpKernel {
public:
Upsample(OpKernelInfo info) : UpsampleBase(info), OpKernel(info), scales_cached_(false) {
if (info.GetInputCount() > 1) {
const Tensor* scale;
bool get_scale = info.TryGetConstantInput(1, &scale);
if (get_scale) {
ParseScalesData(scale, scales_);
scales_cached_ = true;
}
}
}
Status Compute(OpKernelContext* context) const override;
Status BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const;
private:
void ParseScalesData(const Tensor* scale, std::vector<float>& scales) const {
const float* scale_data = scale->template Data<float>();
int64_t scales_size = scale->Shape().Size();
@ -82,9 +74,17 @@ private:
memcpy(scales.data(), scale_data, scales_size * sizeof(float));
ScalesValidation(scales, mode_);
}
};
private:
bool scales_cached_;
template <typename T>
class Upsample : public UpsampleBase, public OpKernel {
public:
Upsample(OpKernelInfo info) : UpsampleBase(info), OpKernel(info) {
}
Status Compute(OpKernelContext* context) const override;
Status BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const;
};
} // namespace onnxruntime

View file

@ -498,10 +498,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int64_t, DynamicSlice);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, Compress);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, float, Upsample);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, double, Upsample);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, MLFloat16, Upsample);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, int32_t, Upsample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, float, Upsample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, double, Upsample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, MLFloat16, Upsample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, int32_t, Upsample);
static void RegisterCudaKernels(std::function<void(KernelCreateInfo&&)> fn) {
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, MemcpyFromHost)>());
@ -761,10 +761,10 @@ static void RegisterCudaKernels(std::function<void(KernelCreateInfo&&)> fn) {
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int32_t, DynamicSlice)>());
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, int64_t, DynamicSlice)>());
fn(BuildKernel<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, Compress)>());
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, float, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, double, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, MLFloat16, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, int32_t, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, float, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, double, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, MLFloat16, Upsample)>());
fn(BuildKernel<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 7, 9, int32_t, Upsample)>());
}
} // namespace cuda

View file

@ -11,13 +11,15 @@ namespace onnxruntime {
namespace cuda {
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \
Upsample, \
kOnnxDomain, \
7, \
9, \
T, \
kCudaExecutionProvider, \
KernelDefBuilder() \
.InputMemoryType<OrtMemTypeCPUInput>(1) \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
Upsample<T>);
@ -27,7 +29,7 @@ REGISTER_KERNEL_TYPED(MLFloat16)
REGISTER_KERNEL_TYPED(int32_t)
template <typename T>
Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
Status Upsample<T>::BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const {
const Tensor* X = context->Input<Tensor>(0);
ORT_ENFORCE(nullptr != X);
const std::vector<int64_t>& X_dims = X->Shape().GetDims();
@ -35,12 +37,12 @@ Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
if (rank == 0)
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Upsample: input tensor cannot be scalar.");
if (rank != scales_.size())
if (rank != scales.size())
return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Upsample: input tensor's dimension does not match the scales.");
std::vector<int64_t> Y_dims;
for (std::size_t i = 0; i < rank; i++) {
Y_dims.push_back(static_cast<int64_t>(scales_[i] * X_dims[i]));
Y_dims.push_back(static_cast<int64_t>(scales[i] * X_dims[i]));
}
Tensor* Y = context->Output(0, Y_dims);
typedef typename ToCudaType<T>::MappedType CudaT;
@ -55,14 +57,13 @@ Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
CudaAsyncBuffer<fast_divmod> output_div_pitches(this, device_id, rank);
gsl::span<fast_divmod> div_strides_span = output_div_pitches.CpuSpan();
CudaAsyncBuffer<fast_divmod> scales_div(this, device_id, rank);
gsl::span<fast_divmod> scales_div_span = scales_div.CpuSpan();
for (int i = 0; i < rank; ++i) {
input_stride_span[i] = input_pitches[i];
div_strides_span[i] = fast_divmod(gsl::narrow_cast<int>(output_pitches[i]));
scales_div_span[i] = fast_divmod(gsl::narrow_cast<int>(ceil(scales_[i])));
scales_div_span[i] = fast_divmod(gsl::narrow_cast<int>(ceil(scales[i])));
}
input_strides.CopyToGpu();
output_div_pitches.CopyToGpu();
@ -88,5 +89,21 @@ Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
return Status::OK();
}
template <typename T>
Status Upsample<T>::ComputeInternal(OpKernelContext* context) const {
// Opset 7
if (OpKernel::Node().InputDefs().size() == 1 || scales_cached_) {
return BaseCompute(context, scales_);
}
// Opset 9
const Tensor* scales = context->Input<Tensor>(1);
ORT_ENFORCE(scales != nullptr);
int64_t scales_size = scales->Shape().Size();
std::vector<float> scales_arrary(scales_size);
ParseScalesData(scales, scales_arrary);
return BaseCompute(context, scales_arrary);
}
} // namespace cuda
} // namespace onnxruntime

View file

@ -17,6 +17,7 @@ class Upsample : public UpsampleBase, public CudaKernel {
}
Status ComputeInternal(OpKernelContext* context) const override;
Status BaseCompute(OpKernelContext* context, const std::vector<float>& scales) const;
};
} // namespace cuda

View file

@ -304,7 +304,6 @@ int real_main(int argc, char* argv[]) {
{"operator_rnn_single_layer", "disable reason"},
{"prelu_broadcast", "disable reason"},
{"prelu_example", "disable reason"},
{"upsample_nearest", "opset 9 not supported yet"},
{"sinh_example", "opset 9 not supported yet"},
{"cosh_example", "opset 9 not supported yet"},
{"asinh_example", "opset 9 not supported yet"},