wehub-resource-sync
233 lines
8.3 KiB
C++
233 lines
8.3 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 "paddle/phi/kernels/group_norm_grad_kernel.h"
|
|
|
|
#include <algorithm>
|
|
#include <array>
|
|
#include <numeric>
|
|
#include <string>
|
|
|
|
#include "paddle/common/layout.h"
|
|
#include "paddle/phi/backends/cpu/cpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/extensions.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void GroupNormGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const optional<DenseTensor>& scale,
|
|
const optional<DenseTensor>& bias,
|
|
const DenseTensor& y,
|
|
const DenseTensor& mean,
|
|
const DenseTensor& var,
|
|
const DenseTensor& d_y,
|
|
double epsilon,
|
|
int groups,
|
|
const std::string& data_layout_str,
|
|
DenseTensor* d_x,
|
|
DenseTensor* d_scale,
|
|
DenseTensor* d_bias) {
|
|
if (x.numel() == 0) {
|
|
dev_ctx.template Alloc<T>(d_x);
|
|
if (d_scale) {
|
|
// If batch dim is 0, we should set d_scale to zero, or else NAN
|
|
if (x.dims().size() > 0 && x.dims()[0] == 0) {
|
|
Full<T, Context>(dev_ctx, d_scale->dims(), 0, d_scale);
|
|
|
|
} else {
|
|
Full<T, Context>(dev_ctx, d_scale->dims(), NAN, d_scale);
|
|
}
|
|
}
|
|
if (d_bias) {
|
|
Full<T, Context>(dev_ctx, d_bias->dims(), 0, d_bias);
|
|
}
|
|
return;
|
|
}
|
|
const DataLayout data_layout = StringToDataLayout(data_layout_str);
|
|
const auto scale_ptr = scale.get_ptr();
|
|
const auto bias_ptr = bias.get_ptr();
|
|
const auto& x_dims = y.dims();
|
|
const int C = static_cast<int>(
|
|
data_layout == DataLayout::NCHW ? x_dims[1] : x_dims[x_dims.size() - 1]);
|
|
const int group_size = C / groups;
|
|
|
|
funcs::SetConstant<CPUContext, T> set_zero;
|
|
|
|
auto* x_data = y.data<T>();
|
|
auto* y_data = d_y.data<T>();
|
|
auto* var_data = var.data<T>();
|
|
|
|
T* d_x_data = nullptr;
|
|
if (d_x) {
|
|
dev_ctx.template Alloc<T>(d_x);
|
|
d_x_data = d_x->data<T>();
|
|
}
|
|
|
|
T* d_scale_data = nullptr;
|
|
if (d_scale) {
|
|
dev_ctx.template Alloc<T>(d_scale);
|
|
set_zero(dev_ctx, d_scale, static_cast<T>(0));
|
|
d_scale_data = d_scale->data<T>();
|
|
}
|
|
|
|
T* d_bias_data = nullptr;
|
|
if (d_bias) {
|
|
dev_ctx.template Alloc<T>(d_bias);
|
|
set_zero(dev_ctx, d_bias, static_cast<T>(0));
|
|
d_bias_data = d_bias->data<T>();
|
|
}
|
|
|
|
const T* scale_data = nullptr;
|
|
if (scale_ptr) scale_data = scale_ptr->data<T>();
|
|
const T* bias_data = nullptr;
|
|
if (bias_ptr) bias_data = bias_ptr->data<T>();
|
|
|
|
int64_t imsize = 1;
|
|
if (data_layout == DataLayout::NCHW) {
|
|
for (int i = 2; i < x_dims.size(); ++i) {
|
|
imsize *= x_dims[i];
|
|
}
|
|
} else {
|
|
for (int i = 1; i < x_dims.size() - 1; ++i) {
|
|
imsize *= x_dims[i];
|
|
}
|
|
}
|
|
auto* iter_x_data = x_data;
|
|
auto* iter_d_x_data = d_x_data;
|
|
auto* iter_y_data = y_data;
|
|
for (int bid = 0; bid < x_dims[0]; bid++) {
|
|
for (int gid = 0; gid < groups; gid++) {
|
|
T x_var = var_data[bid * groups + gid];
|
|
T var_inv = 1.0 / sqrt(x_var + epsilon);
|
|
int64_t number = std::min(static_cast<int64_t>(group_size),
|
|
C - static_cast<int64_t>(gid) * group_size);
|
|
T number_inv = 1.0 / (number * imsize);
|
|
auto* tmp_x = iter_x_data;
|
|
auto* tmp_y = iter_y_data;
|
|
auto* tmp_d_x = iter_d_x_data;
|
|
auto* x_src_data = iter_x_data;
|
|
auto* y_src_data = iter_y_data;
|
|
auto* iter_x_data_backup = iter_x_data;
|
|
auto* iter_y_data_backup = iter_y_data;
|
|
auto* iter_d_x_data_backup = iter_d_x_data;
|
|
T dp_scale = 0, dp_bias = 0;
|
|
|
|
if (data_layout == DataLayout::NCHW) {
|
|
for (int64_t cid = 0; cid < number; cid++) {
|
|
for (int64_t imid = 0; imid < imsize;
|
|
imid++, iter_x_data++, iter_y_data++) {
|
|
T val = iter_x_data[0];
|
|
if (bias_data) val -= bias_data[gid * group_size + cid];
|
|
T dval = iter_y_data[0];
|
|
dp_scale += val * dval;
|
|
if (scale_data)
|
|
dp_bias += dval * scale_data[gid * group_size + cid];
|
|
|
|
if (scale_data && scale_data[gid * group_size + cid] != 0)
|
|
val /= scale_data[gid * group_size + cid];
|
|
if (d_bias_data) d_bias_data[gid * group_size + cid] += dval;
|
|
if (d_scale_data)
|
|
d_scale_data[gid * group_size + cid] += val * dval;
|
|
}
|
|
}
|
|
if (d_x_data) {
|
|
for (int64_t cid = 0; cid < number; cid++) {
|
|
for (int64_t imid = 0; imid < imsize;
|
|
imid++, iter_d_x_data++, tmp_x++, tmp_y++) {
|
|
T v_y = tmp_x[0];
|
|
T dly = tmp_y[0];
|
|
T dss = dp_scale;
|
|
T dbs = dp_bias;
|
|
T v_scale = 1., v_bias = 0.;
|
|
if (scale_data) v_scale = scale_data[gid * group_size + cid];
|
|
if (bias_data) v_bias = bias_data[gid * group_size + cid];
|
|
v_y -= v_bias;
|
|
if (v_scale != 0) v_y /= v_scale;
|
|
iter_d_x_data[0] =
|
|
(dly * v_scale - number_inv * dss * v_y - number_inv * dbs) *
|
|
var_inv;
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
for (int64_t cid = 0; cid < number; cid++) {
|
|
iter_x_data = x_src_data + cid;
|
|
iter_y_data = y_src_data + cid;
|
|
for (int64_t imid = 0; imid < imsize;
|
|
imid++, iter_x_data += C, iter_y_data += C) {
|
|
T val = iter_x_data[0];
|
|
if (bias_data) val -= bias_data[gid * group_size + cid];
|
|
T dval = iter_y_data[0];
|
|
dp_scale += val * dval;
|
|
if (scale_data)
|
|
dp_bias += dval * scale_data[gid * group_size + cid];
|
|
|
|
if (scale_data && scale_data[gid * group_size + cid] != 0)
|
|
val /= scale_data[gid * group_size + cid];
|
|
if (d_bias_data) d_bias_data[gid * group_size + cid] += dval;
|
|
if (d_scale_data)
|
|
d_scale_data[gid * group_size + cid] += val * dval;
|
|
}
|
|
}
|
|
if (d_x_data) {
|
|
for (int64_t cid = 0; cid < number; cid++) {
|
|
tmp_x = x_src_data + cid;
|
|
tmp_y = y_src_data + cid;
|
|
iter_d_x_data = tmp_d_x + cid;
|
|
for (int64_t imid = 0; imid < imsize;
|
|
imid++, iter_d_x_data += C, tmp_x += C, tmp_y += C) {
|
|
T v_y = tmp_x[0];
|
|
T dly = tmp_y[0];
|
|
T dss = dp_scale;
|
|
T dbs = dp_bias;
|
|
T v_scale = 1.0, v_bias = 0.;
|
|
if (scale_data) v_scale = scale_data[gid * group_size + cid];
|
|
if (bias_data) v_bias = bias_data[gid * group_size + cid];
|
|
v_y -= v_bias;
|
|
if (v_scale != 0) v_y /= v_scale;
|
|
iter_d_x_data[0] =
|
|
(dly * v_scale - number_inv * dss * v_y - number_inv * dbs) *
|
|
var_inv;
|
|
}
|
|
}
|
|
}
|
|
iter_x_data = iter_x_data_backup + group_size;
|
|
iter_y_data = iter_y_data_backup + group_size;
|
|
if (d_x_data) {
|
|
iter_d_x_data = iter_d_x_data_backup + group_size;
|
|
}
|
|
}
|
|
}
|
|
if (data_layout == DataLayout::NHWC) {
|
|
iter_x_data = x_data + static_cast<int64_t>(bid + 1) * C * imsize;
|
|
if (d_x_data) {
|
|
iter_d_x_data = d_x_data + static_cast<int64_t>(bid + 1) * C * imsize;
|
|
}
|
|
iter_y_data = y_data + static_cast<int64_t>(bid + 1) * C * imsize;
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(
|
|
group_norm_grad, CPU, ALL_LAYOUT, phi::GroupNormGradKernel, float, double) {
|
|
}
|