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
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