add approximation for gelu, test=develop (#22961)
add approximation for gelu, default value is False (only kernel with eigen is added, remove code for computing gelu with MKLDNN temporarily)revert-23830-2.0-beta
Feiyu Chan
5 years ago
committed by
GitHub
parent
eec10aaba2
commit
01ab8a0619
@ -0,0 +1,144 @@
|
|||||||
|
/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#include <memory>
|
||||||
|
#include <string>
|
||||||
|
#include <unordered_map>
|
||||||
|
|
||||||
|
#include "paddle/fluid/operators/gelu_op.h"
|
||||||
|
#include "paddle/fluid/platform/float16.h"
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
|
||||||
|
class GeluOp : public framework::OperatorWithKernel {
|
||||||
|
public:
|
||||||
|
GeluOp(const std::string &type, const framework::VariableNameMap &inputs,
|
||||||
|
const framework::VariableNameMap &outputs,
|
||||||
|
const framework::AttributeMap &attrs)
|
||||||
|
: OperatorWithKernel(type, inputs, outputs, attrs) {}
|
||||||
|
|
||||||
|
void InferShape(framework::InferShapeContext *ctx) const override {
|
||||||
|
PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
|
||||||
|
platform::errors::InvalidArgument(
|
||||||
|
"Input(%s) of GeluOp should not be null.", "X"));
|
||||||
|
PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
|
||||||
|
platform::errors::InvalidArgument(
|
||||||
|
"Output(%s) of GeluOp should not be null.", "Out"));
|
||||||
|
|
||||||
|
ctx->ShareDim("X", /*->*/ "Out");
|
||||||
|
ctx->ShareLoD("X", /*->*/ "Out");
|
||||||
|
}
|
||||||
|
|
||||||
|
protected:
|
||||||
|
framework::OpKernelType GetExpectedKernelType(
|
||||||
|
const framework::ExecutionContext &ctx) const override {
|
||||||
|
return framework::OpKernelType(
|
||||||
|
OperatorWithKernel::IndicateVarDataType(ctx, "X"), ctx.GetPlace());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class GeluGradOp : public framework::OperatorWithKernel {
|
||||||
|
public:
|
||||||
|
using framework::OperatorWithKernel::OperatorWithKernel;
|
||||||
|
|
||||||
|
void InferShape(framework::InferShapeContext *ctx) const override {
|
||||||
|
PADDLE_ENFORCE_EQ(
|
||||||
|
ctx->HasInput(framework::GradVarName("Out")), true,
|
||||||
|
platform::errors::InvalidArgument(
|
||||||
|
"Input(%s) of GeluGradOp should not be null.", "DOut"));
|
||||||
|
PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
|
||||||
|
platform::errors::InvalidArgument(
|
||||||
|
"Input(%s) of GeluGradOp should not be null.", "X"));
|
||||||
|
PADDLE_ENFORCE_EQ(
|
||||||
|
ctx->HasOutput(framework::GradVarName("X")), true,
|
||||||
|
platform::errors::InvalidArgument(
|
||||||
|
"Output(%s) of GeluGradOp should not be null.", "DX"));
|
||||||
|
auto x_grad_name = framework::GradVarName("X");
|
||||||
|
ctx->SetOutputDim(x_grad_name, ctx->GetInputDim("X"));
|
||||||
|
ctx->ShareLoD("X", /*->*/ x_grad_name);
|
||||||
|
}
|
||||||
|
|
||||||
|
protected:
|
||||||
|
framework::OpKernelType GetExpectedKernelType(
|
||||||
|
const framework::ExecutionContext &ctx) const override {
|
||||||
|
return framework::OpKernelType(
|
||||||
|
OperatorWithKernel::IndicateVarDataType(ctx, "X"), ctx.GetPlace());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
class GeluOpMaker : public framework::OpProtoAndCheckerMaker {
|
||||||
|
public:
|
||||||
|
void Make() override {
|
||||||
|
AddInput("X", "Input of Gelu operator");
|
||||||
|
AddOutput("Out", "Output of Gelu operator");
|
||||||
|
AddAttr<bool>("approximate",
|
||||||
|
"(bool, default false) use approximation of gelu")
|
||||||
|
.SetDefault(false);
|
||||||
|
AddAttr<bool>("use_mkldnn",
|
||||||
|
"(bool, default false) Only used in mkldnn kernel")
|
||||||
|
.SetDefault(false);
|
||||||
|
AddAttr<bool>("use_cudnn",
|
||||||
|
"(bool, default false) Only used in cudnn kernel, need "
|
||||||
|
"install cudnn")
|
||||||
|
.SetDefault(false);
|
||||||
|
AddAttr<bool>("is_test",
|
||||||
|
"(bool, default false) Set to true for inference only, false "
|
||||||
|
"for training. Some layers may run faster when this is true.")
|
||||||
|
.SetDefault(false);
|
||||||
|
AddComment(R"DOC(
|
||||||
|
Gelu Activation Operator.
|
||||||
|
|
||||||
|
For more details, please refer to [Gaussian Error Linear Units](https://arxiv.org/pdf/1606.08415.pdf).
|
||||||
|
|
||||||
|
when using approximation
|
||||||
|
$out = \\frac{1}{2}x(1+tanh(\\sqrt{\\frac{2}{\\pi}}(x+0.044715x^{3}))$
|
||||||
|
|
||||||
|
or else
|
||||||
|
$out = \\frac{1 + erf(\\frac{x}{\\sqrt{2}})}{2} x$
|
||||||
|
|
||||||
|
)DOC");
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
class GeluGradOpMaker : public framework::SingleGradOpMaker<T> {
|
||||||
|
public:
|
||||||
|
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
|
||||||
|
|
||||||
|
protected:
|
||||||
|
void Apply(GradOpPtr<T> grad_op) const override {
|
||||||
|
grad_op->SetType("gelu_grad");
|
||||||
|
grad_op->SetInput("X", this->Input("X"));
|
||||||
|
grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
|
||||||
|
grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
|
||||||
|
grad_op->SetAttrMap(this->Attrs());
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
||||||
|
|
||||||
|
namespace ops = paddle::operators;
|
||||||
|
|
||||||
|
REGISTER_OPERATOR(gelu, ops::GeluOp, ops::GeluOpMaker,
|
||||||
|
ops::GeluGradOpMaker<paddle::framework::OpDesc>,
|
||||||
|
ops::GeluGradOpMaker<paddle::imperative::OpBase>);
|
||||||
|
REGISTER_OPERATOR(gelu_grad, ops::GeluGradOp);
|
||||||
|
REGISTER_OP_CPU_KERNEL(
|
||||||
|
gelu, ops::GeluKernel<paddle::platform::CPUDeviceContext, float>,
|
||||||
|
ops::GeluKernel<paddle::platform::CPUDeviceContext, double>);
|
||||||
|
REGISTER_OP_CPU_KERNEL(
|
||||||
|
gelu_grad, ops::GeluGradKernel<paddle::platform::CPUDeviceContext, float>,
|
||||||
|
ops::GeluGradKernel<paddle::platform::CPUDeviceContext, double>);
|
||||||
@ -0,0 +1,28 @@
|
|||||||
|
/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
|
||||||
|
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. */
|
||||||
|
|
||||||
|
#include "paddle/fluid/operators/gelu_op.h"
|
||||||
|
#include "paddle/fluid/platform/float16.h"
|
||||||
|
|
||||||
|
namespace ops = paddle::operators;
|
||||||
|
REGISTER_OP_CUDA_KERNEL(
|
||||||
|
gelu, ops::GeluKernel<paddle::platform::CUDADeviceContext, float>,
|
||||||
|
ops::GeluKernel<paddle::platform::CUDADeviceContext, double>,
|
||||||
|
ops::GeluKernel<paddle::platform::CUDADeviceContext,
|
||||||
|
paddle::platform::float16>);
|
||||||
|
REGISTER_OP_CUDA_KERNEL(
|
||||||
|
gelu_grad, ops::GeluGradKernel<paddle::platform::CUDADeviceContext, float>,
|
||||||
|
ops::GeluGradKernel<paddle::platform::CUDADeviceContext, double>,
|
||||||
|
ops::GeluGradKernel<paddle::platform::CUDADeviceContext,
|
||||||
|
paddle::platform::float16>);
|
||||||
@ -0,0 +1,119 @@
|
|||||||
|
/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
||||||
|
|
||||||
|
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
|
||||||
|
#ifndef _USE_MATH_DEFINES
|
||||||
|
#define _USE_MATH_DEFINES
|
||||||
|
#endif
|
||||||
|
#include <algorithm>
|
||||||
|
#include <cmath>
|
||||||
|
#include "paddle/fluid/framework/eigen.h"
|
||||||
|
#include "paddle/fluid/framework/op_registry.h"
|
||||||
|
#include "paddle/fluid/operators/math/blas.h"
|
||||||
|
#include "paddle/fluid/platform/float16.h"
|
||||||
|
|
||||||
|
#ifdef PADDLE_WITH_MKLDNN
|
||||||
|
#include "paddle/fluid/platform/mkldnn_helper.h"
|
||||||
|
#endif
|
||||||
|
|
||||||
|
namespace paddle {
|
||||||
|
namespace operators {
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
struct GeluFunctor {
|
||||||
|
template <typename Device, typename X, typename Out>
|
||||||
|
void operator()(Device d, X x, Out out, bool approximate) const {
|
||||||
|
if (approximate) {
|
||||||
|
// gelu(x) = 0.5 * x * (1 + tanh(sqrt(2 / \pi) * (x + 0.044715 * x^{3})))
|
||||||
|
auto temp = (static_cast<T>(M_2_SQRTPI * M_SQRT1_2) *
|
||||||
|
(x + static_cast<T>(0.044715) * x.cube()))
|
||||||
|
.tanh();
|
||||||
|
out.device(d) = x * static_cast<T>(0.5) * (static_cast<T>(1) + temp);
|
||||||
|
} else {
|
||||||
|
// gelu(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
|
||||||
|
auto temp = (x * static_cast<T>(M_SQRT1_2)).erf();
|
||||||
|
out.device(d) = x * static_cast<T>(0.5) * (static_cast<T>(1) + temp);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
struct GeluGradFunctor {
|
||||||
|
template <typename Device, typename X, typename dOut, typename dX>
|
||||||
|
void operator()(Device d, X x, dOut dout, dX dx, bool approximate) const {
|
||||||
|
if (approximate) {
|
||||||
|
const T kAlpha = static_cast<T>(M_2_SQRTPI * M_SQRT1_2);
|
||||||
|
const T kBeta = kAlpha * static_cast<T>(0.044715) * static_cast<T>(3);
|
||||||
|
const auto y =
|
||||||
|
(kAlpha * ((static_cast<T>(0.044715) * x.cube()) + x)).tanh();
|
||||||
|
dx.device(d) = static_cast<T>(0.5) * dout *
|
||||||
|
(static_cast<T>(1) + y +
|
||||||
|
(x - x * y.square()) * (kAlpha + kBeta * x.square()));
|
||||||
|
} else {
|
||||||
|
// gelu_grad(x) = dout * 0.5 * (1 + erf(x / sqrt(2)) + x * sqrt(2 / pi) *
|
||||||
|
// exp(- x^2 / 2)
|
||||||
|
auto first =
|
||||||
|
static_cast<T>(0.5) *
|
||||||
|
(static_cast<T>(1) + ((x * static_cast<T>(M_SQRT1_2)).erf()));
|
||||||
|
|
||||||
|
auto second = static_cast<T>(0.5 * M_2_SQRTPI * M_SQRT1_2) * x *
|
||||||
|
(-static_cast<T>(0.5) * x.square()).exp();
|
||||||
|
dx.device(d) = dout * (first + second);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <typename DeviceContext, typename T>
|
||||||
|
class GeluKernel : public framework::OpKernel<T> {
|
||||||
|
public:
|
||||||
|
void Compute(const framework::ExecutionContext& context) const override {
|
||||||
|
auto* out = context.Output<framework::Tensor>("Out");
|
||||||
|
auto* in = context.Input<framework::Tensor>("X");
|
||||||
|
auto approximate = context.Attr<bool>("approximate");
|
||||||
|
out->mutable_data<T>(in->place());
|
||||||
|
|
||||||
|
auto eigen_out = framework::EigenVector<T>::Flatten(*out);
|
||||||
|
auto eigen_in = framework::EigenVector<T>::Flatten(*in);
|
||||||
|
auto& place =
|
||||||
|
*context.template device_context<DeviceContext>().eigen_device();
|
||||||
|
|
||||||
|
GeluFunctor<T> functor;
|
||||||
|
functor(place, eigen_in, eigen_out, approximate);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
template <typename DeviceContext, typename T>
|
||||||
|
class GeluGradKernel : public framework::OpKernel<T> {
|
||||||
|
public:
|
||||||
|
void Compute(const framework::ExecutionContext& context) const override {
|
||||||
|
auto* x = context.Input<framework::Tensor>("X");
|
||||||
|
auto* dout =
|
||||||
|
context.Input<framework::Tensor>(framework::GradVarName("Out"));
|
||||||
|
auto* dx = context.Output<framework::Tensor>(framework::GradVarName("X"));
|
||||||
|
auto approximate = context.Attr<bool>("approximate");
|
||||||
|
dx->mutable_data<T>(dout->place());
|
||||||
|
|
||||||
|
auto eigen_x = framework::EigenVector<T>::Flatten(*x);
|
||||||
|
auto eigen_dout = framework::EigenVector<T>::Flatten(*dout);
|
||||||
|
auto eigen_dx = framework::EigenVector<T>::Flatten(*dx);
|
||||||
|
auto& place =
|
||||||
|
*context.template device_context<DeviceContext>().eigen_device();
|
||||||
|
|
||||||
|
GeluGradFunctor<T> functor;
|
||||||
|
functor(place, eigen_x, eigen_dout, eigen_dx, approximate);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
} // namespace operators
|
||||||
|
} // namespace paddle
|
||||||
Loading…
Reference in new issue