Caffe2 - C++ API
A deep learning, cross platform ML framework
elementwise_linear_op.cc
1 
17 #include "elementwise_linear_op.h"
18 
19 namespace caffe2 {
20 
21 template<>
22 bool ElementwiseLinearOp<float, CPUContext>::RunOnDevice(){
23  const auto& X = Input(0);
24  const auto& a = Input(1);
25  const auto& b = Input(2);
26  auto* Y = Output(0);
27 
28  const auto canonical_axis = X.canonical_axis_index(axis_);
29  const int N = X.size_to_dim(canonical_axis);
30  const int D = X.size_from_dim(canonical_axis);
31 
32  CAFFE_ENFORCE_EQ(a.ndim(), 1, a.ndim());
33  CAFFE_ENFORCE_EQ(a.dim(0), D, a.ndim());
34  CAFFE_ENFORCE_EQ(b.ndim(), 1, b.ndim());
35  CAFFE_ENFORCE_EQ(b.dim(0), D, b.ndim());
36 
37  Y->ResizeLike(X);
38 
39  const float* X_data = X.data<float>();
40  const float* a_data = a.data<float>();
41  const float* b_data = b.data<float>();
42  float* Y_data = Y->mutable_data<float>();
43 
44  int p = 0;
45  for (int n = 0; n < N; ++n) {
46  for (int d = 0; d < D; ++d) {
47  Y_data[p] = X_data[p] * a_data[d] + b_data[d];
48  p++;
49  }
50  }
51  return true;
52 }
53 
54 template<>
55 bool ElementwiseLinearGradientOp<float, CPUContext>::RunOnDevice(){
56  const auto& g_o = Input(0);
57  const auto& X = Input(1);
58  const auto& a = Input(2);
59 
60  const auto canonical_axis = X.canonical_axis_index(axis_);
61  const int N = X.size_to_dim(canonical_axis);
62  const int D = X.size_from_dim(canonical_axis);
63 
64  CAFFE_ENFORCE_EQ(a.ndim(), 1, a.ndim());
65  CAFFE_ENFORCE_EQ(a.dim(0), D, a.ndim());
66 
67  auto *g_X = Output(0);
68  auto *g_a = Output(1);
69  auto *g_b = Output(2);
70  g_X->ResizeLike(X);
71  g_a->ResizeLike(a);
72  g_b->ResizeLike(a);
73 
74  const float* g_o_data = g_o.data<float>();
75  const float* X_data = X.data<float>();
76  const float* a_data = a.data<float>();
77  float* g_X_data = g_X->mutable_data<float>();
78  float* g_a_data = g_a->mutable_data<float>();
79  float* g_b_data = g_b->mutable_data<float>();
80 
81  math::Set<float, CPUContext>(g_a->size(), 0.f, g_a_data, &context_);
82  math::Set<float, CPUContext>(g_b->size(), 0.f, g_b_data, &context_);
83 
84  int p = 0;
85  for (int n = 0; n < N; ++n) {
86  for (int d = 0; d < D; ++d) {
87  g_X_data[p] = g_o_data[p] * a_data[d];
88  g_a_data[d] += g_o_data[p] * X_data[p];
89  g_b_data[d] += g_o_data[p];
90  p++;
91  }
92  }
93  return true;
94 }
95 
96 REGISTER_CPU_OPERATOR(
97  ElementwiseLinear,
98  ElementwiseLinearOp<float, CPUContext>);
99 REGISTER_CPU_OPERATOR(
100  ElementwiseLinearGradient,
101  ElementwiseLinearGradientOp<float, CPUContext>);
102 
103 OPERATOR_SCHEMA(ElementwiseLinear)
104  .NumInputs(3)
105  .NumOutputs(1)
106  .SetDoc(R"DOC(
107 Given inputs X of size (N x D), w of size D and b of size D,
108 the op computes Y of size (N X D) where Y_{nd} = X_{nd} * w_d + b_d
109  )DOC")
110  .Input(0, "X", "2D input tensor of size (N X D) data")
111  .Input(1, "w", "1D scaling factors of size D")
112  .Input(2, "b", "1D biases of size D")
113  .Output(0, "Y", "2D output tensor")
114  .Arg(
115  "axis",
116  "default to 1; describes the axis of the inputs; "
117  "defaults to one because the 0th axis most likely describes "
118  "the batch_size");
119 
120 OPERATOR_SCHEMA(ElementwiseLinearGradient)
121  .NumInputs(3)
122  .NumOutputs(3);
123 
125  using GradientMakerBase::GradientMakerBase;
126  vector<OperatorDef> GetGradientDefs() override {
127  return SingleGradientDef(
128  "ElementwiseLinearGradient",
129  "",
130  vector<string>{GO(0), I(0), I(1)},
131  vector<string>{GI(0), GI(1), GI(2)});
132  }
133 };
134 
135 REGISTER_GRADIENT(
136  ElementwiseLinear,
138 );
139 
140 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.
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 ...