diff --git a/torchvision/csrc/cpu/ROIPool_cpu.cpp b/torchvision/csrc/cpu/roi_pool_kernel.cpp similarity index 92% rename from torchvision/csrc/cpu/ROIPool_cpu.cpp rename to torchvision/csrc/cpu/roi_pool_kernel.cpp index 34da4f1d1cc..389e9c90248 100644 --- a/torchvision/csrc/cpu/ROIPool_cpu.cpp +++ b/torchvision/csrc/cpu/roi_pool_kernel.cpp @@ -1,7 +1,8 @@ -#include -#include -#include -#include +#include + +#include "roi_pool_kernel.h" + +namespace { template inline void add(T* address, const T& val) { @@ -9,7 +10,7 @@ inline void add(T* address, const T& val) { } template -void RoIPoolForward( +void roi_pool_forward_kernel_impl( const T* input, const T spatial_scale, int channels, @@ -78,7 +79,7 @@ void RoIPoolForward( } template -void RoIPoolBackward( +void roi_pool_backward_kernel_impl( const T* grad_output, const int* argmax_data, int num_rois, @@ -120,7 +121,9 @@ void RoIPoolBackward( } // num_rois } -std::tuple ROIPool_forward_cpu( +} // namespace + +std::tuple roi_pool_forward_cpu( const at::Tensor& input, const at::Tensor& rois, double spatial_scale, @@ -131,7 +134,7 @@ std::tuple ROIPool_forward_cpu( at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2}; - at::CheckedFrom c = "ROIPool_forward_cpu"; + at::CheckedFrom c = "roi_pool_forward_cpu"; at::checkAllSameType(c, {input_t, rois_t}); int num_rois = rois.size(0); @@ -151,8 +154,8 @@ std::tuple ROIPool_forward_cpu( auto input_ = input.contiguous(), rois_ = rois.contiguous(); AT_DISPATCH_FLOATING_TYPES_AND_HALF( - input.scalar_type(), "ROIPool_forward", [&] { - RoIPoolForward( + input.scalar_type(), "roi_pool_forward", [&] { + roi_pool_forward_kernel_impl( input_.data_ptr(), spatial_scale, channels, @@ -168,7 +171,7 @@ std::tuple ROIPool_forward_cpu( return std::make_tuple(output, argmax); } -at::Tensor ROIPool_backward_cpu( +at::Tensor roi_pool_backward_cpu( const at::Tensor& grad, const at::Tensor& rois, const at::Tensor& argmax, @@ -188,7 +191,7 @@ at::Tensor ROIPool_backward_cpu( at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2}; - at::CheckedFrom c = "ROIPool_backward_cpu"; + at::CheckedFrom c = "roi_pool_backward_cpu"; at::checkAllSameType(c, {grad_t, rois_t}); auto num_rois = rois.size(0); @@ -209,8 +212,8 @@ at::Tensor ROIPool_backward_cpu( auto rois_ = rois.contiguous(); AT_DISPATCH_FLOATING_TYPES_AND_HALF( - grad.scalar_type(), "ROIPool_backward", [&] { - RoIPoolBackward( + grad.scalar_type(), "roi_pool_backward", [&] { + roi_pool_backward_kernel_impl( grad.data_ptr(), argmax.data_ptr(), num_rois, diff --git a/torchvision/csrc/cpu/roi_pool_kernel.h b/torchvision/csrc/cpu/roi_pool_kernel.h new file mode 100644 index 00000000000..66fd993d5b4 --- /dev/null +++ b/torchvision/csrc/cpu/roi_pool_kernel.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include "../macros.h" + +VISION_API std::tuple roi_pool_forward_cpu( + const at::Tensor& input, + const at::Tensor& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width); + +VISION_API at::Tensor roi_pool_backward_cpu( + const at::Tensor& grad, + const at::Tensor& rois, + const at::Tensor& argmax, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width); diff --git a/torchvision/csrc/cpu/vision_cpu.h b/torchvision/csrc/cpu/vision_cpu.h index a2647c57aa5..a772fa13f01 100644 --- a/torchvision/csrc/cpu/vision_cpu.h +++ b/torchvision/csrc/cpu/vision_cpu.h @@ -3,22 +3,3 @@ #include "../macros.h" // TODO: Delete this file once all the methods are gone - -VISION_API std::tuple ROIPool_forward_cpu( - const at::Tensor& input, - const at::Tensor& rois, - double spatial_scale, - int64_t pooled_height, - int64_t pooled_width); - -VISION_API at::Tensor ROIPool_backward_cpu( - const at::Tensor& grad, - const at::Tensor& rois, - const at::Tensor& argmax, - double spatial_scale, - int64_t pooled_height, - int64_t pooled_width, - int64_t batch_size, - int64_t channels, - int64_t height, - int64_t width); diff --git a/torchvision/csrc/cuda/ROIPool_cuda.cu b/torchvision/csrc/cuda/roi_pool_kernel.cu similarity index 91% rename from torchvision/csrc/cuda/ROIPool_cuda.cu rename to torchvision/csrc/cuda/roi_pool_kernel.cu index 3131b9eea7e..c10dd0cf403 100644 --- a/torchvision/csrc/cuda/ROIPool_cuda.cu +++ b/torchvision/csrc/cuda/roi_pool_kernel.cu @@ -1,13 +1,15 @@ -#include -#include #include #include +#include #include #include "cuda_helpers.h" +#include "roi_pool_kernel.h" + +namespace { template -__global__ void RoIPoolForward( +__global__ void roi_pool_forward_kernel_impl( int nthreads, const T* input, const T spatial_scale, @@ -72,7 +74,7 @@ __global__ void RoIPoolForward( } template -__global__ void RoIPoolBackward( +__global__ void roi_pool_backward_kernel_impl( int nthreads, const T* grad_output, const int* argmax_data, @@ -115,7 +117,9 @@ __global__ void RoIPoolBackward( } } -std::tuple ROIPool_forward_cuda( +} // namespace + +std::tuple roi_pool_forward_cuda( const at::Tensor& input, const at::Tensor& rois, double spatial_scale, @@ -128,7 +132,7 @@ std::tuple ROIPool_forward_cuda( at::TensorArg input_t{input, "input", 1}, rois_t{rois, "rois", 2}; - at::CheckedFrom c = "ROIPool_forward_cuda"; + at::CheckedFrom c = "roi_pool_forward_cuda"; at::checkAllSameGPU(c, {input_t, rois_t}); at::checkAllSameType(c, {input_t, rois_t}); @@ -160,8 +164,8 @@ std::tuple ROIPool_forward_cuda( auto input_ = input.contiguous(), rois_ = rois.contiguous(); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "ROIPool_forward", [&] { - RoIPoolForward<<>>( + AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "roi_pool_forward", [&] { + roi_pool_forward_kernel_impl<<>>( output_size, input_.data_ptr(), spatial_scale, @@ -178,7 +182,7 @@ std::tuple ROIPool_forward_cuda( return std::make_tuple(output, argmax); } -at::Tensor ROIPool_backward_cuda( +at::Tensor roi_pool_backward_cuda( const at::Tensor& grad, const at::Tensor& rois, const at::Tensor& argmax, @@ -197,7 +201,7 @@ at::Tensor ROIPool_backward_cuda( at::TensorArg grad_t{grad, "grad", 1}, rois_t{rois, "rois", 2}, argmax_t{argmax, "argmax", 3}; - at::CheckedFrom c = "ROIPool_backward_cuda"; + at::CheckedFrom c = "roi_pool_backward_cuda"; at::checkAllSameGPU(c, {grad_t, rois_t, argmax_t}); at::checkAllSameType(c, {grad_t, rois_t}); @@ -228,8 +232,8 @@ at::Tensor ROIPool_backward_cuda( auto argmax_ = argmax.contiguous(), rois_ = rois.contiguous(); - AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad.scalar_type(), "ROIPool_backward", [&] { - RoIPoolBackward<<>>( + AT_DISPATCH_FLOATING_TYPES_AND_HALF(grad.scalar_type(), "roi_pool_backward", [&] { + roi_pool_backward_kernel_impl<<>>( grad.numel(), grad.data_ptr(), argmax_.data_ptr(), diff --git a/torchvision/csrc/cuda/roi_pool_kernel.h b/torchvision/csrc/cuda/roi_pool_kernel.h new file mode 100644 index 00000000000..3a99f7521bd --- /dev/null +++ b/torchvision/csrc/cuda/roi_pool_kernel.h @@ -0,0 +1,23 @@ +#pragma once + +#include +#include "../macros.h" + +VISION_API std::tuple roi_pool_forward_cuda( + const at::Tensor& input, + const at::Tensor& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width); + +VISION_API at::Tensor roi_pool_backward_cuda( + const at::Tensor& grad, + const at::Tensor& rois, + const at::Tensor& argmax, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width); diff --git a/torchvision/csrc/cuda/vision_cuda.h b/torchvision/csrc/cuda/vision_cuda.h index 1ec187c3348..a772fa13f01 100644 --- a/torchvision/csrc/cuda/vision_cuda.h +++ b/torchvision/csrc/cuda/vision_cuda.h @@ -3,22 +3,3 @@ #include "../macros.h" // TODO: Delete this file once all the methods are gone - -VISION_API std::tuple ROIPool_forward_cuda( - const at::Tensor& input, - const at::Tensor& rois, - const double spatial_scale, - const int64_t pooled_height, - const int64_t pooled_width); - -VISION_API at::Tensor ROIPool_backward_cuda( - const at::Tensor& grad, - const at::Tensor& rois, - const at::Tensor& argmax, - const double spatial_scale, - const int64_t pooled_height, - const int64_t pooled_width, - const int64_t batch_size, - const int64_t channels, - const int64_t height, - const int64_t width); diff --git a/torchvision/csrc/ROIPool.h b/torchvision/csrc/roi_pool.cpp similarity index 93% rename from torchvision/csrc/ROIPool.h rename to torchvision/csrc/roi_pool.cpp index 7950005f1bd..c8d70bd8940 100644 --- a/torchvision/csrc/ROIPool.h +++ b/torchvision/csrc/roi_pool.cpp @@ -1,18 +1,10 @@ -#pragma once +#include "roi_pool.h" +#include -#include "cpu/vision_cpu.h" - -#ifdef WITH_CUDA -#include "autocast.h" -#include "cuda/vision_cuda.h" -#endif -#ifdef WITH_HIP -#include "autocast.h" -#include "hip/vision_cuda.h" +#if defined(WITH_CUDA) || defined(WITH_HIP) +#include #endif -// TODO: put this stuff in torchvision namespace - std::tuple roi_pool( const at::Tensor& input, const at::Tensor& rois, @@ -26,7 +18,7 @@ std::tuple roi_pool( } #if defined(WITH_CUDA) || defined(WITH_HIP) -std::tuple ROIPool_autocast( +std::tuple roi_pool_autocast( const at::Tensor& input, const at::Tensor& rois, double spatial_scale, @@ -73,6 +65,8 @@ at::Tensor _roi_pool_backward( width); } +namespace { + class ROIPoolFunction : public torch::autograd::Function { public: static torch::autograd::variable_list forward( @@ -165,7 +159,9 @@ class ROIPoolBackwardFunction } }; -std::tuple ROIPool_autograd( +} // namespace + +std::tuple roi_pool_autograd( const at::Tensor& input, const at::Tensor& rois, double spatial_scale, @@ -177,7 +173,7 @@ std::tuple ROIPool_autograd( return std::make_tuple(result[0], result[1]); } -at::Tensor ROIPool_backward_autograd( +at::Tensor roi_pool_backward_autograd( const at::Tensor& grad, const at::Tensor& rois, const at::Tensor& argmax, diff --git a/torchvision/csrc/roi_pool.h b/torchvision/csrc/roi_pool.h new file mode 100644 index 00000000000..f528ce6d7e0 --- /dev/null +++ b/torchvision/csrc/roi_pool.h @@ -0,0 +1,61 @@ +#pragma once + +#include "cpu/roi_pool_kernel.h" + +#ifdef WITH_CUDA +#include "cuda/roi_pool_kernel.h" +#endif +#ifdef WITH_HIP +#include "hip/roi_pool_kernel.h" +#endif + +// C++ Forward +std::tuple roi_pool( + const at::Tensor& input, + const at::Tensor& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width); + +// Autocast Forward +#if defined(WITH_CUDA) || defined(WITH_HIP) +std::tuple roi_pool_autocast( + const at::Tensor& input, + const at::Tensor& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width); +#endif + +// C++ Backward +at::Tensor _roi_pool_backward( + const at::Tensor& grad, + const at::Tensor& rois, + const at::Tensor& argmax, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width); + +// Autograd Forward and Backward +std::tuple roi_pool_autograd( + const at::Tensor& input, + const at::Tensor& rois, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width); + +at::Tensor roi_pool_backward_autograd( + const at::Tensor& grad, + const at::Tensor& rois, + const at::Tensor& argmax, + double spatial_scale, + int64_t pooled_height, + int64_t pooled_width, + int64_t batch_size, + int64_t channels, + int64_t height, + int64_t width); diff --git a/torchvision/csrc/vision.cpp b/torchvision/csrc/vision.cpp index c41663f0736..d764ec9334b 100644 --- a/torchvision/csrc/vision.cpp +++ b/torchvision/csrc/vision.cpp @@ -8,13 +8,13 @@ #include #endif -#include "ROIPool.h" #include "deform_conv2d.h" #include "empty_tensor_op.h" #include "nms.h" #include "ps_roi_align.h" #include "ps_roi_pool.h" #include "roi_align.h" +#include "roi_pool.h" // If we are in a Windows environment, we need to define // initialization functions for the _custom_ops extension @@ -71,8 +71,8 @@ TORCH_LIBRARY_IMPL(torchvision, CPU, m) { m.impl("_ps_roi_pool_backward", ps_roi_pool_backward_cpu); m.impl("roi_align", roi_align_forward_cpu); m.impl("_roi_align_backward", roi_align_backward_cpu); - m.impl("roi_pool", ROIPool_forward_cpu); - m.impl("_roi_pool_backward", ROIPool_backward_cpu); + m.impl("roi_pool", roi_pool_forward_cpu); + m.impl("_roi_pool_backward", roi_pool_backward_cpu); } // TODO: Place this in a hypothetical separate torchvision_cuda library @@ -87,8 +87,8 @@ TORCH_LIBRARY_IMPL(torchvision, CUDA, m) { m.impl("_ps_roi_pool_backward", ps_roi_pool_backward_cuda); m.impl("roi_align", roi_align_forward_cuda); m.impl("_roi_align_backward", roi_align_backward_cuda); - m.impl("roi_pool", ROIPool_forward_cuda); - m.impl("_roi_pool_backward", ROIPool_backward_cuda); + m.impl("roi_pool", roi_pool_forward_cuda); + m.impl("_roi_pool_backward", roi_pool_backward_cuda); } #endif @@ -100,7 +100,7 @@ TORCH_LIBRARY_IMPL(torchvision, Autocast, m) { m.impl("ps_roi_align", ps_roi_align_autocast); m.impl("ps_roi_pool", ps_roi_pool_autocast); m.impl("roi_align", roi_align_autocast); - m.impl("roi_pool", ROIPool_autocast); + m.impl("roi_pool", roi_pool_autocast); } #endif @@ -113,6 +113,6 @@ TORCH_LIBRARY_IMPL(torchvision, Autograd, m) { m.impl("_ps_roi_pool_backward", ps_roi_pool_backward_autograd); m.impl("roi_align", roi_align_autograd); m.impl("_roi_align_backward", roi_align_backward_autograd); - m.impl("roi_pool", ROIPool_autograd); - m.impl("_roi_pool_backward", ROIPool_backward_autograd); + m.impl("roi_pool", roi_pool_autograd); + m.impl("_roi_pool_backward", roi_pool_backward_autograd); }