mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
ed5b584f65
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/15087 Codemod generated with clangr shard mode, 25 files per diff, motivation: https://github.com/pytorch/pytorch/pull/12407 Reviewed By: ezyang Differential Revision: D13419765 fbshipit-source-id: 34d695309a66723281429610a12544598c507d74
95 lines
2.8 KiB
C++
95 lines
2.8 KiB
C++
#include "caffe2/operators/thresholded_relu_op.h"
|
|
|
|
#include "caffe2/utils/eigen_utils.h"
|
|
#include "caffe2/utils/math.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
template <>
|
|
bool ThresholdedReluOp<float, CPUContext>::RunOnDevice() {
|
|
auto& X = Input(0);
|
|
|
|
auto* Y = Output(0, X.sizes(), at::dtype<float>());
|
|
|
|
ConstEigenVectorArrayMap<float> Xvec(X.data<float>(), X.numel());
|
|
EigenVectorArrayMap<float> Yvec(
|
|
Y->template mutable_data<float>(), Y->numel());
|
|
Yvec = (Xvec > alpha_).select(Xvec, 0.f);
|
|
/* Naive implementation
|
|
const float* Xdata = X.data<float>();
|
|
float* Ydata = Y->template mutable_data<float>();
|
|
for (int i = 0; i < X.size(); ++i) {
|
|
Xdata[i] -= alpha_;
|
|
Ydata[i] = std::max(Xdata[i], 0.0f);
|
|
}
|
|
*/
|
|
return true;
|
|
}
|
|
|
|
template <>
|
|
bool ThresholdedReluGradientOp<float, CPUContext>::RunOnDevice() {
|
|
auto& Y = Input(0);
|
|
auto& dY = Input(1);
|
|
|
|
CAFFE_ENFORCE_EQ(dY.numel(), Y.numel());
|
|
auto* dX = Output(0, Y.sizes(), at::dtype<float>());
|
|
|
|
const float* Ydata = Y.data<float>();
|
|
const float* dYdata = dY.data<float>();
|
|
float* dXdata = dX->template mutable_data<float>();
|
|
EigenVectorArrayMap<float> dXvec(dXdata, dX->numel());
|
|
ConstEigenVectorArrayMap<float> Yvec(Ydata, Y.numel());
|
|
ConstEigenVectorArrayMap<float> dYvec(dYdata, dY.numel());
|
|
dXvec = dYvec * Yvec.cwiseSign();
|
|
/* Non vectorized implementation
|
|
for (int i = 0; i < Y.size(); ++i) {
|
|
dXdata[i] = Ydata[i] > 0 ? dYdata[i] : 0;
|
|
}
|
|
*/
|
|
return true;
|
|
}
|
|
|
|
REGISTER_CPU_OPERATOR(ThresholdedRelu, ThresholdedReluOp<float, CPUContext>);
|
|
REGISTER_CPU_OPERATOR(
|
|
ThresholdedReluGradient,
|
|
ThresholdedReluGradientOp<float, CPUContext>);
|
|
|
|
// Input: X, output: Y
|
|
OPERATOR_SCHEMA(ThresholdedRelu)
|
|
.NumInputs(1)
|
|
.NumOutputs(1)
|
|
.AllowInplace({{0, 0}})
|
|
.CostInferenceFunction(PointwiseCostInference<2>)
|
|
.IdenticalTypeAndShape()
|
|
.SetDoc(R"DOC(
|
|
ThresholdedRelu takes one input data (Tensor) and produces one output data
|
|
(Tensor) where the rectified linear function, y = x for x > alpha, y = 0
|
|
otherwise, is applied to the tensor elementwise.
|
|
)DOC")
|
|
.Arg("alpha", "(float) defaults to 1.0.")
|
|
.Input(0, "X", "1D input tensor")
|
|
.Output(0, "Y", "1D input tensor");
|
|
|
|
// Input: Y, dY, output: dX
|
|
OPERATOR_SCHEMA(ThresholdedReluGradient)
|
|
.NumInputs(2)
|
|
.NumOutputs(1)
|
|
.AllowInplace({{1, 0}})
|
|
.SetDoc(R"DOC(
|
|
ThresholdedReluGradient takes both Y and dY and uses this to update dX
|
|
according to the chain rule and derivatives of the rectified linear function.
|
|
)DOC");
|
|
|
|
class GetThresholdedReluGradient : public GradientMakerBase {
|
|
using GradientMakerBase::GradientMakerBase;
|
|
vector<OperatorDef> GetGradientDefs() override {
|
|
return SingleGradientDef(
|
|
def_.type() + "Gradient",
|
|
"",
|
|
vector<string>{O(0), GO(0)},
|
|
vector<string>{GI(0)});
|
|
}
|
|
};
|
|
REGISTER_GRADIENT(ThresholdedRelu, GetThresholdedReluGradient);
|
|
|
|
} // namespace caffe2
|