Caffe2 - C++ API
A deep learning, cross platform ML framework
thresholded_relu_op.cc
1 
17 #include "caffe2/operators/thresholded_relu_op.h"
18 
19 #include "caffe2/utils/math.h"
20 
21 namespace caffe2 {
22 
23 template <>
24 bool ThresholdedReluOp<float, CPUContext>::RunOnDevice() {
25  auto& X = Input(0);
26  auto* Y = Output(0);
27  Y->ResizeLike(X);
28 
29  ConstEigenVectorArrayMap<float> Xvec(X.data<float>(), X.size());
30  EigenVectorArrayMap<float> Yvec(Y->mutable_data<float>(), Y->size());
31  Yvec = (Xvec > alpha_).select(Xvec, 0.f);
32  /* Naive implementation
33  const float* Xdata = X.data<float>();
34  float* Ydata = Y->mutable_data<float>();
35  for (int i = 0; i < X.size(); ++i) {
36  Xdata[i] -= alpha_;
37  Ydata[i] = std::max(Xdata[i], 0.0f);
38  }
39  */
40  return true;
41 }
42 
43 template <>
44 bool ThresholdedReluGradientOp<float, CPUContext>::RunOnDevice() {
45  auto& Y = Input(0);
46  auto& dY = Input(1);
47  auto* dX = Output(0);
48  CAFFE_ENFORCE_EQ(dY.size(), Y.size());
49  dX->ResizeLike(Y);
50 
51  const float* Ydata = Y.data<float>();
52  const float* dYdata = dY.data<float>();
53  float* dXdata = dX->mutable_data<float>();
54  EigenVectorArrayMap<float> dXvec(dXdata, dX->size());
55  ConstEigenVectorArrayMap<float> Yvec(Ydata, Y.size());
56  ConstEigenVectorArrayMap<float> dYvec(dYdata, dY.size());
57  dXvec = dYvec * Yvec.cwiseSign();
58  /* Non vectorized implementation
59  for (int i = 0; i < Y.size(); ++i) {
60  dXdata[i] = Ydata[i] > 0 ? dYdata[i] : 0;
61  }
62  */
63  return true;
64 }
65 
66 REGISTER_CPU_OPERATOR(ThresholdedRelu, ThresholdedReluOp<float, CPUContext>);
67 REGISTER_CPU_OPERATOR(
68  ThresholdedReluGradient,
69  ThresholdedReluGradientOp<float, CPUContext>);
70 
71 // Input: X, output: Y
72 OPERATOR_SCHEMA(ThresholdedRelu)
73  .NumInputs(1)
74  .NumOutputs(1)
75  .AllowInplace({{0, 0}})
76  .CostInferenceFunction(PointwiseCostInference<2>)
77  .IdenticalTypeAndShape()
78  .SetDoc(R"DOC(
79 ThresholdedRelu takes one input data (Tensor) and produces one output data
80 (Tensor) where the rectified linear function, y = x for x > alpha, y = 0
81 otherwise, is applied to the tensor elementwise.
82 )DOC")
83  .Arg("alpha", "(float) defaults to 1.0.")
84  .Input(0, "X", "1D input tensor")
85  .Output(0, "Y", "1D input tensor");
86 
87 // Input: Y, dY, output: dX
88 OPERATOR_SCHEMA(ThresholdedReluGradient)
89  .NumInputs(2)
90  .NumOutputs(1)
91  .AllowInplace({{1, 0}})
92  .SetDoc(R"DOC(
93 ThresholdedReluGradient takes both Y and dY and uses this to update dX
94 according to the chain rule and derivatives of the rectified linear function.
95 )DOC");
96 
97 class GetThresholdedReluGradient : public GradientMakerBase {
98  using GradientMakerBase::GradientMakerBase;
99  vector<OperatorDef> GetGradientDefs() override {
100  return SingleGradientDef(
101  def_.type() + "Gradient",
102  "",
103  vector<string>{O(0), GO(0)},
104  vector<string>{GI(0)});
105  }
106 };
107 REGISTER_GRADIENT(ThresholdedRelu, GetThresholdedReluGradient);
108 
109 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.