forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
SpatialMaxUnpooling.cu
104 lines (87 loc) · 3.37 KB
/
SpatialMaxUnpooling.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
#ifndef THC_GENERIC_FILE
#define THC_GENERIC_FILE "THCUNN/generic/SpatialMaxUnpooling.cu"
#else
void THNN_(SpatialMaxUnpooling_updateOutput)(
THCState *state,
THCTensor *input,
THCTensor *output,
THCIndexTensor *indices,
int owidth, int oheight)
{
THCUNN_assertSameGPU(state, 3, input, output, indices);
THCUNN_argCheck(state, !input->is_empty() && (input->dim() == 3 || input->dim() == 4), 2, input,
"non-empty 3D or 4D (batch mode) tensor expected for input, but got: %s");
THCUNN_check_shape_indices(state, indices, input);
int64_t nInputCols, nInputRows, nInputPlane, batchSize;
if (input->dim() == 3) {
nInputCols = input->size(2);
nInputRows = input->size(1);
nInputPlane = input->size(0);
batchSize = 1;
}
else
{
nInputCols = input->size(3);
nInputRows = input->size(2);
nInputPlane = input->size(1);
batchSize = input->size(0);
}
input = THCTensor_(newContiguous)(state, input);
indices = THCIndexTensor_(newContiguous)(state, indices);
THCTensor_(resize4d)(state, output, batchSize, nInputPlane, oheight, owidth);
THCTensor_(zero)(state, output);
int count = THCTensor_(nElement)(state, input);
MaxUnpoolForward <<< GET_BLOCKS(count), CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>
(count, THCTensor_(data)(state, input), THCIndexTensor_(data)(state, indices),
batchSize, nInputPlane, nInputRows, nInputCols, oheight, owidth, THCTensor_(data)(state, output));
THCudaCheck(cudaGetLastError());
if(input->dim() == 3)
THCTensor_(resize3d)(state, output, nInputPlane, oheight, owidth);
THCTensor_(free)(state, input);
THCIndexTensor_(free)(state, indices);
}
void THNN_(SpatialMaxUnpooling_updateGradInput)(
THCState *state,
THCTensor *input,
THCTensor *gradOutput,
THCTensor *gradInput,
THCIndexTensor *indices,
int owidth, int oheight)
{
THCUNN_assertSameGPU(state, 4, input, gradOutput, indices, gradInput);
THCUNN_check_shape_indices(state, indices, input);
int64_t nInputCols, nInputRows, nInputPlane, batchSize;
int dimw = 2;
int dimh = 1;
if (input->dim() == 3) {
nInputPlane = input->size(0);
batchSize = 1;
}
else
{
++dimw;
++dimh;
nInputPlane = input->size(1);
batchSize = input->size(0);
}
nInputCols = input->size(dimw);
nInputRows = input->size(dimh);
if(owidth!=gradOutput->size(dimw) || oheight!=gradOutput->size(dimh)){
THError("Inconsistent gradOutput size. oheight= %d, owidth= %d, gradOutput: %dx%d",
oheight, owidth,gradOutput->size(dimh),gradOutput->size(dimw));
}
input = THCTensor_(newContiguous)(state, input);
indices = THCIndexTensor_(newContiguous)(state, indices);
gradOutput = THCTensor_(newContiguous)(state, gradOutput);
THCTensor_(resizeAs)(state, gradInput, input);
int count = THCTensor_(nElement)(state, input);
MaxUnpoolBackward <<< GET_BLOCKS(count), CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state) >>>
(count, THCTensor_(data)(state, gradOutput), THCIndexTensor_(data)(state, indices),
batchSize, nInputPlane, nInputRows, nInputCols, oheight, owidth, THCTensor_(data)(state, gradInput));
THCudaCheck(cudaGetLastError());
// clean
THCTensor_(free)(state, input);
THCIndexTensor_(free)(state, indices);
THCTensor_(free)(state, gradOutput);
}
#endif