prelu_op.h 2.8 KB
Newer Older
Z
zchen0211 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once
#include "paddle/framework/eigen.h"
#include "paddle/framework/op_registry.h"
Z
zchen0211 已提交
18
#include "paddle/platform/transform.h"
Z
zchen0211 已提交
19 20 21 22 23

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
Z
zchen0211 已提交
24
using platform::Transform;
Z
zchen0211 已提交
25

Z
zchen0211 已提交
26
template <typename T>
Z
zchen0211 已提交
27
class PReluFunctor {
Z
zchen0211 已提交
28
 public:
Z
zchen0211 已提交
29
  explicit PReluFunctor(const T* alpha) : alpha_(alpha) {}
Z
zchen0211 已提交
30

Z
prelu  
zchen0211 已提交
31 32 33
  HOSTDEVICE T operator()(const T& x) const {
    if (x > 0)
      return x;
Z
zchen0211 已提交
34
    else
Z
zchen0211 已提交
35
      return x * (*alpha_);
Z
zchen0211 已提交
36 37 38
  }

 private:
Z
zchen0211 已提交
39
  const T* alpha_;
Z
zchen0211 已提交
40 41
};

Z
prelu  
zchen0211 已提交
42
template <typename Place, typename T>
Y
Yu Yang 已提交
43
class PReluKernel : public framework::OpKernel<T> {
Z
zchen0211 已提交
44 45
 public:
  void Compute(const framework::ExecutionContext& context) const override {
Z
prelu  
zchen0211 已提交
46 47 48
    auto* x = context.Input<Tensor>("X");
    auto* alpha = context.Input<Tensor>("Alpha");
    auto* out = context.Output<Tensor>("Out");
Z
zchen0211 已提交
49

Z
prelu  
zchen0211 已提交
50 51
    const T* x_ptr = x->data<T>();
    T* o_ptr = out->mutable_data<T>(context.GetPlace());
Z
zchen0211 已提交
52

Z
zchen0211 已提交
53
    auto* alpha_ptr = alpha->data<T>();
Z
zchen0211 已提交
54

Z
prelu  
zchen0211 已提交
55
    int numel = x->numel();
Z
zchen0211 已提交
56

57 58 59
    Transform<Place> trans;
    trans(context.device_context(), x_ptr, x_ptr + numel, o_ptr,
          PReluFunctor<T>(alpha_ptr));
Z
zchen0211 已提交
60 61 62
  }
};

Z
zchen0211 已提交
63
template <typename T>
Z
zchen0211 已提交
64
class PReluGradFunctor {
Z
zchen0211 已提交
65
 public:
Z
zchen0211 已提交
66
  explicit PReluGradFunctor(const T* alpha) : alpha_(alpha) {}
Z
zchen0211 已提交
67

Z
prelu  
zchen0211 已提交
68 69 70
  HOSTDEVICE T operator()(const T& out, const T& dout) const {
    if (out > 0)
      return dout;
Z
zchen0211 已提交
71
    else
Z
zchen0211 已提交
72
      return dout * (*alpha_);
Z
zchen0211 已提交
73 74 75
  }

 private:
Z
zchen0211 已提交
76
  const T* alpha_;
Z
zchen0211 已提交
77 78
};

Z
prelu  
zchen0211 已提交
79
template <typename Place, typename T>
Y
Yu Yang 已提交
80
class PReluGradKernel : public framework::OpKernel<T> {
Z
zchen0211 已提交
81 82
 public:
  void Compute(const framework::ExecutionContext& context) const override {
Z
prelu  
zchen0211 已提交
83 84
    auto* dx = context.Output<Tensor>(framework::GradVarName("X"));
    auto* dout = context.Input<Tensor>(framework::GradVarName("Out"));
Z
zchen0211 已提交
85

Z
prelu  
zchen0211 已提交
86 87
    auto* out = context.Input<Tensor>("Out");
    auto* alpha = context.Input<Tensor>("Alpha");
Z
zchen0211 已提交
88
    auto* alpha_ptr = alpha->data<T>();
Z
zchen0211 已提交
89

Z
prelu  
zchen0211 已提交
90 91 92 93
    T* dx_ptr = dx->mutable_data<T>(context.GetPlace());
    const T* dout_ptr = dout->data<T>();
    const T* out_ptr = out->data<T>();
    int numel = dx->numel();
Z
zchen0211 已提交
94

95 96 97
    Transform<Place> trans;
    trans(context.device_context(), out_ptr, out_ptr + numel, dout_ptr, dx_ptr,
          PReluGradFunctor<T>(alpha_ptr));
Z
zchen0211 已提交
98

99
    // TODO(Zhuoyuan): add dalpha upgrade when GPU kernels ready
Z
zchen0211 已提交
100 101 102 103 104
  }
};

}  // namespace operators
}  // namespace paddle