wehub-resource-sync
793 lines
33 KiB
C++
793 lines
33 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/grid_sample_grad_kernel.h"
|
|
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/cpu/grid_sample_utils.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/eigen/common.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
namespace phi {
|
|
|
|
template <typename T>
|
|
static inline void ClipWithMask(const CPUContext& dev_ctx,
|
|
const int max_val, // height-1 or width-1
|
|
bool align_corners,
|
|
std::string padding_mode,
|
|
DenseTensor* grid_slice,
|
|
DenseTensor* grid_scale) {
|
|
auto& place = *dev_ctx.eigen_device();
|
|
grid_scale->Resize(grid_slice->dims());
|
|
dev_ctx.Alloc<T>(grid_scale);
|
|
|
|
auto grid_slice_t = EigenTensor<T, 3>::From(*grid_slice);
|
|
auto factor = static_cast<T>(max_val * 0.5);
|
|
if (!align_corners) {
|
|
factor = static_cast<T>((max_val + 1) * 0.5);
|
|
}
|
|
auto grid_scale_t = EigenTensor<T, 3>::From(*grid_scale).setConstant(factor);
|
|
|
|
if (padding_mode == "border") {
|
|
// auto bounded_lo = grid_slice_t.cwiseMax(static_cast<T>(0));
|
|
auto res = grid_slice_t.cwiseMax(static_cast<T>(0))
|
|
.cwiseMin(static_cast<T>(max_val));
|
|
|
|
auto in_bound = (res == grid_slice_t);
|
|
grid_scale_t.device(place) = grid_scale_t * in_bound.template cast<T>();
|
|
grid_slice_t.device(place) = res;
|
|
} else if (padding_mode == "reflection") {
|
|
if (align_corners) {
|
|
auto double_range = static_cast<T>(max_val * 2);
|
|
auto is_neg = (grid_slice_t < static_cast<T>(0));
|
|
auto grid_abs = grid_slice_t.abs();
|
|
auto extra = grid_abs - (grid_abs / double_range).floor() * double_range;
|
|
auto one_more_flip = (extra > (double_range - extra));
|
|
grid_scale_t.device(place) =
|
|
grid_scale_t * ((is_neg == one_more_flip).template cast<T>() -
|
|
(is_neg != one_more_flip).template cast<T>());
|
|
grid_slice_t.device(place) = extra.cwiseMin(double_range - extra);
|
|
if (max_val == 0) {
|
|
grid_slice_t.device(place) = grid_slice_t.constant(static_cast<T>(0));
|
|
}
|
|
} else {
|
|
auto double_range = static_cast<T>((max_val + 1) * 2);
|
|
auto grid_abs = (grid_slice_t + static_cast<T>(0.5)).abs();
|
|
auto is_neg = ((grid_slice_t + static_cast<T>(0.5)) < static_cast<T>(0));
|
|
auto extra = grid_abs - (grid_abs / double_range).floor() * double_range;
|
|
auto one_more_flip = (extra > (double_range - extra));
|
|
auto reflected =
|
|
extra.cwiseMin(double_range - extra) - static_cast<T>(0.5);
|
|
auto clipped = reflected.cwiseMax(static_cast<T>(0))
|
|
.cwiseMin(static_cast<T>(max_val));
|
|
auto in_bound = (clipped == reflected).template cast<T>();
|
|
grid_scale_t.device(place) =
|
|
grid_scale_t *
|
|
((is_neg == one_more_flip).template cast<T>() -
|
|
(is_neg != one_more_flip).template cast<T>()) *
|
|
in_bound;
|
|
grid_slice_t.device(place) = clipped;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static inline void ClipWithMask3D(const CPUContext& dev_ctx,
|
|
const int max_val, // height-1 or width-1
|
|
bool align_corners,
|
|
std::string padding_mode,
|
|
DenseTensor* grid_slice,
|
|
DenseTensor* grid_scale) {
|
|
auto& place = *dev_ctx.eigen_device();
|
|
grid_scale->Resize(grid_slice->dims());
|
|
dev_ctx.Alloc<T>(grid_scale);
|
|
|
|
auto grid_slice_t = EigenTensor<T, 4>::From(*grid_slice);
|
|
auto factor = static_cast<T>(max_val * 0.5);
|
|
if (!align_corners) {
|
|
factor = static_cast<T>((max_val + 1) * 0.5);
|
|
}
|
|
auto grid_scale_t = EigenTensor<T, 4>::From(*grid_scale).setConstant(factor);
|
|
|
|
if (padding_mode == "border") {
|
|
// auto bounded_lo = grid_slice_t.cwiseMax(static_cast<T>(0));
|
|
auto res = grid_slice_t.cwiseMax(static_cast<T>(0))
|
|
.cwiseMin(static_cast<T>(max_val));
|
|
|
|
auto in_bound = (res == grid_slice_t);
|
|
grid_scale_t.device(place) = grid_scale_t * in_bound.template cast<T>();
|
|
grid_slice_t.device(place) = res;
|
|
} else if (padding_mode == "reflection") {
|
|
if (align_corners) {
|
|
auto double_range = static_cast<T>(max_val * 2);
|
|
auto is_neg = (grid_slice_t < static_cast<T>(0));
|
|
auto grid_abs = grid_slice_t.abs();
|
|
auto extra = grid_abs - (grid_abs / double_range).floor() * double_range;
|
|
auto one_more_flip = (extra > (double_range - extra));
|
|
grid_scale_t.device(place) =
|
|
grid_scale_t * ((is_neg == one_more_flip).template cast<T>() -
|
|
(is_neg != one_more_flip).template cast<T>());
|
|
grid_slice_t.device(place) = extra.cwiseMin(double_range - extra);
|
|
if (max_val == 0) {
|
|
grid_slice_t.device(place) = grid_slice_t.constant(static_cast<T>(0));
|
|
}
|
|
} else {
|
|
auto double_range = static_cast<T>((max_val + 1) * 2);
|
|
auto grid_abs = (grid_slice_t + static_cast<T>(0.5)).abs();
|
|
auto is_neg = ((grid_slice_t + static_cast<T>(0.5)) < static_cast<T>(0));
|
|
auto extra = grid_abs - (grid_abs / double_range).floor() * double_range;
|
|
auto one_more_flip = (extra > (double_range - extra));
|
|
auto reflected =
|
|
extra.cwiseMin(double_range - extra) - static_cast<T>(0.5);
|
|
auto clipped = reflected.cwiseMax(static_cast<T>(0))
|
|
.cwiseMin(static_cast<T>(max_val));
|
|
auto in_bound = (clipped == reflected).template cast<T>();
|
|
grid_scale_t.device(place) =
|
|
grid_scale_t *
|
|
((is_neg == one_more_flip).template cast<T>() -
|
|
(is_neg != one_more_flip).template cast<T>()) *
|
|
in_bound;
|
|
grid_slice_t.device(place) = clipped;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void CalcGridLocationsWithGrad(const CPUContext& dev_ctx,
|
|
const DenseTensor& grid,
|
|
const int in_h,
|
|
const int in_w,
|
|
bool align_corners,
|
|
std::string padding_mode,
|
|
DenseTensor* grid_x,
|
|
DenseTensor* grid_y,
|
|
DenseTensor* grid_x_scale,
|
|
DenseTensor* grid_y_scale) {
|
|
const int n = static_cast<int>(grid.dims()[0]);
|
|
const int out_h = static_cast<int>(grid.dims()[1]);
|
|
const int out_w = static_cast<int>(grid.dims()[2]);
|
|
|
|
// split grid with shape (n, h, w, 2) into (x, y) by the 3rd Dim
|
|
grid_x->Resize({n, out_h, out_w});
|
|
grid_y->Resize({n, out_h, out_w});
|
|
T* grid_x_data = dev_ctx.Alloc<T>(grid_x);
|
|
T* grid_y_data = dev_ctx.Alloc<T>(grid_y);
|
|
|
|
const T* grid_data = grid.data<T>();
|
|
for (int64_t i = 0; i < static_cast<int64_t>(n) * out_h * out_w; i++) {
|
|
grid_x_data[i] = grid_data[2 * i];
|
|
grid_y_data[i] = grid_data[(2 * i) + 1];
|
|
}
|
|
|
|
Unnormalize<T>(dev_ctx, grid_x, in_w - 1, align_corners);
|
|
Unnormalize<T>(dev_ctx, grid_y, in_h - 1, align_corners);
|
|
|
|
ClipWithMask<T>(
|
|
dev_ctx, in_w - 1, align_corners, padding_mode, grid_x, grid_x_scale);
|
|
ClipWithMask<T>(
|
|
dev_ctx, in_h - 1, align_corners, padding_mode, grid_y, grid_y_scale);
|
|
}
|
|
|
|
template <typename T>
|
|
static void Calc3DGridLocationsWithGrad(const CPUContext& dev_ctx,
|
|
const DenseTensor& grid,
|
|
const int in_d,
|
|
const int in_h,
|
|
const int in_w,
|
|
bool align_corners,
|
|
std::string padding_mode,
|
|
DenseTensor* grid_x,
|
|
DenseTensor* grid_y,
|
|
DenseTensor* grid_z,
|
|
DenseTensor* grid_x_scale,
|
|
DenseTensor* grid_y_scale,
|
|
DenseTensor* grid_z_scale) {
|
|
const int n = static_cast<int>(grid.dims()[0]);
|
|
const int out_d = static_cast<int>(grid.dims()[1]);
|
|
const int out_h = static_cast<int>(grid.dims()[2]);
|
|
const int out_w = static_cast<int>(grid.dims()[3]);
|
|
|
|
// split grid with shape (n, d, h, w, 3) into (x, y, z) by the 3rd Dim
|
|
grid_x->Resize({n, out_d, out_h, out_w});
|
|
grid_y->Resize({n, out_d, out_h, out_w});
|
|
grid_z->Resize({n, out_d, out_h, out_w});
|
|
T* grid_x_data = dev_ctx.Alloc<T>(grid_x);
|
|
T* grid_y_data = dev_ctx.Alloc<T>(grid_y);
|
|
T* grid_z_data = dev_ctx.Alloc<T>(grid_z);
|
|
|
|
const T* grid_data = grid.data<T>();
|
|
for (int64_t i = 0; i < static_cast<int64_t>(n) * out_d * out_h * out_w;
|
|
i++) {
|
|
grid_x_data[i] = grid_data[3 * i];
|
|
grid_y_data[i] = grid_data[(3 * i) + 1];
|
|
grid_z_data[i] = grid_data[(3 * i) + 2];
|
|
}
|
|
|
|
Unnormalize3D<T>(dev_ctx, grid_x, in_w - 1, align_corners);
|
|
Unnormalize3D<T>(dev_ctx, grid_y, in_h - 1, align_corners);
|
|
Unnormalize3D<T>(dev_ctx, grid_z, in_d - 1, align_corners);
|
|
|
|
ClipWithMask3D<T>(
|
|
dev_ctx, in_w - 1, align_corners, padding_mode, grid_x, grid_x_scale);
|
|
ClipWithMask3D<T>(
|
|
dev_ctx, in_h - 1, align_corners, padding_mode, grid_y, grid_y_scale);
|
|
ClipWithMask3D<T>(
|
|
dev_ctx, in_d - 1, align_corners, padding_mode, grid_z, grid_z_scale);
|
|
}
|
|
|
|
template <typename T>
|
|
static void GatherOutputGradToInputGrad(const DenseTensor& output_grad,
|
|
DenseTensor* input_grad,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& d1,
|
|
const DenseTensor& d2) {
|
|
const int n = static_cast<int>(output_grad.dims()[0]);
|
|
const int c = static_cast<int>(output_grad.dims()[1]);
|
|
const int out_h = static_cast<int>(output_grad.dims()[2]);
|
|
const int out_w = static_cast<int>(output_grad.dims()[3]);
|
|
const int in_h = static_cast<int>(input_grad->dims()[2]);
|
|
const int in_w = static_cast<int>(input_grad->dims()[3]);
|
|
auto x_t = EigenTensor<T, 3>::From(x);
|
|
auto y_t = EigenTensor<T, 3>::From(y);
|
|
auto d1_t = EigenTensor<T, 3>::From(d1);
|
|
auto d2_t = EigenTensor<T, 3>::From(d2);
|
|
auto input_grad_t = EigenTensor<T, 4>::From(*input_grad);
|
|
auto output_grad_t = EigenTensor<T, 4>::From(output_grad);
|
|
|
|
for (int i = 0; i < n; i++) {
|
|
for (int k = 0; k < out_h; k++) {
|
|
for (int l = 0; l < out_w; l++) {
|
|
if (IsInBound<int>(static_cast<int>(x_t(i, k, l)),
|
|
static_cast<int>(y_t(i, k, l)),
|
|
(in_w - 1),
|
|
(in_h - 1))) {
|
|
for (int j = 0; j < c; j++) {
|
|
input_grad_t(i,
|
|
j,
|
|
static_cast<int>(y_t(i, k, l)),
|
|
static_cast<int>(x_t(i, k, l))) +=
|
|
output_grad_t(i, j, k, l) * d1_t(i, k, l) * d2_t(i, k, l);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void Gather3DOutputGradToInputGrad(const DenseTensor& output_grad,
|
|
DenseTensor* input_grad,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& z,
|
|
const DenseTensor& d1,
|
|
const DenseTensor& d2,
|
|
const DenseTensor& d3) {
|
|
const int n = static_cast<int>(output_grad.dims()[0]);
|
|
const int c = static_cast<int>(output_grad.dims()[1]);
|
|
const int out_d = static_cast<int>(output_grad.dims()[2]);
|
|
const int out_h = static_cast<int>(output_grad.dims()[3]);
|
|
const int out_w = static_cast<int>(output_grad.dims()[4]);
|
|
const int in_d = static_cast<int>(input_grad->dims()[2]);
|
|
const int in_h = static_cast<int>(input_grad->dims()[3]);
|
|
const int in_w = static_cast<int>(input_grad->dims()[4]);
|
|
auto x_t = EigenTensor<T, 4>::From(x);
|
|
auto y_t = EigenTensor<T, 4>::From(y);
|
|
auto z_t = EigenTensor<T, 4>::From(z);
|
|
auto d1_t = EigenTensor<T, 4>::From(d1);
|
|
auto d2_t = EigenTensor<T, 4>::From(d2);
|
|
auto d3_t = EigenTensor<T, 4>::From(d3);
|
|
auto input_grad_t = EigenTensor<T, 5>::From(*input_grad);
|
|
auto output_grad_t = EigenTensor<T, 5>::From(output_grad);
|
|
|
|
for (int i = 0; i < n; i++) {
|
|
for (int m = 0; m < out_d; m++) {
|
|
for (int k = 0; k < out_h; k++) {
|
|
for (int l = 0; l < out_w; l++) {
|
|
if (IsInBound3D<int>(static_cast<int>(x_t(i, m, k, l)),
|
|
static_cast<int>(y_t(i, m, k, l)),
|
|
static_cast<int>(z_t(i, m, k, l)),
|
|
(in_w - 1),
|
|
(in_h - 1),
|
|
(in_d - 1))) {
|
|
for (int j = 0; j < c; j++) {
|
|
input_grad_t(i,
|
|
j,
|
|
static_cast<int>(z_t(i, m, k, l)),
|
|
static_cast<int>(y_t(i, m, k, l)),
|
|
static_cast<int>(x_t(i, m, k, l))) +=
|
|
output_grad_t(i, j, m, k, l) * d1_t(i, m, k, l) *
|
|
d2_t(i, m, k, l) * d3_t(i, m, k, l);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void GatherBilinearGrad(const CPUContext& dev_ctx,
|
|
const DenseTensor& input,
|
|
const DenseTensor& output_grad,
|
|
DenseTensor* grid_x,
|
|
DenseTensor* grid_y,
|
|
DenseTensor* grid_x_scale,
|
|
DenseTensor* grid_y_scale,
|
|
DenseTensor* input_grad,
|
|
DenseTensor* grid_grad) {
|
|
const int n = static_cast<int>(grid_x->dims()[0]);
|
|
const int out_h = static_cast<int>(grid_x->dims()[1]);
|
|
const int out_w = static_cast<int>(grid_x->dims()[2]);
|
|
const int c = static_cast<int>(input.dims()[1]);
|
|
|
|
DenseTensor x_w, x_e, y_n, y_s;
|
|
DenseTensor d_w, d_e, d_n, d_s;
|
|
DenseTensor v_wn, v_en, v_ws, v_es;
|
|
|
|
AllNeighbors<T>(dev_ctx,
|
|
input,
|
|
grid_x, // grid_x
|
|
grid_y, // grid_y
|
|
&x_w,
|
|
&x_e,
|
|
&y_n,
|
|
&y_s,
|
|
&d_w,
|
|
&d_e,
|
|
&d_n,
|
|
&d_s,
|
|
&v_wn,
|
|
&v_en,
|
|
&v_ws,
|
|
&v_es);
|
|
|
|
// gather output grad value to input grad by corner point coords and weight
|
|
GatherOutputGradToInputGrad<T>(output_grad, input_grad, x_w, y_n, d_e, d_s);
|
|
GatherOutputGradToInputGrad<T>(output_grad, input_grad, x_w, y_s, d_e, d_n);
|
|
GatherOutputGradToInputGrad<T>(output_grad, input_grad, x_e, y_n, d_w, d_s);
|
|
GatherOutputGradToInputGrad<T>(output_grad, input_grad, x_e, y_s, d_w, d_n);
|
|
|
|
auto v_wn_t = EigenTensor<T, 4>::From(v_wn);
|
|
auto v_en_t = EigenTensor<T, 4>::From(v_en);
|
|
auto v_ws_t = EigenTensor<T, 4>::From(v_ws);
|
|
auto v_es_t = EigenTensor<T, 4>::From(v_es);
|
|
|
|
auto d_w_t = EigenTensor<T, 3>::From(d_w);
|
|
auto d_e_t = EigenTensor<T, 3>::From(d_e);
|
|
auto d_n_t = EigenTensor<T, 3>::From(d_n);
|
|
auto d_s_t = EigenTensor<T, 3>::From(d_s);
|
|
|
|
auto output_grad_t = EigenTensor<T, 4>::From(output_grad);
|
|
|
|
if (grid_grad != nullptr) {
|
|
DenseTensor grid_grad_x, grid_grad_y;
|
|
grid_grad_x.Resize({n, out_h, out_w});
|
|
grid_grad_y.Resize({n, out_h, out_w});
|
|
dev_ctx.Alloc<T>(&grid_grad_x);
|
|
dev_ctx.Alloc<T>(&grid_grad_y);
|
|
auto grid_grad_x_t =
|
|
EigenTensor<T, 3>::From(grid_grad_x).setConstant(static_cast<T>(0.0));
|
|
auto grid_grad_y_t =
|
|
EigenTensor<T, 3>::From(grid_grad_y).setConstant(static_cast<T>(0.0));
|
|
for (int i = 0; i < n; i++) {
|
|
for (int j = 0; j < c; j++) {
|
|
for (int k = 0; k < out_h; k++) {
|
|
for (int l = 0; l < out_w; l++) {
|
|
grid_grad_x_t(i, k, l) +=
|
|
((v_en_t(i, j, k, l) - v_wn_t(i, j, k, l)) * d_s_t(i, k, l) +
|
|
(v_es_t(i, j, k, l) - v_ws_t(i, j, k, l)) * d_n_t(i, k, l)) *
|
|
output_grad_t(i, j, k, l);
|
|
grid_grad_y_t(i, k, l) +=
|
|
((v_ws_t(i, j, k, l) - v_wn_t(i, j, k, l)) * d_e_t(i, k, l) +
|
|
(v_es_t(i, j, k, l) - v_en_t(i, j, k, l)) * d_w_t(i, k, l)) *
|
|
output_grad_t(i, j, k, l);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// const T x_max = static_cast<T>(in_w - 1);
|
|
// const T y_max = static_cast<T>(in_h - 1);
|
|
|
|
auto grid_x_scale_t = EigenTensor<T, 3>::From(*grid_x_scale);
|
|
auto grid_y_scale_t = EigenTensor<T, 3>::From(*grid_y_scale);
|
|
grid_grad_x_t = grid_grad_x_t * grid_x_scale_t;
|
|
grid_grad_y_t = grid_grad_y_t * grid_y_scale_t;
|
|
|
|
// gather grid_grad [x, y] in 3rd Dim
|
|
T* grid_grad_data = grid_grad->data<T>();
|
|
T* grid_grad_x_data = grid_grad_x.data<T>();
|
|
T* grid_grad_y_data = grid_grad_y.data<T>();
|
|
for (int64_t i = 0; i < static_cast<int64_t>(n) * out_h * out_w; i++) {
|
|
grid_grad_data[2 * i] = grid_grad_x_data[i];
|
|
grid_grad_data[2 * i + 1] = grid_grad_y_data[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void Gather3DBilinearGrad(const CPUContext& dev_ctx,
|
|
const DenseTensor& input,
|
|
const DenseTensor& output_grad,
|
|
DenseTensor* grid_x,
|
|
DenseTensor* grid_y,
|
|
DenseTensor* grid_z,
|
|
DenseTensor* grid_x_scale,
|
|
DenseTensor* grid_y_scale,
|
|
DenseTensor* grid_z_scale,
|
|
DenseTensor* input_grad,
|
|
DenseTensor* grid_grad) {
|
|
const int n = static_cast<int>(grid_x->dims()[0]);
|
|
const int out_d = static_cast<int>(grid_x->dims()[1]);
|
|
const int out_h = static_cast<int>(grid_x->dims()[2]);
|
|
const int out_w = static_cast<int>(grid_x->dims()[3]);
|
|
const int c = static_cast<int>(input.dims()[1]);
|
|
|
|
DenseTensor x_w, x_e, y_n, y_s, z_t, z_b;
|
|
DenseTensor d_w, d_e, d_n, d_s, d_t, d_b;
|
|
DenseTensor v_twn, v_ten, v_tws, v_tes, v_bwn, v_ben, v_bws, v_bes;
|
|
|
|
All3DNeighbors<T>(dev_ctx,
|
|
input,
|
|
grid_x,
|
|
grid_y,
|
|
grid_z,
|
|
&x_w,
|
|
&x_e,
|
|
&y_n,
|
|
&y_s,
|
|
&z_t,
|
|
&z_b,
|
|
&d_w,
|
|
&d_e,
|
|
&d_n,
|
|
&d_s,
|
|
&d_t,
|
|
&d_b,
|
|
&v_twn,
|
|
&v_ten,
|
|
&v_tws,
|
|
&v_tes,
|
|
&v_bwn,
|
|
&v_ben,
|
|
&v_bws,
|
|
&v_bes);
|
|
// gather output grad value to input grad by corner point coords and weight
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_w, y_n, z_t, d_e, d_s, d_b);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_w, y_s, z_t, d_e, d_n, d_b);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_e, y_n, z_t, d_w, d_s, d_b);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_e, y_s, z_t, d_w, d_n, d_b);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_w, y_n, z_b, d_e, d_s, d_t);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_w, y_s, z_b, d_e, d_n, d_t);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_e, y_n, z_b, d_w, d_s, d_t);
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
output_grad, input_grad, x_e, y_s, z_b, d_w, d_n, d_t);
|
|
|
|
auto v_twn_t = EigenTensor<T, 5>::From(v_twn);
|
|
auto v_ten_t = EigenTensor<T, 5>::From(v_ten);
|
|
auto v_tws_t = EigenTensor<T, 5>::From(v_tws);
|
|
auto v_tes_t = EigenTensor<T, 5>::From(v_tes);
|
|
auto v_bwn_t = EigenTensor<T, 5>::From(v_bwn);
|
|
auto v_ben_t = EigenTensor<T, 5>::From(v_ben);
|
|
auto v_bws_t = EigenTensor<T, 5>::From(v_bws);
|
|
auto v_bes_t = EigenTensor<T, 5>::From(v_bes);
|
|
|
|
auto d_w_t = EigenTensor<T, 4>::From(d_w);
|
|
auto d_e_t = EigenTensor<T, 4>::From(d_e);
|
|
auto d_n_t = EigenTensor<T, 4>::From(d_n);
|
|
auto d_s_t = EigenTensor<T, 4>::From(d_s);
|
|
auto d_t_t = EigenTensor<T, 4>::From(d_t);
|
|
auto d_b_t = EigenTensor<T, 4>::From(d_b);
|
|
|
|
auto output_grad_t = EigenTensor<T, 5>::From(output_grad);
|
|
if (grid_grad != nullptr) {
|
|
DenseTensor grid_grad_x, grid_grad_y, grid_grad_z;
|
|
grid_grad_x.Resize({n, out_d, out_h, out_w});
|
|
grid_grad_y.Resize({n, out_d, out_h, out_w});
|
|
grid_grad_z.Resize({n, out_d, out_h, out_w});
|
|
dev_ctx.Alloc<T>(&grid_grad_x);
|
|
dev_ctx.Alloc<T>(&grid_grad_y);
|
|
dev_ctx.Alloc<T>(&grid_grad_z);
|
|
auto grid_grad_x_t =
|
|
EigenTensor<T, 4>::From(grid_grad_x).setConstant(static_cast<T>(0.0));
|
|
auto grid_grad_y_t =
|
|
EigenTensor<T, 4>::From(grid_grad_y).setConstant(static_cast<T>(0.0));
|
|
auto grid_grad_z_t =
|
|
EigenTensor<T, 4>::From(grid_grad_z).setConstant(static_cast<T>(0.0));
|
|
for (int i = 0; i < n; i++) {
|
|
for (int j = 0; j < c; j++) {
|
|
for (int m = 0; m < out_d; m++) {
|
|
for (int k = 0; k < out_h; k++) {
|
|
for (int l = 0; l < out_w; l++) {
|
|
grid_grad_x_t(i, m, k, l) +=
|
|
((v_ten_t(i, j, m, k, l) - v_twn_t(i, j, m, k, l)) *
|
|
d_s_t(i, m, k, l) * d_b_t(i, m, k, l) +
|
|
(v_tes_t(i, j, m, k, l) - v_tws_t(i, j, m, k, l)) *
|
|
d_n_t(i, m, k, l) * d_b_t(i, m, k, l) +
|
|
(v_ben_t(i, j, m, k, l) - v_bwn_t(i, j, m, k, l)) *
|
|
d_s_t(i, m, k, l) * d_t_t(i, m, k, l) +
|
|
(v_bes_t(i, j, m, k, l) - v_bws_t(i, j, m, k, l)) *
|
|
d_n_t(i, m, k, l) * d_t_t(i, m, k, l)) *
|
|
output_grad_t(i, j, m, k, l);
|
|
grid_grad_y_t(i, m, k, l) +=
|
|
((v_tws_t(i, j, m, k, l) - v_twn_t(i, j, m, k, l)) *
|
|
d_e_t(i, m, k, l) * d_b_t(i, m, k, l) +
|
|
(v_tes_t(i, j, m, k, l) - v_ten_t(i, j, m, k, l)) *
|
|
d_w_t(i, m, k, l) * d_b_t(i, m, k, l) +
|
|
(v_bws_t(i, j, m, k, l) - v_bwn_t(i, j, m, k, l)) *
|
|
d_e_t(i, m, k, l) * d_t_t(i, m, k, l) +
|
|
(v_bes_t(i, j, m, k, l) - v_ben_t(i, j, m, k, l)) *
|
|
d_w_t(i, m, k, l) * d_t_t(i, m, k, l)) *
|
|
output_grad_t(i, j, m, k, l);
|
|
grid_grad_z_t(i, m, k, l) +=
|
|
((v_bws_t(i, j, m, k, l) - v_tws_t(i, j, m, k, l)) *
|
|
d_e_t(i, m, k, l) * d_n_t(i, m, k, l) +
|
|
(v_bes_t(i, j, m, k, l) - v_tes_t(i, j, m, k, l)) *
|
|
d_w_t(i, m, k, l) * d_n_t(i, m, k, l) +
|
|
(v_bwn_t(i, j, m, k, l) - v_twn_t(i, j, m, k, l)) *
|
|
d_e_t(i, m, k, l) * d_s_t(i, m, k, l) +
|
|
(v_ben_t(i, j, m, k, l) - v_ten_t(i, j, m, k, l)) *
|
|
d_w_t(i, m, k, l) * d_s_t(i, m, k, l)) *
|
|
output_grad_t(i, j, m, k, l);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
auto grid_x_scale_t = EigenTensor<T, 4>::From(*grid_x_scale);
|
|
auto grid_y_scale_t = EigenTensor<T, 4>::From(*grid_y_scale);
|
|
auto grid_z_scale_t = EigenTensor<T, 4>::From(*grid_z_scale);
|
|
|
|
grid_grad_x_t = grid_grad_x_t * grid_x_scale_t;
|
|
grid_grad_y_t = grid_grad_y_t * grid_y_scale_t;
|
|
grid_grad_z_t = grid_grad_z_t * grid_z_scale_t;
|
|
// gather grid_grad [x, y, z] in 4th Dim
|
|
T* grid_grad_data = grid_grad->data<T>();
|
|
T* grid_grad_x_data = grid_grad_x.data<T>();
|
|
T* grid_grad_y_data = grid_grad_y.data<T>();
|
|
T* grid_grad_z_data = grid_grad_z.data<T>();
|
|
for (int64_t i = 0; i < static_cast<int64_t>(n) * out_d * out_h * out_w;
|
|
i++) {
|
|
grid_grad_data[3 * i] = grid_grad_x_data[i];
|
|
grid_grad_data[3 * i + 1] = grid_grad_y_data[i];
|
|
grid_grad_data[3 * i + 2] = grid_grad_z_data[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void GatherOutputGradToInputGrad(const DenseTensor& output_grad,
|
|
DenseTensor* input_grad,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y) {
|
|
const int n = static_cast<int>(output_grad.dims()[0]);
|
|
const int c = static_cast<int>(output_grad.dims()[1]);
|
|
const int out_h = static_cast<int>(output_grad.dims()[2]);
|
|
const int out_w = static_cast<int>(output_grad.dims()[3]);
|
|
const int in_h = static_cast<int>(input_grad->dims()[2]);
|
|
const int in_w = static_cast<int>(input_grad->dims()[3]);
|
|
auto x_t = EigenTensor<T, 3>::From(x);
|
|
auto y_t = EigenTensor<T, 3>::From(y);
|
|
auto input_grad_t = EigenTensor<T, 4>::From(*input_grad);
|
|
auto output_grad_t = EigenTensor<T, 4>::From(output_grad);
|
|
for (int i = 0; i < n; i++) {
|
|
for (int k = 0; k < out_h; k++) {
|
|
for (int l = 0; l < out_w; l++) {
|
|
if (IsInBound<int>(static_cast<int>(std::nearbyint(x_t(i, k, l))),
|
|
static_cast<int>(std::nearbyint(y_t(i, k, l))),
|
|
(in_w - 1),
|
|
(in_h - 1))) {
|
|
for (int j = 0; j < c; j++) {
|
|
input_grad_t(i,
|
|
j,
|
|
static_cast<int>(std::nearbyint(y_t(i, k, l))),
|
|
static_cast<int>(std::nearbyint(x_t(i, k, l)))) +=
|
|
output_grad_t(i, j, k, l);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
static void Gather3DOutputGradToInputGrad(const DenseTensor& output_grad,
|
|
DenseTensor* input_grad,
|
|
const DenseTensor& x,
|
|
const DenseTensor& y,
|
|
const DenseTensor& z) {
|
|
const int n = static_cast<int>(output_grad.dims()[0]);
|
|
const int c = static_cast<int>(output_grad.dims()[1]);
|
|
const int out_d = static_cast<int>(output_grad.dims()[2]);
|
|
const int out_h = static_cast<int>(output_grad.dims()[3]);
|
|
const int out_w = static_cast<int>(output_grad.dims()[4]);
|
|
const int in_d = static_cast<int>(input_grad->dims()[2]);
|
|
const int in_h = static_cast<int>(input_grad->dims()[3]);
|
|
const int in_w = static_cast<int>(input_grad->dims()[4]);
|
|
auto x_t = EigenTensor<T, 4>::From(x);
|
|
auto y_t = EigenTensor<T, 4>::From(y);
|
|
auto z_t = EigenTensor<T, 4>::From(z);
|
|
auto input_grad_t = EigenTensor<T, 5>::From(*input_grad);
|
|
auto output_grad_t = EigenTensor<T, 5>::From(output_grad);
|
|
for (int i = 0; i < n; i++) {
|
|
for (int m = 0; m < out_d; m++) {
|
|
for (int k = 0; k < out_h; k++) {
|
|
for (int l = 0; l < out_w; l++) {
|
|
if (IsInBound3D<int>(
|
|
static_cast<int>(std::nearbyint(x_t(i, m, k, l))),
|
|
static_cast<int>(std::nearbyint(y_t(i, m, k, l))),
|
|
static_cast<int>(std::nearbyint(z_t(i, m, k, l))),
|
|
(in_w - 1),
|
|
(in_h - 1),
|
|
(in_d - 1))) {
|
|
for (int j = 0; j < c; j++) {
|
|
input_grad_t(i,
|
|
j,
|
|
static_cast<int>(std::nearbyint(z_t(i, m, k, l))),
|
|
static_cast<int>(std::nearbyint(y_t(i, m, k, l))),
|
|
static_cast<int>(std::nearbyint(x_t(i, m, k, l)))) +=
|
|
output_grad_t(i, j, m, k, l);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void GridSampleGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& grid,
|
|
const DenseTensor& out_grad,
|
|
const std::string& mode,
|
|
const std::string& padding_mode,
|
|
bool align_corners,
|
|
DenseTensor* x_grad,
|
|
DenseTensor* grid_grad) {
|
|
if (out_grad.numel() == 0) {
|
|
if (x_grad) {
|
|
Full<T, Context>(dev_ctx, x_grad->dims(), 0, x_grad);
|
|
}
|
|
if (grid_grad) {
|
|
Full<T, Context>(dev_ctx, grid_grad->dims(), 0, grid_grad);
|
|
}
|
|
return;
|
|
}
|
|
|
|
std::string enum_mode;
|
|
if (mode == "nearest") {
|
|
enum_mode = "nearest";
|
|
} else {
|
|
enum_mode = "bilinear";
|
|
}
|
|
|
|
if (x.dims().size() == 4) {
|
|
const int n = static_cast<int>(grid.dims()[0]);
|
|
const int out_h = static_cast<int>(grid.dims()[1]);
|
|
const int out_w = static_cast<int>(grid.dims()[2]);
|
|
const int c = static_cast<int>(x.dims()[1]);
|
|
const int in_h = static_cast<int>(x.dims()[2]);
|
|
const int in_w = static_cast<int>(x.dims()[3]);
|
|
|
|
x_grad->Resize({n, c, in_h, in_w});
|
|
dev_ctx.template Alloc<T>(x_grad);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, x_grad, static_cast<T>(0));
|
|
|
|
if (grid_grad != nullptr) {
|
|
grid_grad->Resize({n, out_h, out_w, 2});
|
|
dev_ctx.template Alloc<T>(grid_grad);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, grid_grad, static_cast<T>(0));
|
|
}
|
|
|
|
DenseTensor grid_x, grid_y;
|
|
DenseTensor grid_x_scale, grid_y_scale;
|
|
CalcGridLocationsWithGrad<T>(dev_ctx,
|
|
grid,
|
|
in_h,
|
|
in_w,
|
|
align_corners,
|
|
padding_mode,
|
|
&grid_x,
|
|
&grid_y,
|
|
&grid_x_scale,
|
|
&grid_y_scale);
|
|
if (enum_mode == "nearest") {
|
|
GatherOutputGradToInputGrad<T>(out_grad, x_grad, grid_x, grid_y);
|
|
|
|
} else if (enum_mode == "bilinear") {
|
|
GatherBilinearGrad<T>(dev_ctx,
|
|
x,
|
|
out_grad,
|
|
&grid_x,
|
|
&grid_y,
|
|
&grid_x_scale,
|
|
&grid_y_scale,
|
|
x_grad,
|
|
grid_grad);
|
|
}
|
|
} else {
|
|
const int n = static_cast<int>(grid.dims()[0]);
|
|
const int out_d = static_cast<int>(grid.dims()[1]);
|
|
const int out_h = static_cast<int>(grid.dims()[2]);
|
|
const int out_w = static_cast<int>(grid.dims()[3]);
|
|
const int c = static_cast<int>(x.dims()[1]);
|
|
const int in_d = static_cast<int>(x.dims()[2]);
|
|
const int in_h = static_cast<int>(x.dims()[3]);
|
|
const int in_w = static_cast<int>(x.dims()[4]);
|
|
|
|
x_grad->Resize({n, c, in_d, in_h, in_w});
|
|
dev_ctx.template Alloc<T>(x_grad);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, x_grad, static_cast<T>(0));
|
|
|
|
if (grid_grad != nullptr) {
|
|
grid_grad->Resize({n, out_d, out_h, out_w, 3});
|
|
dev_ctx.template Alloc<T>(grid_grad);
|
|
funcs::SetConstant<Context, T>()(dev_ctx, grid_grad, static_cast<T>(0));
|
|
}
|
|
DenseTensor grid_x, grid_y, grid_z;
|
|
DenseTensor grid_x_scale, grid_y_scale, grid_z_scale;
|
|
|
|
Calc3DGridLocationsWithGrad<T>(dev_ctx,
|
|
grid,
|
|
in_d,
|
|
in_h,
|
|
in_w,
|
|
align_corners,
|
|
padding_mode,
|
|
&grid_x,
|
|
&grid_y,
|
|
&grid_z,
|
|
&grid_x_scale,
|
|
&grid_y_scale,
|
|
&grid_z_scale);
|
|
if (enum_mode == "nearest") {
|
|
Gather3DOutputGradToInputGrad<T>(
|
|
out_grad, x_grad, grid_x, grid_y, grid_z);
|
|
|
|
} else if (enum_mode == "bilinear") {
|
|
Gather3DBilinearGrad<T>(dev_ctx,
|
|
x,
|
|
out_grad,
|
|
&grid_x,
|
|
&grid_y,
|
|
&grid_z,
|
|
&grid_x_scale,
|
|
&grid_y_scale,
|
|
&grid_z_scale,
|
|
x_grad,
|
|
grid_grad);
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(grid_sample_grad,
|
|
CPU,
|
|
ALL_LAYOUT,
|
|
phi::GridSampleGradKernel,
|
|
float,
|
|
double) {}
|