Skip to content

Commit

Permalink
ROIPool + Dispatcher + Autocast + Code Cleanup (pytorch#2922)
Browse files Browse the repository at this point in the history 
* Fixing types.

* Dispatcher + Autocast.

* + Autograd.

* Formating.

* Fixing return casting with autocast.

* Clean up and refactor ROIPool implementation:
- Remove primitive const declaration from method names.
- Using references when possible.

* Restore include headers.

* New line at end of file.
  • Loading branch information
@datumbox @bryant1410
datumbox authored and bryant1410 committed Nov 22, 2020
1 parent 28f68e9 commit a18bd38
Show file tree
Hide file tree
Showing 6 changed files with 210 additions and 127 deletions.
179 changes: 126 additions & 53 deletions torchvision/csrc/ROIPool.h
View file
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
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
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

0 comments on commit a18bd38

Please sign in to comment.