forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
conv_transpose_gradient_op.cc
50 lines (44 loc) · 1.61 KB
/
conv_transpose_gradient_op.cc
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
#include "caffe2/operators/conv_transpose_op.h"
#include "caffe2/operators/conv_transpose_op_impl.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(
ConvTransposeGradient,
ConvTransposeGradientOp<float, CPUContext>);
OPERATOR_SCHEMA(ConvTransposeGradient).NumInputs(3).NumOutputs(1, 3);
class GetConvTransposeGradient : public GradientMakerBase {
using GradientMakerBase::GradientMakerBase;
vector<OperatorDef> GetGradientDefs() override {
auto compute_dX =
!ArgumentHelper::GetSingleArgument(def_, "no_gradient_to_input", false);
CAFFE_ENFORCE(3 == def_.input_size() || 2 == def_.input_size());
if (def_.input_size() == 3 && compute_dX) {
return SingleGradientDef(
"ConvTransposeGradient",
"",
vector<string>{I(0), I(1), GO(0)},
vector<string>{GI(1), GI(2), GI(0)});
} else if (def_.input_size() == 3) {
return SingleGradientDef(
"ConvTransposeGradient",
"",
vector<string>{I(0), I(1), GO(0)},
vector<string>{GI(1), GI(2)});
} else if (compute_dX) {
return SingleGradientDef(
"ConvTransposeGradient",
"",
vector<string>{I(0), I(1), GO(0)},
vector<string>{GI(1), GI(0)},
vector<Argument>{MakeArgument<bool>("no_bias", true)});
} else {
return SingleGradientDef(
"ConvTransposeGradient",
"",
vector<string>{I(0), I(1), GO(0)},
vector<string>{GI(1)},
vector<Argument>{MakeArgument<bool>("no_bias", true)});
}
}
};
REGISTER_GRADIENT(ConvTranspose, GetConvTransposeGradient);
} // namespace caffe2