forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
SpatialUpSamplingBilinear.cu
106 lines (98 loc) · 4.55 KB
/
SpatialUpSamplingBilinear.cu
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
#ifndef THC_GENERIC_FILE
#define THC_GENERIC_FILE "THCUNN/generic/SpatialUpSamplingBilinear.cu"
#else
#include <THCUNN/upsampling.h>
#include "ATen/cuda/CUDAContext.h"
static inline void THNN_(SpatialUpSamplingBilinear_shapeCheck)
(THCState *state,
THCTensor *input, THCTensor *gradOutput,
int nBatch, int nChannels,
int inputHeight, int inputWidth,
int outputHeight, int outputWidth) {
THArgCheck(inputHeight > 0 && inputWidth > 0
&& outputHeight > 0 && outputWidth > 0, 2,
"input and output sizes should be greater than 0,"
" but got input (H: %d, W: %d) output (H: %d, W: %d)",
inputHeight, inputWidth, outputHeight, outputWidth);
if (input != NULL) {
THCUNN_argCheck(state, !input->is_empty() && input->dim() == 4, 2, input,
"non-empty 4D input tensor expected but got: %s");
}
if (gradOutput != NULL) {
THCUNN_check_dim_size(state, gradOutput, 4, 0, nBatch);
THCUNN_check_dim_size(state, gradOutput, 4, 1, nChannels);
THCUNN_check_dim_size(state, gradOutput, 4, 2, outputHeight);
THCUNN_check_dim_size(state, gradOutput, 4, 3, outputWidth);
}
}
void THNN_(SpatialUpSamplingBilinear_updateOutput)(
THCState *state,
THCTensor *input,
THCTensor *output,
int outputHeight,
int outputWidth,
bool align_corners)
{
int nbatch = THCTensor_(size)(state, input, 0);
int channels = THCTensor_(size)(state, input, 1);
int inputHeight = THCTensor_(size)(state, input, 2);
int inputWidth = THCTensor_(size)(state, input, 3);
THNN_(SpatialUpSamplingBilinear_shapeCheck)
(state, input, NULL,
nbatch, channels,
inputHeight, inputWidth,
outputHeight, outputWidth);
THCUNN_assertSameGPU(state, 2, input, output);
THCTensor_(resize4d)(state, output,
THCTensor_(size)(state, input, 0),
THCTensor_(size)(state, input, 1),
outputHeight, outputWidth);
THCTensor_(zero)(state, output);
THCDeviceTensor<scalar_t, 4> idata = toDeviceTensor<scalar_t, 4>(state, input);
THCDeviceTensor<scalar_t, 4> odata = toDeviceTensor<scalar_t, 4>(state, output);
THAssert(inputHeight > 0 && inputWidth > 0 && outputHeight > 0 && outputWidth > 0);
const accreal rheight = area_pixel_compute_scale<accreal>(inputHeight, outputHeight, align_corners);
const accreal rwidth = area_pixel_compute_scale<accreal>(inputWidth, outputWidth, align_corners);
const int num_kernels = outputHeight * outputWidth;
const int num_threads =
at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock;
cudaStream_t stream = THCState_getCurrentStream(state);
caffe_gpu_interp2_kernel<scalar_t, accreal> <<<THCCeilDiv(num_kernels, num_threads), num_threads ,
0 , stream>>>(num_kernels, rheight, rwidth, align_corners, idata, odata);
THCudaCheck(cudaGetLastError());
}
void THNN_(SpatialUpSamplingBilinear_updateGradInput)(
THCState *state,
THCTensor *gradOutput,
THCTensor *gradInput,
int nbatch,
int nchannels,
int inputHeight,
int inputWidth,
int outputHeight,
int outputWidth,
bool align_corners)
{
THNN_(SpatialUpSamplingBilinear_shapeCheck)
(state, NULL, gradOutput,
nbatch, nchannels,
inputHeight, inputWidth,
outputHeight, outputWidth);
gradOutput = THCTensor_(newContiguous)(state, gradOutput);
THCUNN_assertSameGPU(state, 2, gradOutput, gradInput);
THCTensor_(resize4d)(state, gradInput, nbatch, nchannels, inputHeight, inputWidth);
THCTensor_(zero)(state, gradInput);
THCDeviceTensor<scalar_t, 4> data1 = toDeviceTensor<scalar_t, 4>(state, gradInput);
THCDeviceTensor<scalar_t, 4> data2 = toDeviceTensor<scalar_t, 4>(state, gradOutput);
const accreal rheight = area_pixel_compute_scale<accreal>(inputHeight, outputHeight, align_corners);
const accreal rwidth = area_pixel_compute_scale<accreal>(inputWidth, outputWidth, align_corners);
const int num_kernels = outputHeight * outputWidth;
const int num_threads =
at::cuda::getCurrentDeviceProperties()->maxThreadsPerBlock;
cudaStream_t stream = THCState_getCurrentStream(state);
caffe_gpu_interp2_kernel_backward<scalar_t ,accreal> <<<THCCeilDiv(num_kernels, num_threads),
num_threads, 0, stream>>>(num_kernels, rheight, rwidth, align_corners, data1, data2);
THCudaCheck(cudaGetLastError());
THCTensor_(free)(state, gradOutput);
}
#endif