提交 67731297 编写于 作者: C chengduoZH

code refine

上级 df0e74db
...@@ -97,15 +97,15 @@ class LayerNormKernel : public framework::OpKernel<T> { ...@@ -97,15 +97,15 @@ class LayerNormKernel : public framework::OpKernel<T> {
auto &dev_ctx = ctx.template device_context<DeviceContext>(); auto &dev_ctx = ctx.template device_context<DeviceContext>();
math::RowwiseMean<DeviceContext, T> row_mean; math::RowwiseMean<DeviceContext, T> row_mean;
// functor-> get mean // get mean
row_mean(dev_ctx, x, mean); row_mean(dev_ctx, x, mean);
// functor-> get variance // get variance
ElementwiseComputeEx<SubAndSquareFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<SubAndSquareFunctor<T>, DeviceContext, T>(
ctx, &x, mean, /*axis*/ 0, SubAndSquareFunctor<T>(), &out); ctx, &x, mean, /*axis*/ 0, SubAndSquareFunctor<T>(), &out);
row_mean(dev_ctx, out, var); row_mean(dev_ctx, out, var);
// functor-> get norm_out // get x_norm
ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
ctx, &x, mean, /*axis*/ 0, SubFunctor<T>(), &out); ctx, &x, mean, /*axis*/ 0, SubFunctor<T>(), &out);
ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
...@@ -129,9 +129,11 @@ class LayerNormGradKernel : public framework::OpKernel<T> { ...@@ -129,9 +129,11 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
void Compute(const framework::ExecutionContext &ctx) const override { void Compute(const framework::ExecutionContext &ctx) const override {
const float epsilon = ctx.Attr<float>("epsilon"); const float epsilon = ctx.Attr<float>("epsilon");
auto x = *ctx.Input<Tensor>("X"); auto x = *ctx.Input<Tensor>("X");
auto mean = *ctx.Input<Tensor>("Mean"); auto *y = ctx.Input<Tensor>("Y");
auto var = *ctx.Input<Tensor>("Variance"); auto *mean = ctx.Input<Tensor>("Mean");
auto scale = *ctx.Input<Tensor>("Scale"); auto *var = ctx.Input<Tensor>("Variance");
auto *scale = ctx.Input<Tensor>("Scale");
auto *bias = ctx.Input<Tensor>("Bias");
auto d_y = *ctx.Input<Tensor>(framework::GradVarName("Y")); auto d_y = *ctx.Input<Tensor>(framework::GradVarName("Y"));
const auto begin_norm_axis = ctx.Attr<int>("begin_norm_axis"); const auto begin_norm_axis = ctx.Attr<int>("begin_norm_axis");
...@@ -155,15 +157,20 @@ class LayerNormGradKernel : public framework::OpKernel<T> { ...@@ -155,15 +157,20 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
if (d_scale || d_x) { if (d_scale || d_x) {
x.Resize(matrix_shape); x.Resize(matrix_shape);
temp.mutable_data<T>(matrix_shape, ctx.GetPlace()); temp.mutable_data<T>(matrix_shape, ctx.GetPlace());
temp_norm.mutable_data<T>(matrix_shape, ctx.GetPlace());
if (!(bias && scale)) {
temp_norm.ShareDataWith(*y);
temp_norm.Resize(matrix_shape);
} else {
temp_norm.mutable_data<T>(matrix_shape, ctx.GetPlace());
// get x_norm // get x_norm
ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
ctx, &x, &mean, /*axis*/ 0, SubFunctor<T>(), &temp_norm); ctx, &x, mean, /*axis*/ 0, SubFunctor<T>(), &temp_norm);
ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
ctx, &temp_norm, &var, /*axis*/ 0, ctx, &temp_norm, var, /*axis*/ 0,
DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), &temp_norm); DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), &temp_norm);
} }
}
if (d_bias) { if (d_bias) {
d_bias->mutable_data<T>(ctx.GetPlace()); d_bias->mutable_data<T>(ctx.GetPlace());
...@@ -188,7 +195,7 @@ class LayerNormGradKernel : public framework::OpKernel<T> { ...@@ -188,7 +195,7 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
if (d_scale) { if (d_scale) {
// dy_dx // dy_dx
ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
ctx, &d_y, &scale, /*axis*/ 1, MulFunctor<T>(), &temp); ctx, &d_y, scale, /*axis*/ 1, MulFunctor<T>(), &temp);
framework::Copy(temp, ctx.GetPlace(), ctx.device_context(), d_x); framework::Copy(temp, ctx.GetPlace(), ctx.device_context(), d_x);
// dy_dmean_dx // dy_dmean_dx
...@@ -199,7 +206,6 @@ class LayerNormGradKernel : public framework::OpKernel<T> { ...@@ -199,7 +206,6 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
// dy_var_dx // dy_var_dx
ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
ctx, &temp, &temp_norm, /*axis*/ 0, MulFunctor<T>(), &temp); ctx, &temp, &temp_norm, /*axis*/ 0, MulFunctor<T>(), &temp);
} else { } else {
// dy_dx // dy_dx
framework::Copy(d_y, ctx.GetPlace(), ctx.device_context(), d_x); framework::Copy(d_y, ctx.GetPlace(), ctx.device_context(), d_x);
...@@ -216,12 +222,12 @@ class LayerNormGradKernel : public framework::OpKernel<T> { ...@@ -216,12 +222,12 @@ class LayerNormGradKernel : public framework::OpKernel<T> {
// dy_var_dx // dy_var_dx
row_mean(dev_ctx, temp, &temp_vec); row_mean(dev_ctx, temp, &temp_vec);
ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<MulFunctor<T>, DeviceContext, T>(
ctx, &temp_norm, &temp_vec, /*axis*/ 0, MulFunctor<T>(), &temp_norm); ctx, &temp_norm, &temp_vec, /*axis*/ 0, MulFunctor<T>(), &temp);
ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<SubFunctor<T>, DeviceContext, T>(
ctx, d_x, &temp_norm, /*axis*/ 0, SubFunctor<T>(), d_x); ctx, d_x, &temp, /*axis*/ 0, SubFunctor<T>(), d_x);
ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>( ElementwiseComputeEx<DivAndSqrtFunctor<T>, DeviceContext, T>(
ctx, d_x, &var, /*axis*/ 0, ctx, d_x, var, /*axis*/ 0,
DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), d_x); DivAndSqrtFunctor<T>(static_cast<T>(epsilon)), d_x);
d_x->Resize(dx_dim); d_x->Resize(dx_dim);
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册