1 #include "elementwise_linear_op.h" 6 bool ElementwiseLinearOp<float, CPUContext>::RunOnDevice(){
7 const auto& X = Input(0);
8 const auto& a = Input(1);
9 const auto& b = Input(2);
12 const auto canonical_axis = X.canonical_axis_index(axis_);
13 const int N = X.size_to_dim(canonical_axis);
14 const int D = X.size_from_dim(canonical_axis);
16 CAFFE_ENFORCE_EQ(a.ndim(), 1, a.ndim());
17 CAFFE_ENFORCE_EQ(a.dim(0), D, a.ndim());
18 CAFFE_ENFORCE_EQ(b.ndim(), 1, b.ndim());
19 CAFFE_ENFORCE_EQ(b.dim(0), D, b.ndim());
23 const float* X_data = X.data<
float>();
24 const float* a_data = a.data<
float>();
25 const float* b_data = b.data<
float>();
26 float* Y_data = Y->mutable_data<
float>();
29 for (
int n = 0; n < N; ++n) {
30 for (
int d = 0; d < D; ++d) {
31 Y_data[p] = X_data[p] * a_data[d] + b_data[d];
39 bool ElementwiseLinearGradientOp<float, CPUContext>::RunOnDevice(){
40 const auto& g_o = Input(0);
41 const auto& X = Input(1);
42 const auto& a = Input(2);
44 const auto canonical_axis = X.canonical_axis_index(axis_);
45 const int N = X.size_to_dim(canonical_axis);
46 const int D = X.size_from_dim(canonical_axis);
48 CAFFE_ENFORCE_EQ(a.ndim(), 1, a.ndim());
49 CAFFE_ENFORCE_EQ(a.dim(0), D, a.ndim());
51 auto *g_X = Output(0);
52 auto *g_a = Output(1);
53 auto *g_b = Output(2);
58 const float* g_o_data = g_o.data<
float>();
59 const float* X_data = X.data<
float>();
60 const float* a_data = a.data<
float>();
61 float* g_X_data = g_X->mutable_data<
float>();
62 float* g_a_data = g_a->mutable_data<
float>();
63 float* g_b_data = g_b->mutable_data<
float>();
65 math::Set<float, CPUContext>(g_a->size(), 0.f, g_a_data, &context_);
66 math::Set<float, CPUContext>(g_b->size(), 0.f, g_b_data, &context_);
69 for (
int n = 0; n < N; ++n) {
70 for (
int d = 0; d < D; ++d) {
71 g_X_data[p] = g_o_data[p] * a_data[d];
72 g_a_data[d] += g_o_data[p] * X_data[p];
73 g_b_data[d] += g_o_data[p];
80 REGISTER_CPU_OPERATOR(
82 ElementwiseLinearOp<float, CPUContext>);
83 REGISTER_CPU_OPERATOR(
84 ElementwiseLinearGradient,
85 ElementwiseLinearGradientOp<float, CPUContext>);
87 OPERATOR_SCHEMA(ElementwiseLinear)
91 Given inputs X of size (N x D), w of size D and b of size D, 92 the op computes Y of size (N X D) where Y_{nd} = X_{nd} * w_d + b_d 94 .Input(0, "X",
"2D input tensor of size (N X D) data")
95 .Input(1,
"w",
"1D scaling factors of size D")
96 .Input(2,
"b",
"1D biases of size D")
97 .Output(0,
"Y",
"2D output tensor")
100 "default to 1; describes the axis of the inputs; " 101 "defaults to one because the 0th axis most likely describes " 104 OPERATOR_SCHEMA(ElementwiseLinearGradient)
109 using GradientMakerBase::GradientMakerBase;
110 vector<OperatorDef> GetGradientDefs()
override {
112 "ElementwiseLinearGradient",
114 vector<string>{GO(0), I(0), I(1)},
115 vector<string>{GI(0), GI(1), GI(2)});
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...