Caffe2 - C++ API
A deep learning, cross platform ML framework
leaky_relu_op.cc
1 
17 #include "caffe2/operators/leaky_relu_op.h"
18 
19 #include "caffe2/utils/math.h"
20 
21 namespace caffe2 {
22 
23 template <>
24 bool LeakyReluOp<float, CPUContext>::RunOnDevice() {
25  const auto& X = Input(0);
26  auto* Y = Output(0);
27  Y->ResizeLike(X);
28  ConstEigenVectorMap<float> Xvec(X.template data<float>(), X.size());
29  EigenVectorMap<float> Yvec(Y->template mutable_data<float>(), Y->size());
30  Yvec = Xvec.cwiseMax(0.f) + Xvec.cwiseMin(0.f) * alpha_;
31  return true;
32 }
33 
34 template <>
35 bool LeakyReluGradientOp<float, CPUContext>::RunOnDevice() {
36  const auto& Y = Input(0);
37  const auto& dY = Input(1);
38  auto* dX = Output(0);
39  dX->ResizeLike(Y);
40  CAFFE_ENFORCE_EQ(Y.size(), dY.size());
41  ConstEigenVectorMap<float> Yvec(Y.template data<float>(), Y.size());
42  ConstEigenVectorMap<float> dYvec(dY.template data<float>(), dY.size());
43  EigenVectorMap<float> dXvec(dX->template mutable_data<float>(), dX->size());
44  Eigen::VectorXf gtZero = (Yvec.array() >= 0.0f).cast<float>();
45  dXvec = dYvec.array() * gtZero.array() -
46  dYvec.array() * (gtZero.array() - 1.0f) * alpha_;
47  return true;
48 }
49 
50 REGISTER_CPU_OPERATOR(LeakyRelu, LeakyReluOp<float, CPUContext>);
51 REGISTER_CPU_OPERATOR(
52  LeakyReluGradient,
53  LeakyReluGradientOp<float, CPUContext>);
54 
55 OPERATOR_SCHEMA(LeakyRelu)
56  .NumInputs(1)
57  .NumOutputs(1)
58  .Arg("alpha", "Coefficient of leakage, default value is 0.01")
59  .AllowInplace({{0, 0}})
60  .CostInferenceFunction(PointwiseCostInference<2>)
61  .IdenticalTypeAndShape()
62  .SetDoc(R"DOC(
63 LeakyRelu takes input data (Tensor<T>) and an argument alpha, and produces one
64 output data (Tensor<T>) where the function `f(x) = alpha * x for x < 0`,
65 `f(x) = x for x >= 0`, is applied to the data tensor elementwise.
66 )DOC")
67  .Input(0, "X", "1D input tensor")
68  .Output(0, "Y", "1D input tensor");
69 OPERATOR_SCHEMA(LeakyReluGradient)
70  .NumInputs(2)
71  .NumOutputs(1)
72  .AllowInplace({{1, 0}})
73  .Arg("alpha", "Coefficient of leakage");
74 
76  using GradientMakerBase::GradientMakerBase;
77  vector<OperatorDef> GetGradientDefs() override {
78  return SingleGradientDef(
79  "LeakyReluGradient",
80  "",
81  vector<string>{O(0), GO(0)},
82  vector<string>{GI(0)});
83  }
84 };
85 
86 REGISTER_GRADIENT(LeakyRelu, GetLeakyReluGradient);
87 
88 } // 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 ...