ps_roi_pool.cpp 5.7 KB
Newer Older
1
2
#include "ps_roi_pool.h"
#include <torch/extension.h>
3

4
5
#if defined(WITH_CUDA) || defined(WITH_HIP)
#include <ATen/autocast_mode.h>
6
#endif
7

8
9
namespace vision {
namespace ops {
10
11

std::tuple<at::Tensor, at::Tensor> ps_roi_pool(
12
13
    const at::Tensor& input,
    const at::Tensor& rois,
14
15
16
17
18
19
20
21
22
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width) {
  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);
}

23
#if defined(WITH_CUDA) || defined(WITH_HIP)
24
std::tuple<at::Tensor, at::Tensor> ps_roi_pool_autocast(
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
    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 = ps_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()));
41
}
42
#endif
43

44
at::Tensor _ps_roi_pool_backward(
45
46
    const at::Tensor& grad,
    const at::Tensor& rois,
47
48
49
50
51
52
53
54
55
56
57
58
59
    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(
60
61
      grad,
      rois,
62
      channel_mapping,
63
64
65
66
67
68
69
70
71
      spatial_scale,
      pooled_height,
      pooled_width,
      batch_size,
      channels,
      height,
      width);
}

72
73
namespace {

74
75
class PSROIPoolFunction : public torch::autograd::Function<PSROIPoolFunction> {
 public:
76
77
  static torch::autograd::variable_list forward(
      torch::autograd::AutogradContext* ctx,
78
79
80
81
82
      const torch::autograd::Variable& input,
      const torch::autograd::Variable& rois,
      double spatial_scale,
      int64_t pooled_height,
      int64_t pooled_width) {
83
84
85
86
    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();
87
88
89
90
    at::AutoNonVariableTypeMode g;
    auto result =
        ps_roi_pool(input, rois, spatial_scale, pooled_height, pooled_width);

91
92
93
94
    auto output = std::get<0>(result);
    auto channel_mapping = std::get<1>(result);
    ctx->save_for_backward({rois, channel_mapping});
    ctx->mark_non_differentiable({channel_mapping});
95

96
97
98
    return {output, channel_mapping};
  }

99
100
  static torch::autograd::variable_list backward(
      torch::autograd::AutogradContext* ctx,
101
      const torch::autograd::variable_list& grad_output) {
102
103
104
105
106
    // Use data saved in forward
    auto saved = ctx->get_saved_variables();
    auto rois = saved[0];
    auto channel_mapping = saved[1];
    auto input_shape = ctx->saved_data["input_shape"].toIntList();
107
    auto grad_in = _ps_roi_pool_backward(
108
109
110
111
112
113
114
115
116
117
        grad_output[0],
        rois,
        channel_mapping,
        ctx->saved_data["spatial_scale"].toDouble(),
        ctx->saved_data["pooled_height"].toInt(),
        ctx->saved_data["pooled_width"].toInt(),
        input_shape[0],
        input_shape[1],
        input_shape[2],
        input_shape[3]);
118

119
120
121
122
123
    return {grad_in,
            torch::autograd::Variable(),
            torch::autograd::Variable(),
            torch::autograd::Variable(),
            torch::autograd::Variable()};
124
125
126
  }
};

127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
// TODO: There should be an easier way to do this
class PSROIPoolBackwardFunction
    : public torch::autograd::Function<PSROIPoolBackwardFunction> {
 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& 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) {
    at::AutoNonVariableTypeMode g;
    auto grad_in = _ps_roi_pool_backward(
        grad,
        rois,
        channel_mapping,
        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 ps_roi_pool not supported");
  }
};

166
167
168
} // namespace

std::tuple<at::Tensor, at::Tensor> ps_roi_pool_autograd(
169
170
    const at::Tensor& input,
    const at::Tensor& rois,
171
172
173
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width) {
174
175
  auto result = PSROIPoolFunction::apply(
      input, rois, spatial_scale, pooled_height, pooled_width);
176
177
178
179

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

180
at::Tensor ps_roi_pool_backward_autograd(
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
    const at::Tensor& grad,
    const at::Tensor& rois,
    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) {
  return PSROIPoolBackwardFunction::apply(
      grad,
      rois,
      channel_mapping,
      spatial_scale,
      pooled_height,
      pooled_width,
      batch_size,
      channels,
      height,
      width)[0];
202
}
203
204
205

} // namespace ops
} // namespace vision