ps_roi_pool.cpp 1.78 KB
Newer Older
1
#include "ps_roi_pool.h"
2
3

#include <torch/types.h>
4

5
6
namespace vision {
namespace ops {
7
8

std::tuple<at::Tensor, at::Tensor> ps_roi_pool(
9
10
    const at::Tensor& input,
    const at::Tensor& rois,
11
12
13
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width) {
Kai Zhang's avatar
Kai Zhang committed
14
  C10_LOG_API_USAGE_ONCE("torchvision.csrc.ops.ps_roi_pool.ps_roi_pool");
15
16
17
18
19
20
  static auto op = c10::Dispatcher::singleton()
                       .findSchemaOrThrow("torchvision::ps_roi_pool", "")
                       .typed<decltype(ps_roi_pool)>();
  return op.call(input, rois, spatial_scale, pooled_height, pooled_width);
}

21
22
namespace detail {

23
at::Tensor _ps_roi_pool_backward(
24
25
    const at::Tensor& grad,
    const at::Tensor& rois,
26
27
28
29
30
31
32
33
34
35
36
37
38
    const at::Tensor& channel_mapping,
    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::_ps_roi_pool_backward", "")
          .typed<decltype(_ps_roi_pool_backward)>();
  return op.call(
39
40
      grad,
      rois,
41
      channel_mapping,
42
43
44
45
46
47
48
49
50
      spatial_scale,
      pooled_height,
      pooled_width,
      batch_size,
      channels,
      height,
      width);
}

51
52
} // namespace detail

53
TORCH_LIBRARY_FRAGMENT(torchvision, m) {
54
55
56
57
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torchvision::ps_roi_pool(Tensor input, Tensor rois, float spatial_scale, int pooled_height, int pooled_width) -> (Tensor, Tensor)"));
  m.def(TORCH_SELECTIVE_SCHEMA(
      "torchvision::_ps_roi_pool_backward(Tensor grad, Tensor rois, Tensor channel_mapping, float spatial_scale, int pooled_height, int pooled_width, int batch_size, int channels, int height, int width) -> Tensor"));
58
59
}

60
61
} // namespace ops
} // namespace vision