Skip to content

ROIPool + Dispatcher + Autocast + Code Cleanup #2922

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Oct 30, 2020
Merged

ROIPool + Dispatcher + Autocast + Code Cleanup #2922

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
588c819
Fixing types.
datumbox Oct 27, 2020
7abc3ce
Merge branch 'master' into feature/dispatcher_roipool
datumbox Oct 28, 2020
733878e
Dispatcher + Autocast.
datumbox Oct 28, 2020
81a0420
+ Autograd.
datumbox Oct 28, 2020
b71f0d3
Formating.
datumbox Oct 28, 2020
bfe9bbc
Fixing return casting with autocast.
datumbox Oct 28, 2020
dde8cf9
Clean up and refactor ROIPool implementation:
datumbox Oct 28, 2020
e8fba5e
Restore include headers.
datumbox Oct 30, 2020
1f8547d
New line at end of file.
datumbox Oct 30, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 126 additions & 53 deletions torchvision/csrc/ROIPool.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,59 +3,64 @@
#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"
#endif

std::tuple<at::Tensor, at::Tensor> ROIPool_forward(
// TODO: put this stuff in torchvision namespace

std::tuple<at::Tensor, at::Tensor> roi_pool(
const at::Tensor& input,
const at::Tensor& rois,
const double spatial_scale,
const int64_t pooled_height,
const int64_t pooled_width) {
if (input.is_cuda()) {
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width) {
static auto op = c10::Dispatcher::singleton()
.findSchemaOrThrow("torchvision::roi_pool", "")
.typed<decltype(roi_pool)>();
return op.call(input, rois, spatial_scale, pooled_height, pooled_width);
}

#if defined(WITH_CUDA) || defined(WITH_HIP)
return ROIPool_forward_cuda(
input, rois, spatial_scale, pooled_height, pooled_width);
#else
TORCH_CHECK(false, "Not compiled with GPU support");
#endif
}
return ROIPool_forward_cpu(
input, rois, spatial_scale, pooled_height, pooled_width);
std::tuple<at::Tensor, at::Tensor> ROIPool_autocast(
const at::Tensor& input,
const at::Tensor& rois,
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width) {
c10::impl::ExcludeDispatchKeyGuard no_autocast(c10::DispatchKey::Autocast);
auto result = roi_pool(
at::autocast::cached_cast(at::kFloat, input),
at::autocast::cached_cast(at::kFloat, rois),
spatial_scale,
pooled_height,
pooled_width);

return std::make_tuple(
std::get<0>(result).to(input.scalar_type()),
std::get<1>(result).to(input.scalar_type()));
}
#endif

at::Tensor ROIPool_backward(
at::Tensor _roi_pool_backward(
const at::Tensor& grad,
const at::Tensor& rois,
const at::Tensor& argmax,
const float spatial_scale,
const int pooled_height,
const int pooled_width,
const int batch_size,
const int channels,
const int height,
const int width) {
if (grad.is_cuda()) {
#if defined(WITH_CUDA) || defined(WITH_HIP)
return ROIPool_backward_cuda(
grad,
rois,
argmax,
spatial_scale,
pooled_height,
pooled_width,
batch_size,
channels,
height,
width);
#else
TORCH_CHECK(false, "Not compiled with GPU support");
#endif
}
return ROIPool_backward_cpu(
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width,
int64_t batch_size,
int64_t channels,
int64_t height,
int64_t width) {
static auto op = c10::Dispatcher::singleton()
.findSchemaOrThrow("torchvision::_roi_pool_backward", "")
.typed<decltype(_roi_pool_backward)>();
return op.call(
grad,
rois,
argmax,
Expand All @@ -72,33 +77,36 @@ class ROIPoolFunction : public torch::autograd::Function<ROIPoolFunction> {
public:
static torch::autograd::variable_list forward(
torch::autograd::AutogradContext* ctx,
torch::autograd::Variable input,
torch::autograd::Variable rois,
const double spatial_scale,
const int64_t pooled_height,
const int64_t pooled_width) {
const torch::autograd::Variable& input,
const torch::autograd::Variable& rois,
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width) {
ctx->saved_data["spatial_scale"] = spatial_scale;
ctx->saved_data["pooled_height"] = pooled_height;
ctx->saved_data["pooled_width"] = pooled_width;
ctx->saved_data["input_shape"] = input.sizes();
auto result = ROIPool_forward(
input, rois, spatial_scale, pooled_height, pooled_width);
at::AutoNonVariableTypeMode g;
auto result =
roi_pool(input, rois, spatial_scale, pooled_height, pooled_width);

auto output = std::get<0>(result);
auto argmax = std::get<1>(result);
ctx->save_for_backward({rois, argmax});
ctx->mark_non_differentiable({argmax});

return {output, argmax};
}

static torch::autograd::variable_list backward(
torch::autograd::AutogradContext* ctx,
torch::autograd::variable_list grad_output) {
const torch::autograd::variable_list& grad_output) {
// Use data saved in forward
auto saved = ctx->get_saved_variables();
auto rois = saved[0];
auto argmax = saved[1];
auto input_shape = ctx->saved_data["input_shape"].toIntList();
auto grad_in = ROIPool_backward(
auto grad_in = _roi_pool_backward(
grad_output[0],
rois,
argmax,
Expand All @@ -109,6 +117,7 @@ class ROIPoolFunction : public torch::autograd::Function<ROIPoolFunction> {
input_shape[1],
input_shape[2],
input_shape[3]);

return {grad_in,
torch::autograd::Variable(),
torch::autograd::Variable(),
Expand All @@ -117,13 +126,77 @@ class ROIPoolFunction : public torch::autograd::Function<ROIPoolFunction> {
}
};

std::tuple<at::Tensor, at::Tensor> roi_pool(
// TODO: There should be an easier way to do this
class ROIPoolBackwardFunction
: public torch::autograd::Function<ROIPoolBackwardFunction> {
public:
static torch::autograd::variable_list forward(
torch::autograd::AutogradContext* ctx,
const torch::autograd::Variable& grad,
const torch::autograd::Variable& rois,
const torch::autograd::Variable& 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) {
at::AutoNonVariableTypeMode g;
auto grad_in = _roi_pool_backward(
grad,
rois,
argmax,
spatial_scale,
pooled_height,
pooled_width,
batch_size,
channels,
height,
width);

return {grad_in};
}

static torch::autograd::variable_list backward(
torch::autograd::AutogradContext* ctx,
const torch::autograd::variable_list& grad_output) {
TORCH_CHECK(0, "double backwards on roi_pool not supported");
}
};

std::tuple<at::Tensor, at::Tensor> ROIPool_autograd(
const at::Tensor& input,
const at::Tensor& rois,
const double spatial_scale,
const int64_t pooled_height,
const int64_t pooled_width) {
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width) {
auto result = ROIPoolFunction::apply(
input, rois, spatial_scale, pooled_height, pooled_width);
return std::tuple<at::Tensor, at::Tensor>(result[0], result[1]);

return std::make_tuple(result[0], result[1]);
}

at::Tensor ROIPool_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) {
return ROIPoolBackwardFunction::apply(
grad,
rois,
argmax,
spatial_scale,
pooled_height,
pooled_width,
batch_size,
channels,
height,
width)[0];
}
52 changes: 26 additions & 26 deletions torchvision/csrc/cpu/ROIPool_cpu.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,13 @@ template <typename T>
void RoIPoolForward(
const T* input,
const T spatial_scale,
const int channels,
const int height,
const int width,
const int pooled_height,
const int pooled_width,
int channels,
int height,
int width,
int pooled_height,
int pooled_width,
const T* rois,
const int num_rois,
int num_rois,
T* output,
int* argmax_data) {
for (int n = 0; n < num_rois; ++n) {
Expand Down Expand Up @@ -81,18 +81,18 @@ template <typename T>
void RoIPoolBackward(
const T* grad_output,
const int* argmax_data,
const int num_rois,
const int channels,
const int height,
const int width,
const int pooled_height,
const int pooled_width,
int num_rois,
int channels,
int height,
int width,
int pooled_height,
int pooled_width,
T* grad_input,
const T* rois,
const int n_stride,
const int c_stride,
const int h_stride,
const int w_stride) {
int n_stride,
int c_stride,
int h_stride,
int w_stride) {
for (int n = 0; n < num_rois; ++n) {
const T* offset_rois = rois + n * 5;
int roi_batch_ind = offset_rois[0];
Expand Down Expand Up @@ -123,9 +123,9 @@ void RoIPoolBackward(
std::tuple<at::Tensor, at::Tensor> ROIPool_forward_cpu(
const at::Tensor& input,
const at::Tensor& rois,
const float spatial_scale,
const int pooled_height,
const int pooled_width) {
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width) {
TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");

Expand Down Expand Up @@ -172,13 +172,13 @@ at::Tensor ROIPool_backward_cpu(
const at::Tensor& grad,
const at::Tensor& rois,
const at::Tensor& argmax,
const float spatial_scale,
const int pooled_height,
const int pooled_width,
const int batch_size,
const int channels,
const int height,
const int width) {
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width,
int64_t batch_size,
int64_t channels,
int64_t height,
int64_t width) {
// Check if input tensors are CPU tensors
TORCH_CHECK(grad.device().is_cpu(), "grad must be a CPU tensor");
TORCH_CHECK(rois.device().is_cpu(), "rois must be a CPU tensor");
Expand Down
20 changes: 10 additions & 10 deletions torchvision/csrc/cpu/vision_cpu.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,21 @@
VISION_API std::tuple<at::Tensor, at::Tensor> ROIPool_forward_cpu(
const at::Tensor& input,
const at::Tensor& rois,
const float spatial_scale,
const int pooled_height,
const int pooled_width);
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,
const float spatial_scale,
const int pooled_height,
const int pooled_width,
const int batch_size,
const int channels,
const int height,
const int width);
double spatial_scale,
int64_t pooled_height,
int64_t pooled_width,
int64_t batch_size,
int64_t channels,
int64_t height,
int64_t width);

VISION_API at::Tensor ROIAlign_forward_cpu(
const at::Tensor& input,
Expand Down
Loading