From 2f7f83c655a209ae8361b76de157a961a30a954e Mon Sep 17 00:00:00 2001 From: Yulong Wang Date: Tue, 17 Sep 2019 12:11:54 +0800 Subject: [PATCH] ROIAlign: add cuda kernel implementation (#1823) --- .../cpu/object_detection/roialign.cc | 16 +- .../providers/cpu/object_detection/roialign.h | 39 ++- .../providers/cuda/cuda_execution_provider.cc | 4 + .../cuda/object_detection/roialign.cc | 75 ++++++ .../cuda/object_detection/roialign.h | 24 ++ .../cuda/object_detection/roialign_impl.cu | 230 ++++++++++++++++++ .../cuda/object_detection/roialign_impl.h | 31 +++ 7 files changed, 399 insertions(+), 20 deletions(-) create mode 100644 onnxruntime/core/providers/cuda/object_detection/roialign.cc create mode 100644 onnxruntime/core/providers/cuda/object_detection/roialign.h create mode 100644 onnxruntime/core/providers/cuda/object_detection/roialign_impl.cu create mode 100644 onnxruntime/core/providers/cuda/object_detection/roialign_impl.h diff --git a/onnxruntime/core/providers/cpu/object_detection/roialign.cc b/onnxruntime/core/providers/cpu/object_detection/roialign.cc index 9eeec05144..9975eb98d5 100644 --- a/onnxruntime/core/providers/cpu/object_detection/roialign.cc +++ b/onnxruntime/core/providers/cpu/object_detection/roialign.cc @@ -168,7 +168,7 @@ void RoiAlignForward( const T* bottom_rois, int64_t num_roi_cols, T* top_data, - const std::string& mode, + RoiAlignMode mode, const int64_t* batch_indices_ptr, const ThreadPool* ttp) { int64_t n_rois = nthreads / channels / pooled_width / pooled_height; @@ -231,7 +231,7 @@ void RoiAlignForward( int64_t index = index_n_c + ph * pooled_width + pw; T output_val = 0.; - if (mode == "avg") { // avg pooling + if (mode == RoiAlignMode::avg) { // avg pooling for (int64_t iy = 0; iy < roi_bin_grid_h; iy++) { for (int64_t ix = 0; ix < roi_bin_grid_w; ix++) { PreCalc pc = pre_calc[pre_calc_index]; @@ -334,22 +334,22 @@ Status RoiAlign::Compute(OpKernelContext* context) const { return status; } - auto& Y = *context->Output(0, {num_rois, x_dims[1], output_height_, output_width_}); + auto& Y = *context->Output(0, {num_rois, x_dims[1], this->output_height_, this->output_width_}); int64_t output_size = Y.Shape().Size(); RoiAlignForward( output_size, // num threads X_ptr->Data(), - spatial_scale_, + this->spatial_scale_, x_dims[1], // num channels x_dims[2], // height x_dims[3], // width - output_height_, - output_width_, - sampling_ratio_, + this->output_height_, + this->output_width_, + this->sampling_ratio_, rois_ptr->Data(), num_roi_cols, Y.template MutableData(), - mode_, + this->mode_, batch_indices_ptr->Data(), static_cast(context)->GetOperatorThreadPool()); diff --git a/onnxruntime/core/providers/cpu/object_detection/roialign.h b/onnxruntime/core/providers/cpu/object_detection/roialign.h index 3a0f277720..ac7ab0409e 100644 --- a/onnxruntime/core/providers/cpu/object_detection/roialign.h +++ b/onnxruntime/core/providers/cpu/object_detection/roialign.h @@ -10,18 +10,23 @@ namespace onnxruntime { Status CheckROIAlignValidInput(const Tensor* X_ptr, const Tensor* rois_ptr, const Tensor* batch_indices_ptr); +enum struct RoiAlignMode { + avg = 0, + max +}; + template -class RoiAlign final : public OpKernel { +class RoiAlignBase { public: - explicit RoiAlign(const OpKernelInfo& info) : OpKernel(info) { + explicit RoiAlignBase(const OpKernelInfo& info) { // mode - std::string mode_tmp; - if (info.GetAttr("mode", &mode_tmp).IsOK()) { - mode_ = mode_tmp; - std::transform(mode_.begin(), mode_.end(), mode_.begin(), [](auto& i) { return static_cast(::tolower(i)); }); - if (mode_ != "avg" && mode_ != "max") { - ORT_THROW("Invalid mode of value ", mode_, " specified. It should be either avg or max"); + std::string mode; + if (info.GetAttr("mode", &mode).IsOK()) { + std::transform(mode.begin(), mode.end(), mode.begin(), [](auto& i) { return static_cast(::tolower(i)); }); + if (mode != "avg" && mode != "max") { + ORT_THROW("Invalid mode of value ", mode, " specified. It should be either avg or max"); } + mode_ = mode == "avg" ? RoiAlignMode::avg : RoiAlignMode::max; } // output_height @@ -50,15 +55,25 @@ class RoiAlign final : public OpKernel { } } - Status Compute(OpKernelContext* context) const override; - - private: - std::string mode_{"avg"}; + protected: + RoiAlignMode mode_{RoiAlignMode::avg}; int64_t output_height_{1}; int64_t output_width_{1}; int64_t sampling_ratio_{0}; float spatial_scale_{1.0f}; + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RoiAlignBase); +}; + +template +class RoiAlign final : public OpKernel, public RoiAlignBase { + public: + explicit RoiAlign(const OpKernelInfo& info) : OpKernel(info), RoiAlignBase(info) {} + + Status Compute(OpKernelContext* context) const override; + + private: ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RoiAlign); }; } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index e26db9e070..38082c0676 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc @@ -537,6 +537,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 9, MLFloat16, Less); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, Dropout); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, float, RoiAlign); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, double, RoiAlign); static void RegisterCudaKernels(KernelRegistry& kernel_registry) { static const BuildKernelCreateInfoFn function_table[] = { @@ -874,6 +876,8 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, }; for (auto& function_table_entry : function_table) { diff --git a/onnxruntime/core/providers/cuda/object_detection/roialign.cc b/onnxruntime/core/providers/cuda/object_detection/roialign.cc new file mode 100644 index 0000000000..cd63b54cb8 --- /dev/null +++ b/onnxruntime/core/providers/cuda/object_detection/roialign.cc @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "roialign.h" +#include "roialign_impl.h" + +namespace onnxruntime { +namespace cuda { + +#define REGISTER_KERNEL_TYPED(T) \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + RoiAlign, \ + kOnnxDomain, \ + 10, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()), \ + RoiAlign); + +template +Status RoiAlign::ComputeInternal(OpKernelContext* context) const { + // X + const auto* X_ptr = context->Input(0); + // rois + const auto* rois_ptr = context->Input(1); + // batch indices + const auto* batch_indices_ptr = context->Input(2); + + const auto& x_dims = X_ptr->Shape(); + const auto& rois_dims = rois_ptr->Shape(); + const auto& batch_indices_dims = batch_indices_ptr->Shape(); + + auto num_rois = batch_indices_dims[0]; + auto num_roi_cols = rois_dims[1]; + + auto status = CheckROIAlignValidInput(X_ptr, rois_ptr, batch_indices_ptr); + if (status != Status::OK()) { + return status; + } + + auto& Y = *context->Output(0, {num_rois, x_dims[1], this->output_height_, this->output_width_}); + int64_t output_size = Y.Shape().Size(); + + RoiAlignImpl( + output_size, // num threads + reinterpret_cast::MappedType*>(X_ptr->template Data()), + ToCudaType::FromFloat(this->spatial_scale_), + x_dims[1], // num channels + x_dims[2], // height + x_dims[3], // width + this->output_height_, + this->output_width_, + this->sampling_ratio_, + reinterpret_cast::MappedType*>(rois_ptr->template Data()), + num_roi_cols, + reinterpret_cast::MappedType*>(Y.template MutableData()), + this->mode_ == RoiAlignMode::avg, + batch_indices_ptr->template Data() + ); + + return Status::OK(); +} + +#define SPECIALIZED_COMPUTE(T) \ + REGISTER_KERNEL_TYPED(T) \ + template Status RoiAlign::ComputeInternal(OpKernelContext* ctx) const; + +SPECIALIZED_COMPUTE(float) +SPECIALIZED_COMPUTE(double) +//SPECIALIZED_COMPUTE(MLFloat16) + +} // namespace cuda +}; // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/object_detection/roialign.h b/onnxruntime/core/providers/cuda/object_detection/roialign.h new file mode 100644 index 0000000000..fdd0f95ccf --- /dev/null +++ b/onnxruntime/core/providers/cuda/object_detection/roialign.h @@ -0,0 +1,24 @@ +// 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/providers/cuda/cuda_common.h" +#include "core/providers/cpu/object_detection/roialign.h" + +namespace onnxruntime { +namespace cuda { + +template +struct RoiAlign final : CudaKernel, RoiAlignBase { + RoiAlign(const OpKernelInfo& info) : CudaKernel(info), RoiAlignBase(info) {} + + Status ComputeInternal(OpKernelContext* context) const override; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(RoiAlign); +}; +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/object_detection/roialign_impl.cu b/onnxruntime/core/providers/cuda/object_detection/roialign_impl.cu new file mode 100644 index 0000000000..45a35b291e --- /dev/null +++ b/onnxruntime/core/providers/cuda/object_detection/roialign_impl.cu @@ -0,0 +1,230 @@ +/** +* Copyright (c) 2016-present, Facebook, Inc. +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ +/* Modifications Copyright (c) Microsoft. */ + +#include "roialign_impl.h" +#include "core/providers/cuda/cu_inc/common.cuh" + +namespace onnxruntime { +namespace cuda { + +template +__device__ T bilinear_interpolate( + const T* bottom_data, + const int height, + const int width, + T y, + T x, + const bool is_mode_avg, + const int index /* index for debug only*/) { + // deal with cases that inverse elements are out of feature map boundary + if (y < -1.0 || y > height || x < -1.0 || x > width) { + // empty + return 0; + } + + if (y <= 0) { + y = 0; + } + if (x <= 0) { + x = 0; + } + + int y_low = (int)y; + int x_low = (int)x; + int y_high; + int x_high; + + if (y_low >= height - 1) { + y_high = y_low = height - 1; + y = (T)y_low; + } else { + y_high = y_low + 1; + } + + if (x_low >= width - 1) { + x_high = x_low = width - 1; + x = (T)x_low; + } else { + x_high = x_low + 1; + } + + T ly = y - y_low; + T lx = x - x_low; + T hy = 1. - ly, hx = 1. - lx; + // do bilinear interpolation + T v1 = bottom_data[y_low * width + x_low]; + T v2 = bottom_data[y_low * width + x_high]; + T v3 = bottom_data[y_high * width + x_low]; + T v4 = bottom_data[y_high * width + x_high]; + T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; + + T val = is_mode_avg + ? (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4) // mode Avg + : max(max(max(w1 * v1, w2 * v2), w3 * v3), w4 * v4); // mode Max + + return val; +} + +template +__global__ void RoIAlignForward( + const int64_t nthreads, + const T* bottom_data, + const T spatial_scale, + const int64_t channels, + const int64_t height, + const int64_t width, + const int64_t pooled_height, + const int64_t pooled_width, + const int64_t sampling_ratio, + const T* bottom_rois, + int64_t roi_cols, + T* top_data, + const bool is_mode_avg, + const int64_t* batch_indices_ptr) { + for (size_t index = blockIdx.x * blockDim.x + threadIdx.x; index < nthreads; index += blockDim.x * gridDim.x) { + // (n, c, ph, pw) is an element in the pooled output + int pw = index % pooled_width; + int ph = (index / pooled_width) % pooled_height; + int c = (index / pooled_width / pooled_height) % channels; + int n = index / pooled_width / pooled_height / channels; + + // RoI could have 4 or 5 columns + const T* offset_bottom_rois = bottom_rois + n * roi_cols; + const auto roi_batch_ind = batch_indices_ptr[n]; + + bool continuous_coordinate = false; + // Do not using rounding; this implementation detail is critical + T roi_offset = continuous_coordinate ? T(0.5) : T(0); + T roi_start_w = offset_bottom_rois[0] * spatial_scale - roi_offset; + T roi_start_h = offset_bottom_rois[1] * spatial_scale - roi_offset; + T roi_end_w = offset_bottom_rois[2] * spatial_scale - roi_offset; + T roi_end_h = offset_bottom_rois[3] * spatial_scale - roi_offset; + + T roi_width = roi_end_w - roi_start_w; + T roi_height = roi_end_h - roi_start_h; + if (!continuous_coordinate) { // backward compatiblity + // Force malformed ROIs to be 1x1 + roi_width = max(roi_width, (T)1.); + roi_height = max(roi_height, (T)1.); + } + T bin_size_h = static_cast(roi_height) / static_cast(pooled_height); + T bin_size_w = static_cast(roi_width) / static_cast(pooled_width); + + const T* offset_bottom_data = + bottom_data + static_cast((roi_batch_ind * channels + c) * height * width); + + // We use roi_bin_grid to sample the grid and mimic integral + int roi_bin_grid_h = (sampling_ratio > 0) + ? sampling_ratio + : ceil(roi_height / pooled_height); // e.g., = 2 + int roi_bin_grid_w = + (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / pooled_width); + + // We do average (integral) pooling inside a bin + const T count = roi_bin_grid_h * roi_bin_grid_w; // e.g. = 4 + + T output_val = 0.; + bool max_flag = false; + for (int iy = 0; iy < roi_bin_grid_h; iy++) // e.g., iy = 0, 1 + { + const T y = roi_start_h + ph * bin_size_h + + static_cast(iy + .5f) * bin_size_h / + static_cast(roi_bin_grid_h); // e.g., 0.5, 1.5 + for (int ix = 0; ix < roi_bin_grid_w; ix++) { + const T x = roi_start_w + pw * bin_size_w + + static_cast(ix + .5f) * bin_size_w / + static_cast(roi_bin_grid_w); + + T val = bilinear_interpolate( + offset_bottom_data, height, width, y, x, is_mode_avg, index); + + if (is_mode_avg) { + output_val += val; + } else { + if (!max_flag) { + output_val = val; + max_flag = true; + } else { + output_val = max(output_val, val); + } + } + } + } + if (is_mode_avg) { + output_val /= count; + } + + top_data[index] = output_val; + } +} + +template +void RoiAlignImpl( + const int64_t nthreads, + const T* bottom_data, + const T spatial_scale, + const int64_t channels, + const int64_t height, + const int64_t width, + const int64_t pooled_height, + const int64_t pooled_width, + const int64_t sampling_ratio, + const T* bottom_rois, + int64_t roi_cols, + T* top_data, + const bool is_mode_avg, + const int64_t* batch_indices_ptr) { + int blocksPerGrid = (int)(ceil(static_cast(nthreads) / GridDim::maxThreadsPerBlock)); + RoIAlignForward<<>>( + nthreads, + bottom_data, + spatial_scale, + channels, + height, + width, + pooled_height, + pooled_width, + sampling_ratio, + bottom_rois, + roi_cols, + top_data, + is_mode_avg, + batch_indices_ptr); +} + +#define SPECIALIZED_IMPL(T) \ + template void RoiAlignImpl( \ + const int64_t nthreads, \ + const T* bottom_data, \ + const T spatial_scale, \ + const int64_t channels, \ + const int64_t height, \ + const int64_t width, \ + const int64_t pooled_height, \ + const int64_t pooled_width, \ + const int64_t sampling_ratio, \ + const T* bottom_rois, \ + int64_t roi_cols, \ + T* top_data, \ + const bool is_mode_avg, \ + const int64_t* batch_indices_ptr); + +SPECIALIZED_IMPL(float) +SPECIALIZED_IMPL(double) + +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/object_detection/roialign_impl.h b/onnxruntime/core/providers/cuda/object_detection/roialign_impl.h new file mode 100644 index 0000000000..712771bdf6 --- /dev/null +++ b/onnxruntime/core/providers/cuda/object_detection/roialign_impl.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once +#include +#include "core/providers/cuda/shared_inc/cuda_utils.h" +#include "core/framework/data_types.h" +#include "core/common/common.h" + +namespace onnxruntime { +namespace cuda { + +template +void RoiAlignImpl( + const int64_t nthreads, + const T* bottom_data, + const T spatial_scale, + const int64_t channels, + const int64_t height, + const int64_t width, + const int64_t pooled_height, + const int64_t pooled_width, + const int64_t sampling_ratio, + const T* bottom_rois, + int64_t roi_cols, + T* top_data, + const bool is_mode_avg, + const int64_t* batch_indices_ptr); + +} // namespace cuda +} // namespace onnxruntime