wehub-resource-sync
694 lines
26 KiB
C++
694 lines
26 KiB
C++
// Copyright (c) 2022 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 "glog/logging.h"
|
|
|
|
#include "paddle/phi/backends/cpu/cpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/batch_norm_kernel.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/batch_norm_utils.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T>
|
|
using EigenArrayMap =
|
|
Eigen::Map<Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>>;
|
|
template <typename T>
|
|
using ConstEigenArrayMap =
|
|
Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, Eigen::Dynamic>>;
|
|
template <typename T>
|
|
using EigenVectorArrayMap = Eigen::Map<Eigen::Array<T, Eigen::Dynamic, 1>>;
|
|
template <typename T>
|
|
using ConstEigenVectorArrayMap =
|
|
Eigen::Map<const Eigen::Array<T, Eigen::Dynamic, 1>>;
|
|
|
|
template <typename T, typename Context>
|
|
void BatchNormGradFunctor(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& scale,
|
|
const optional<DenseTensor>& bias,
|
|
const optional<DenseTensor>& mean,
|
|
const optional<DenseTensor>& variance,
|
|
const DenseTensor& saved_mean,
|
|
const DenseTensor& saved_variance,
|
|
const optional<DenseTensor>& reserve_space,
|
|
const DenseTensor& y_grad,
|
|
float momentum,
|
|
float epsilon,
|
|
const std::string& data_layout_str,
|
|
bool is_test,
|
|
bool use_global_stats,
|
|
bool trainable_statistics,
|
|
bool is_inplace,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* scale_grad,
|
|
DenseTensor* bias_grad) {
|
|
const auto* d_y = &y_grad;
|
|
|
|
DataLayout data_layout = StringToDataLayout(data_layout_str);
|
|
|
|
auto* d_x = x_grad;
|
|
auto* d_scale = scale_grad;
|
|
auto* d_bias = bias_grad;
|
|
|
|
use_global_stats = is_test || use_global_stats;
|
|
|
|
// batch_norm with inplace as false will take X as grad input, which
|
|
// is same as cuDNN batch_norm backward calculation, batch_norm
|
|
// with inplace as true only take Y as input and X should be calculate
|
|
// by inverse operation of batch_norm on Y
|
|
|
|
if (is_inplace) {
|
|
if (d_x) {
|
|
PADDLE_ENFORCE_EQ(d_x,
|
|
d_y,
|
|
common::errors::InvalidArgument(
|
|
"X@GRAD and Y@GRAD inplaced in non-inplace mode"));
|
|
}
|
|
} else {
|
|
if (d_x) {
|
|
PADDLE_ENFORCE_NE(d_x,
|
|
d_y,
|
|
common::errors::InvalidArgument(
|
|
"X@GRAD and Y@GRAD inplaced in non-inplace mode"));
|
|
}
|
|
}
|
|
|
|
// Get the size for each dimension.
|
|
// NCHW [batch_size, in_channels, in_height, in_width]
|
|
const auto& x_dims = x.dims();
|
|
PADDLE_ENFORCE_GE(
|
|
x_dims.size(),
|
|
2,
|
|
common::errors::InvalidArgument(
|
|
"The size of input X's dimensions should be larger than 1."
|
|
"But received: the size of input X's dimensions is [%d]",
|
|
x_dims.size()));
|
|
PADDLE_ENFORCE_LE(
|
|
x_dims.size(),
|
|
5,
|
|
common::errors::InvalidArgument(
|
|
"The size of input X's dimensions should be less than 6."
|
|
"But received: the size of input X's dimensions is [%d]",
|
|
x_dims.size()));
|
|
const int64_t N = x_dims[0];
|
|
const int64_t C =
|
|
data_layout == DataLayout::NCHW ? x_dims[1] : x_dims[x_dims.size() - 1];
|
|
const int64_t sample_size = x.numel() / N / C;
|
|
const int64_t num_batch_channels = N * C;
|
|
const int64_t num_batch_spatial = N * sample_size;
|
|
|
|
// input dimension is 2 and the format is NCHW. The input can be regarded as
|
|
// NHWC format
|
|
if (x_dims.size() == 2 && data_layout == DataLayout::NCHW) {
|
|
data_layout = DataLayout::NHWC;
|
|
}
|
|
|
|
// init output
|
|
if (d_x) {
|
|
dev_ctx.template Alloc<T>(d_x);
|
|
}
|
|
|
|
const T* mean_data = nullptr;
|
|
const T* inv_var_data = nullptr;
|
|
DenseTensor inv_var_tensor;
|
|
if (use_global_stats) {
|
|
const auto* running_mean = mean.get_ptr();
|
|
const auto* running_variance = variance.get_ptr();
|
|
mean_data = running_mean->data<T>();
|
|
inv_var_tensor.Resize({C});
|
|
T* running_inv_var_data = dev_ctx.template Alloc<T>(&inv_var_tensor);
|
|
EigenVectorArrayMap<T> inv_var_tmp(running_inv_var_data, C);
|
|
ConstEigenVectorArrayMap<T> var_arr(running_variance->data<T>(), C);
|
|
|
|
inv_var_tmp = (var_arr + epsilon).sqrt().inverse();
|
|
inv_var_data = running_inv_var_data;
|
|
} else {
|
|
mean_data = saved_mean.data<T>();
|
|
inv_var_data = saved_variance.data<T>();
|
|
}
|
|
|
|
ConstEigenVectorArrayMap<T> mean_arr(mean_data, C);
|
|
ConstEigenVectorArrayMap<T> inv_var_arr(inv_var_data, C);
|
|
|
|
T* d_bias_data = nullptr;
|
|
T* d_scale_data = nullptr;
|
|
if (d_scale && d_bias) {
|
|
d_bias_data = dev_ctx.template Alloc<T>(d_bias);
|
|
d_scale_data = dev_ctx.template Alloc<T>(d_scale);
|
|
}
|
|
|
|
// d_bias = np.sum(d_y, axis=0)
|
|
// d_scale = np.sum((X - mean) / inv_std * dy, axis=0)
|
|
// d_x = (1. / N) * scale * inv_var * (N * d_y - np.sum(d_y, axis=0)
|
|
// - (X - mean) * inv_var * inv_var * np.sum(d_y * (X - mean), axis=0))
|
|
EigenVectorArrayMap<T> d_bias_arr(d_bias_data, C);
|
|
EigenVectorArrayMap<T> d_scale_arr(d_scale_data, C);
|
|
|
|
if (d_scale && d_bias) {
|
|
d_bias_arr.setZero();
|
|
d_scale_arr.setZero();
|
|
}
|
|
|
|
if (d_x && num_batch_spatial == 1 && !use_global_stats) {
|
|
Copy(dev_ctx, *d_y, dev_ctx.GetPlace(), false, d_x);
|
|
return;
|
|
}
|
|
auto* Scale = scale.get_ptr();
|
|
auto* Bias = bias.get_ptr();
|
|
Eigen::Array<T, Eigen::Dynamic, 1> scale_arr(C);
|
|
Eigen::Array<T, Eigen::Dynamic, 1> bias_arr(C);
|
|
if (Scale) {
|
|
scale_arr = ConstEigenVectorArrayMap<T>(Scale->data<T>(), C);
|
|
} else {
|
|
scale_arr.setOnes();
|
|
}
|
|
if (Bias) {
|
|
bias_arr = ConstEigenVectorArrayMap<T>(Bias->data<T>(), C);
|
|
} else {
|
|
bias_arr.setZero();
|
|
}
|
|
|
|
int64_t scale_coeff = use_global_stats ? 1 : num_batch_spatial;
|
|
const auto scale_inv_var_nhw = scale_arr * inv_var_arr / scale_coeff;
|
|
|
|
DenseTensor dy_sum;
|
|
dy_sum.Resize({C});
|
|
auto dy_sum_data = dev_ctx.template Alloc<T>(&dy_sum);
|
|
EigenVectorArrayMap<T> dy_sum_arr(dy_sum_data, C);
|
|
|
|
DenseTensor dy_mul_x_sub_mean_mul_invstd_sum;
|
|
dy_mul_x_sub_mean_mul_invstd_sum.Resize({C});
|
|
auto dy_mul_x_sub_mean_mul_invstd_sum_data =
|
|
dev_ctx.template Alloc<T>(&dy_mul_x_sub_mean_mul_invstd_sum);
|
|
EigenVectorArrayMap<T> dy_mul_x_sub_mean_mul_invstd_sum_arr(
|
|
dy_mul_x_sub_mean_mul_invstd_sum_data, C);
|
|
|
|
dy_sum_arr.setZero();
|
|
dy_mul_x_sub_mean_mul_invstd_sum_arr.setZero();
|
|
|
|
// inplace calculation
|
|
// Y: ((x - est_mean) * (inv_var) * scale + bias
|
|
// formula transform ====>
|
|
// (x * inv_var * scale) + (bias - est_mean * inv_var * scale)
|
|
// X: (y - bias) / scale / (inv_var) + est_mean
|
|
// formula transform ====>
|
|
// (y - bias) / (scale * inv_var) + est_mean
|
|
switch (data_layout) {
|
|
case DataLayout::NCHW: {
|
|
if (is_inplace) {
|
|
auto px = x;
|
|
EigenArrayMap<T> x_data(
|
|
dev_ctx.template Alloc<T>(&px), sample_size, num_batch_channels);
|
|
ConstEigenArrayMap<T> y_data(
|
|
x.data<T>(), sample_size, num_batch_channels);
|
|
for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
|
|
x_data.col(nc) = (y_data.col(nc) - bias_arr(nc % C)) /
|
|
scale_inv_var_nhw(nc % C) / scale_coeff +
|
|
mean_arr(nc % C);
|
|
}
|
|
}
|
|
ConstEigenArrayMap<T> x_arr(x.data<T>(), sample_size, num_batch_channels);
|
|
ConstEigenArrayMap<T> d_y_arr(
|
|
d_y->data<T>(), sample_size, num_batch_channels);
|
|
|
|
for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
|
|
int c = nc % C;
|
|
dy_sum_arr(c) += d_y_arr.col(nc).sum();
|
|
dy_mul_x_sub_mean_mul_invstd_sum_arr(c) +=
|
|
((x_arr.col(nc) - mean_arr(c)) * inv_var_arr(c) * d_y_arr.col(nc))
|
|
.sum();
|
|
}
|
|
|
|
if (d_scale && d_bias) {
|
|
d_bias_arr = dy_sum_arr;
|
|
d_scale_arr = dy_mul_x_sub_mean_mul_invstd_sum_arr;
|
|
}
|
|
|
|
if (d_x) {
|
|
EigenArrayMap<T> d_x_arr(
|
|
dev_ctx.template Alloc<T>(d_x), sample_size, num_batch_channels);
|
|
if (!use_global_stats) {
|
|
for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
|
|
int c = nc % C;
|
|
d_x_arr.col(nc) =
|
|
scale_inv_var_nhw(c) *
|
|
(d_y_arr.col(nc) * num_batch_spatial - dy_sum_arr(c) -
|
|
(x_arr.col(nc) - mean_arr[c]) *
|
|
dy_mul_x_sub_mean_mul_invstd_sum_arr(c) * inv_var_arr(c));
|
|
}
|
|
} else {
|
|
for (int64_t nc = 0; nc < num_batch_channels; ++nc) {
|
|
int c = nc % C;
|
|
d_x_arr.col(nc) = scale_inv_var_nhw(c) * d_y_arr.col(nc);
|
|
}
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
case DataLayout::NHWC: {
|
|
if (is_inplace) {
|
|
auto px = x;
|
|
EigenArrayMap<T> x_data(
|
|
dev_ctx.template Alloc<T>(&px), C, num_batch_spatial);
|
|
ConstEigenArrayMap<T> y_data(x.data<T>(), C, num_batch_spatial);
|
|
for (int64_t nhw = 0; nhw < num_batch_spatial; nhw++) {
|
|
x_data.col(nhw) =
|
|
(y_data.col(nhw) - bias_arr) / scale_inv_var_nhw / scale_coeff +
|
|
mean_arr;
|
|
}
|
|
}
|
|
ConstEigenArrayMap<T> x_arr(x.data<T>(), C, num_batch_spatial);
|
|
ConstEigenArrayMap<T> d_y_arr(d_y->data<T>(), C, num_batch_spatial);
|
|
|
|
for (int64_t nhw = 0; nhw < num_batch_spatial; ++nhw) {
|
|
dy_sum_arr += d_y_arr.col(nhw);
|
|
dy_mul_x_sub_mean_mul_invstd_sum_arr +=
|
|
(x_arr.col(nhw) - mean_arr) * inv_var_arr * d_y_arr.col(nhw);
|
|
}
|
|
|
|
if (d_scale && d_bias) {
|
|
d_bias_arr = dy_sum_arr;
|
|
d_scale_arr = dy_mul_x_sub_mean_mul_invstd_sum_arr;
|
|
}
|
|
|
|
if (d_x) {
|
|
EigenArrayMap<T> d_x_arr(
|
|
dev_ctx.template Alloc<T>(d_x), C, num_batch_spatial);
|
|
if (!use_global_stats) {
|
|
for (int64_t nhw = 0; nhw < num_batch_spatial; ++nhw) {
|
|
d_x_arr.col(nhw) =
|
|
scale_inv_var_nhw *
|
|
(d_y_arr.col(nhw) * num_batch_spatial - dy_sum_arr -
|
|
(x_arr.col(nhw) - mean_arr) *
|
|
dy_mul_x_sub_mean_mul_invstd_sum_arr * inv_var_arr);
|
|
}
|
|
} else {
|
|
for (int64_t nhw = 0; nhw < num_batch_spatial; ++nhw) {
|
|
d_x_arr.col(nhw) = scale_inv_var_nhw * d_y_arr.col(nhw);
|
|
}
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
default:
|
|
PADDLE_THROW(common::errors::InvalidArgument("Unknown storage order: %s",
|
|
data_layout_str));
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void BatchNormGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& scale,
|
|
const optional<DenseTensor>& bias,
|
|
const optional<DenseTensor>& mean,
|
|
const optional<DenseTensor>& variance,
|
|
const DenseTensor& saved_mean,
|
|
const DenseTensor& saved_variance,
|
|
const optional<DenseTensor>& reserve_space,
|
|
const DenseTensor& y_grad,
|
|
float momentum,
|
|
float epsilon,
|
|
const std::string& data_layout,
|
|
bool is_test,
|
|
bool use_global_stats,
|
|
bool trainable_statistics,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* scale_grad,
|
|
DenseTensor* bias_grad) {
|
|
if (x.numel() == 0) {
|
|
dev_ctx.template Alloc<T>(x_grad);
|
|
if (scale_grad)
|
|
Full<T, Context>(dev_ctx, scale_grad->dims(), 0, scale_grad);
|
|
if (bias_grad) Full<T, Context>(dev_ctx, bias_grad->dims(), 0, bias_grad);
|
|
return;
|
|
}
|
|
BatchNormGradFunctor<T, Context>(dev_ctx,
|
|
x,
|
|
scale,
|
|
bias,
|
|
mean,
|
|
variance,
|
|
saved_mean,
|
|
saved_variance,
|
|
reserve_space,
|
|
y_grad,
|
|
momentum,
|
|
epsilon,
|
|
data_layout,
|
|
is_test,
|
|
use_global_stats,
|
|
trainable_statistics,
|
|
false,
|
|
x_grad,
|
|
scale_grad,
|
|
bias_grad);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void BatchNormDoubleGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& scale,
|
|
const optional<DenseTensor>& mean,
|
|
const optional<DenseTensor>& variance,
|
|
const DenseTensor& saved_mean,
|
|
const DenseTensor& saved_variance,
|
|
const DenseTensor& y_grad,
|
|
const optional<DenseTensor>& x_grad_grad,
|
|
const optional<DenseTensor>& scale_grad_grad,
|
|
const optional<DenseTensor>& bias_grad_grad,
|
|
float momentum,
|
|
float epsilon,
|
|
const std::string& data_layout_str,
|
|
bool is_test,
|
|
bool use_global_stats,
|
|
bool trainable_statistics,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* scale_grad,
|
|
DenseTensor* y_grad_grad) {
|
|
const auto* X = &x;
|
|
const auto* Scale = scale.get_ptr();
|
|
const auto* dY = &y_grad;
|
|
const auto* Saved_mean = &saved_mean;
|
|
const auto* Saved_variance = &saved_variance;
|
|
|
|
PADDLE_ENFORCE_EQ(is_test,
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"`is_test = True` CANNOT be used in train program. If "
|
|
"you want to use global status in pre_train model, "
|
|
"please set `use_global_stats = True`"));
|
|
|
|
const auto data_layout = StringToDataLayout(data_layout_str);
|
|
|
|
const auto* ddX = x_grad_grad.get_ptr();
|
|
const auto* ddScale = scale_grad_grad.get_ptr();
|
|
const auto* ddBias = bias_grad_grad.get_ptr();
|
|
|
|
auto* dX = x_grad;
|
|
auto* dScale = scale_grad;
|
|
auto* ddY = y_grad_grad;
|
|
dev_ctx.template Alloc<T>(dX);
|
|
dev_ctx.template Alloc<T>(ddY);
|
|
|
|
const auto& x_dims = X->dims();
|
|
const int64_t C =
|
|
data_layout == DataLayout::NCHW ? x_dims[1] : x_dims[x_dims.size() - 1];
|
|
const int64_t sample_size = X->numel() / C;
|
|
funcs::SetConstant<Context, T> set_constant;
|
|
|
|
const T* mean_data = Saved_mean->data<T>();
|
|
const T* inv_var_data = Saved_variance->data<T>();
|
|
|
|
DenseTensor inv_var_tensor;
|
|
if (use_global_stats) {
|
|
const auto* running_mean = mean.get_ptr();
|
|
const auto* running_variance = variance.get_ptr();
|
|
mean_data = running_mean->data<T>();
|
|
inv_var_tensor.Resize({C});
|
|
|
|
T* running_inv_var_data = dev_ctx.template Alloc<T>(&inv_var_tensor);
|
|
EigenVectorArrayMap<T> inv_var_tmp(running_inv_var_data, C);
|
|
ConstEigenVectorArrayMap<T> var_arr(running_variance->data<T>(), C);
|
|
|
|
inv_var_tmp = (var_arr + epsilon).sqrt().inverse();
|
|
inv_var_data = running_inv_var_data;
|
|
}
|
|
|
|
// transpose NCHW -> NHWC for easy calculate
|
|
DenseTensor transformed_x(X->type());
|
|
DenseTensor transformed_dy(dY->type());
|
|
DenseTensor transformed_ddx(ddX->type());
|
|
|
|
DenseTensor transformed_dx(dX->type());
|
|
DenseTensor transformed_ddy(ddY->type());
|
|
if (data_layout == DataLayout::NCHW && x_dims.size() > 2) {
|
|
VLOG(3) << "Transform batchnorm output from NCHW to NHWC";
|
|
// Input DenseTensor
|
|
ResizeToChannelLast<Context, T>(dev_ctx, X, &transformed_x);
|
|
TransToChannelLast<Context, T>(dev_ctx, X, &transformed_x);
|
|
ResizeToChannelLast<Context, T>(dev_ctx, dY, &transformed_dy);
|
|
TransToChannelLast<Context, T>(dev_ctx, dY, &transformed_dy);
|
|
ResizeToChannelLast<Context, T>(dev_ctx, ddX, &transformed_ddx);
|
|
TransToChannelLast<Context, T>(dev_ctx, ddX, &transformed_ddx);
|
|
// Output DenseTensor
|
|
ResizeToChannelLast<Context, T>(dev_ctx, dX, &transformed_dx);
|
|
ResizeToChannelLast<Context, T>(dev_ctx, ddY, &transformed_ddy);
|
|
} else {
|
|
transformed_x.ShareDataWith(*X);
|
|
transformed_dy.ShareDataWith(*dY);
|
|
transformed_ddx.ShareDataWith(*ddX);
|
|
|
|
transformed_dx.ShareDataWith(*dX);
|
|
transformed_ddy.ShareDataWith(*ddY);
|
|
}
|
|
|
|
ConstEigenArrayMap<T> x_arr(transformed_x.data<T>(), C, sample_size);
|
|
ConstEigenVectorArrayMap<T> mean_arr(mean_data, C);
|
|
ConstEigenVectorArrayMap<T> inv_var_arr(inv_var_data, C);
|
|
|
|
DenseTensor mean_tile;
|
|
mean_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> mean_tile_data(
|
|
dev_ctx.template Alloc<T>(&mean_tile), C, sample_size);
|
|
|
|
DenseTensor inv_var_tile;
|
|
inv_var_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> inv_var_tile_data(
|
|
dev_ctx.template Alloc<T>(&inv_var_tile), C, sample_size);
|
|
|
|
mean_tile_data = mean_arr.replicate(1, sample_size);
|
|
inv_var_tile_data = inv_var_arr.replicate(1, sample_size);
|
|
|
|
DenseTensor Scale_data;
|
|
if (!Scale) {
|
|
Scale_data.Resize({C});
|
|
dev_ctx.template Alloc<T>(&Scale_data);
|
|
set_constant(dev_ctx, &Scale_data, static_cast<T>(1));
|
|
}
|
|
ConstEigenVectorArrayMap<T> scale_arr(
|
|
Scale ? Scale->data<T>() : Scale_data.data<T>(), C);
|
|
|
|
DenseTensor scale_tile;
|
|
scale_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> scale_tile_data(
|
|
dev_ctx.template Alloc<T>(&scale_tile), C, sample_size);
|
|
scale_tile_data = scale_arr.replicate(1, sample_size);
|
|
|
|
ConstEigenArrayMap<T> dy_arr(transformed_dy.data<T>(), C, sample_size);
|
|
ConstEigenArrayMap<T> ddx_arr(transformed_ddx.data<T>(), C, sample_size);
|
|
|
|
DenseTensor x_sub_mean_mul_invstd;
|
|
x_sub_mean_mul_invstd.Resize({C, sample_size});
|
|
|
|
EigenArrayMap<T> x_sub_mean_mul_invstd_arr(
|
|
dev_ctx.template Alloc<T>(&x_sub_mean_mul_invstd), C, sample_size);
|
|
x_sub_mean_mul_invstd_arr = (x_arr - mean_tile_data) * inv_var_tile_data;
|
|
|
|
if (dX) {
|
|
dev_ctx.template Alloc<T>(dX);
|
|
EigenArrayMap<T> dx_arr(
|
|
dev_ctx.template Alloc<T>(&transformed_dx), C, sample_size);
|
|
dx_arr.setZero();
|
|
if (use_global_stats) {
|
|
// math: dx = (ddscale * dy) * inv_var
|
|
if (ddScale) {
|
|
ConstEigenVectorArrayMap<T> ddscale_arr(ddScale->data<T>(), C);
|
|
DenseTensor ddscale_tile;
|
|
ddscale_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> ddscale_tile_data(
|
|
dev_ctx.template Alloc<T>(&ddscale_tile), C, sample_size);
|
|
ddscale_tile_data = ddscale_arr.replicate(1, sample_size);
|
|
|
|
dx_arr = dy_arr * ddscale_tile_data * inv_var_tile_data;
|
|
}
|
|
} else {
|
|
// math: dx = scale * ((x - mean) * inv_var / NxHxW * (np.mean(ddx,
|
|
// axis=(n,h,w)) *
|
|
// np.sum(dy, axis=(n,h,w)) -
|
|
// np.sum(dy * ddx, axis=(n,h,w)) + 3 * np.mean(dy * (x -
|
|
// mean),
|
|
// axis=(n,h,w)) * inv_var.pow(2) *
|
|
// np.sum(ddx * (x - mean), axis=(n,h,w))) + inv_var.pow(3) /
|
|
// NxHxW *
|
|
// np.sum(ddx * (x - mean)) *
|
|
// (np.mean(dy, axis=(n,h,w)) - dy) + inv_var.pow(3) / NxHxW *
|
|
// np.sum(dy,
|
|
// axis=(n,h,w)) * (x - mean) *
|
|
// (np.mean(ddx, axis=(n,h,w)) - ddx)) + ddr * (dy * inv_var -
|
|
// inv_var
|
|
// *
|
|
// np.mean(dy, axis=(n,h,w)) -
|
|
// inv_var.pow(3) * (x - mean) * np.mean(dy * (x - mean),
|
|
// axis=(n,h,w)))
|
|
|
|
if (ddX) {
|
|
dx_arr +=
|
|
(x_sub_mean_mul_invstd_arr * inv_var_tile_data * inv_var_tile_data /
|
|
sample_size)
|
|
.colwise() *
|
|
(ddx_arr.rowwise().sum() * dy_arr.rowwise().sum() / sample_size -
|
|
(dy_arr * ddx_arr).rowwise().sum() +
|
|
3. * (dy_arr * x_sub_mean_mul_invstd_arr).rowwise().sum() *
|
|
(ddx_arr * x_sub_mean_mul_invstd_arr).rowwise().sum() /
|
|
sample_size);
|
|
|
|
dx_arr += (inv_var_tile_data * inv_var_tile_data).colwise() *
|
|
(ddx_arr * x_sub_mean_mul_invstd_arr).rowwise().sum() /
|
|
sample_size * (dy_arr.rowwise().sum() / sample_size - dy_arr);
|
|
|
|
dx_arr += (inv_var_tile_data * inv_var_tile_data).colwise() *
|
|
(dy_arr * x_sub_mean_mul_invstd_arr).rowwise().sum() /
|
|
sample_size *
|
|
(ddx_arr.rowwise().sum() / sample_size - ddx_arr);
|
|
|
|
dx_arr = scale_tile_data * dx_arr;
|
|
}
|
|
if (ddScale) {
|
|
ConstEigenVectorArrayMap<T> ddscale_arr(ddScale->data<T>(), C);
|
|
DenseTensor ddscale_tile;
|
|
ddscale_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> ddscale_tile_data(
|
|
dev_ctx.template Alloc<T>(&ddscale_tile), C, sample_size);
|
|
ddscale_tile_data = ddscale_arr.replicate(1, sample_size);
|
|
|
|
dx_arr +=
|
|
(dy_arr * inv_var_tile_data -
|
|
(dy_arr.rowwise().sum().replicate(1, sample_size) / sample_size) *
|
|
inv_var_tile_data -
|
|
x_sub_mean_mul_invstd_arr * inv_var_tile_data *
|
|
(dy_arr * x_sub_mean_mul_invstd_arr)
|
|
.rowwise()
|
|
.sum()
|
|
.replicate(1, sample_size) /
|
|
sample_size) *
|
|
ddscale_tile_data;
|
|
}
|
|
}
|
|
if (data_layout == DataLayout::NCHW) {
|
|
VLOG(3) << "Transform batchnorm output from NHWC to NCHW";
|
|
TransToChannelFirst<Context, T>(dev_ctx, &transformed_dx, dX);
|
|
}
|
|
}
|
|
if (dScale) {
|
|
EigenVectorArrayMap<T> dscale_arr(dev_ctx.template Alloc<T>(dScale), C);
|
|
dscale_arr.setZero();
|
|
if (use_global_stats) {
|
|
// math: dscale = np.sum(ddx * dy, axis=(n,h,w)) * inv_var
|
|
if (ddX) {
|
|
dscale_arr = (ddx_arr * dy_arr * inv_var_tile_data).rowwise().sum();
|
|
}
|
|
} else {
|
|
// math: dscale = inv_var * (dy - np.mean(dy, axis=(n,h,w) - (x-mean) *
|
|
// inv_var.pow(2) * np.mean(dy * (x-mean), axis=(n,h,w)))) *
|
|
// ddx
|
|
if (ddX) {
|
|
DenseTensor first_grad;
|
|
first_grad.Resize({C, sample_size});
|
|
EigenArrayMap<T> first_grad_arr(
|
|
dev_ctx.template Alloc<T>(&first_grad), C, sample_size);
|
|
first_grad_arr.setZero();
|
|
|
|
first_grad_arr +=
|
|
inv_var_tile_data *
|
|
(dy_arr -
|
|
dy_arr.rowwise().sum().replicate(1, sample_size) / sample_size -
|
|
x_sub_mean_mul_invstd_arr *
|
|
(dy_arr * x_sub_mean_mul_invstd_arr)
|
|
.rowwise()
|
|
.sum()
|
|
.replicate(1, sample_size) /
|
|
sample_size);
|
|
dscale_arr = (first_grad_arr * ddx_arr).rowwise().sum();
|
|
}
|
|
}
|
|
}
|
|
|
|
if (ddY) {
|
|
dev_ctx.template Alloc<T>(ddY);
|
|
EigenArrayMap<T> ddy_arr(
|
|
dev_ctx.template Alloc<T>(&transformed_ddy), C, sample_size);
|
|
ddy_arr.setZero();
|
|
if (use_global_stats) { // NOLINT
|
|
// math: ddy = r * ddx * inv_var + ddbias +
|
|
// ddscale * (x - mean) * inv_var
|
|
if (ddX) {
|
|
ddy_arr = scale_tile_data * ddx_arr * inv_var_tile_data;
|
|
}
|
|
} else {
|
|
// math: ddy = (x - mean) * inv_var * ddscale + ddbias +
|
|
// scale * inv_var * (ddx - (x - mean) * inv_var.pow(2) *
|
|
// np.mean(ddx * (x - mean), axis=(n,h,w)))
|
|
if (ddX) {
|
|
ddy_arr +=
|
|
scale_tile_data * inv_var_tile_data *
|
|
(ddx_arr -
|
|
ddx_arr.rowwise().sum().replicate(1, sample_size) / sample_size -
|
|
x_sub_mean_mul_invstd_arr *
|
|
(ddx_arr * x_sub_mean_mul_invstd_arr)
|
|
.rowwise()
|
|
.sum()
|
|
.replicate(1, sample_size) /
|
|
sample_size);
|
|
}
|
|
}
|
|
if (ddScale) {
|
|
ConstEigenVectorArrayMap<T> ddscale_arr(ddScale->data<T>(), C);
|
|
DenseTensor ddscale_tile;
|
|
ddscale_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> ddscale_tile_data(
|
|
dev_ctx.template Alloc<T>(&ddscale_tile), C, sample_size);
|
|
ddscale_tile_data = ddscale_arr.replicate(1, sample_size);
|
|
|
|
ddy_arr += x_sub_mean_mul_invstd_arr * ddscale_tile_data;
|
|
}
|
|
|
|
if (ddBias) {
|
|
ConstEigenVectorArrayMap<T> ddbias_arr(ddBias->data<T>(), C);
|
|
DenseTensor ddbias_tile;
|
|
ddbias_tile.Resize({C, sample_size});
|
|
EigenArrayMap<T> ddbias_tile_data(
|
|
dev_ctx.template Alloc<T>(&ddbias_tile), C, sample_size);
|
|
ddbias_tile_data = ddbias_arr.replicate(1, sample_size);
|
|
|
|
ddy_arr += ddbias_tile_data;
|
|
}
|
|
|
|
if (data_layout == DataLayout::NCHW) {
|
|
VLOG(3) << "Transform batchnorm output from NHWC to NCHW";
|
|
TransToChannelFirst<Context, T>(dev_ctx, &transformed_ddy, ddY);
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_DECLARE_BN_GRAD_FUNCTOR(float, CPU);
|
|
PD_DECLARE_BN_GRAD_FUNCTOR(double, CPU);
|
|
|
|
PD_REGISTER_KERNEL(
|
|
batch_norm_grad, CPU, ALL_LAYOUT, phi::BatchNormGradKernel, float, double) {
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(batch_norm_double_grad,
|
|
CPU,
|
|
ALL_LAYOUT,
|
|
phi::BatchNormDoubleGradKernel,
|
|
float,
|
|
double) {}
|