vision_cuda.h 3.25 KB
Newer Older
1
2
#pragma once
#include <torch/extension.h>
3
#include "../macros.h"
4

5
6
7
8
VISION_API at::Tensor DeformConv2d_forward_cuda(
    const at::Tensor& input,
    const at::Tensor& weight,
    const at::Tensor& offset,
9
    const at::Tensor& mask,
10
11
12
13
14
15
16
17
    const at::Tensor& bias,
    int64_t stride_h,
    int64_t stride_w,
    int64_t pad_h,
    int64_t pad_w,
    int64_t dilation_h,
    int64_t dilation_w,
    int64_t groups,
18
19
    int64_t deformable_groups,
    bool use_mask);
20

21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
VISION_API std::
    tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor, at::Tensor>
    DeformConv2d_backward_cuda(
        const at::Tensor& grad_out,
        const at::Tensor& input,
        const at::Tensor& weight,
        const at::Tensor& offset,
        const at::Tensor& mask,
        const at::Tensor& bias,
        int64_t stride_h,
        int64_t stride_w,
        int64_t pad_h,
        int64_t pad_w,
        int64_t dilation_h,
        int64_t dilation_w,
        int64_t groups,
        int64_t deformable_groups,
        bool use_mask);
39
40
41
42
43
44
45

VISION_API at::Tensor nms_cuda(
    const at::Tensor& dets,
    const at::Tensor& scores,
    double iou_threshold);

VISION_API std::tuple<at::Tensor, at::Tensor> PSROIAlign_forward_cuda(
46
47
    const at::Tensor& input,
    const at::Tensor& rois,
48
49
50
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width,
51
    int64_t sampling_ratio);
52

53
VISION_API at::Tensor PSROIAlign_backward_cuda(
54
55
    const at::Tensor& grad,
    const at::Tensor& rois,
56
    const at::Tensor& channel_mapping,
57
58
59
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width,
60
    int64_t sampling_ratio,
61
62
63
    int64_t batch_size,
    int64_t channels,
    int64_t height,
64
    int64_t width);
65

66
VISION_API std::tuple<at::Tensor, at::Tensor> PSROIPool_forward_cuda(
67
68
    const at::Tensor& input,
    const at::Tensor& rois,
69
70
71
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width);
72

73
VISION_API at::Tensor PSROIPool_backward_cuda(
74
75
    const at::Tensor& grad,
    const at::Tensor& rois,
76
77
78
79
80
81
82
83
    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);
84

85
VISION_API at::Tensor ROIAlign_forward_cuda(
86
87
    const at::Tensor& input,
    const at::Tensor& rois,
88
89
90
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width,
91
92
    int64_t sampling_ratio,
    bool aligned);
93

94
VISION_API at::Tensor ROIAlign_backward_cuda(
95
96
    const at::Tensor& grad,
    const at::Tensor& rois,
97
98
99
100
101
102
    double spatial_scale,
    int64_t pooled_height,
    int64_t pooled_width,
    int64_t batch_size,
    int64_t channels,
    int64_t height,
103
104
105
    int64_t width,
    int64_t sampling_ratio,
    bool aligned);
106

107
VISION_API std::tuple<at::Tensor, at::Tensor> ROIPool_forward_cuda(
108
    const at::Tensor& input,
109
110
111
112
    const at::Tensor& rois,
    const double spatial_scale,
    const int64_t pooled_height,
    const int64_t pooled_width);
113

114
115
116
117
118
119
120
121
122
123
124
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);