388 lines
15 KiB
C++
388 lines
15 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/common/ddim.h"
|
|
#include "paddle/common/layout.h"
|
|
#include "paddle/phi/kernels/conv_transpose_grad_kernel.h"
|
|
#include "paddle/phi/kernels/cpu/conv_util.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/concat_and_split_functor.h"
|
|
#include "paddle/phi/kernels/funcs/im2col.h"
|
|
#include "paddle/phi/kernels/funcs/slice.h"
|
|
#include "paddle/phi/kernels/funcs/vol2col.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T, typename Context>
|
|
void ConvTransposeGradRawKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& filter,
|
|
const DenseTensor& dout,
|
|
const std::vector<int>& strides,
|
|
const std::vector<int>& paddings,
|
|
const std::string& padding_algorithm,
|
|
int groups,
|
|
const std::vector<int>& dilations,
|
|
const std::string& data_format,
|
|
DenseTensor* dx,
|
|
DenseTensor* dfilter) {
|
|
const DataLayout data_layout = StringToDataLayout(data_format);
|
|
// For filter, we do not use const pointer because we will do reshape,
|
|
// but we should avoid modifying its value.
|
|
DenseTensor filter_ = filter;
|
|
|
|
if ((!dx) && (!dfilter)) {
|
|
return;
|
|
}
|
|
|
|
// 0-size
|
|
if (x.numel() == 0) {
|
|
if (dx) dev_ctx.template Alloc<T>(dx);
|
|
if (dfilter) {
|
|
Full<T, Context>(dev_ctx, dfilter->dims(), 0, dfilter);
|
|
}
|
|
return;
|
|
}
|
|
if (filter.numel() == 0) {
|
|
if (dfilter) dev_ctx.template Alloc<T>(dfilter);
|
|
if (dx) {
|
|
Full<T, Context>(dev_ctx, dx->dims(), 0, dx);
|
|
}
|
|
return;
|
|
}
|
|
|
|
std::vector<int> paddings_ = paddings;
|
|
std::vector<int> dilations_ = dilations;
|
|
|
|
auto x_dims = x.dims();
|
|
auto filter_dims = filter_.dims();
|
|
auto dout_dims = dout.dims();
|
|
const int batch_size = static_cast<int>(x.dims()[0]);
|
|
|
|
DDim in_data_dims;
|
|
if (data_layout != DataLayout::NHWC) {
|
|
in_data_dims = slice_ddim(x_dims, 2, x_dims.size());
|
|
} else {
|
|
in_data_dims = slice_ddim(x_dims, 1, x_dims.size() - 1);
|
|
}
|
|
DDim filter_data_dims = slice_ddim(filter_dims, 2, filter_dims.size());
|
|
std::vector<int> ksize = vectorize<int>(filter_data_dims);
|
|
UpdatePaddingAndDilation(
|
|
&paddings_, &dilations_, padding_algorithm, in_data_dims, strides, ksize);
|
|
|
|
// x_shape_vec: {n, c, h, w} or {n, c, d, h, w} for channel_first
|
|
// x_shape_vec: {n, h, w, c} or {n, d, h, w, c} for channel_last
|
|
std::vector<int64_t> x_shape_vec = vectorize(x.dims());
|
|
// filter_shape_vec: {i_c, o_c, k_h, k_w} or {i_c, o_c, k_d, k_h, k_w}
|
|
std::vector<int64_t> filter_shape_vec = vectorize(filter_.dims());
|
|
|
|
// use col_shape in the im2col and col2im (or vol2col and col2vol)
|
|
// calculation
|
|
// col_shape_vec: {o_c, k_h, k_w, h, w} or {o_c, k_d, k_h, k_w, d, h, w} for
|
|
size_t data_dim = filter_shape_vec.size() - 2;
|
|
std::vector<int64_t> col_shape_vec(1 + 2 * data_dim);
|
|
if (data_layout != DataLayout::NHWC) {
|
|
col_shape_vec[0] = dout_dims[1];
|
|
for (size_t j = 0; j < data_dim; ++j) {
|
|
col_shape_vec[j + 1] = filter_shape_vec[j + 2];
|
|
col_shape_vec[j + 1 + data_dim] = x_shape_vec[j + 2];
|
|
}
|
|
} else {
|
|
col_shape_vec[0] = dout_dims[dout_dims.size() - 1];
|
|
for (size_t j = 0; j < data_dim; ++j) {
|
|
col_shape_vec[j + 1] = filter_shape_vec[j + 2];
|
|
col_shape_vec[j + 1 + data_dim] = x_shape_vec[j + 1];
|
|
}
|
|
}
|
|
DDim col_shape(make_ddim(col_shape_vec));
|
|
|
|
// use col_matrix_shape in the gemm calculation
|
|
// size: (o_c * k_h * k_w, h * w) or (o_c * k_d * k_h * k_w, d * h * w)
|
|
DDim col_matrix_shape = flatten_to_2d(col_shape, data_dim + 1);
|
|
|
|
// output size: (o_c, o_h, o_w) or (o_c, o_d, o_h, o_w) for channel_first
|
|
// output size: (o_h, o_w, o_c) or (o_d, o_h, o_w, o_c) for channel_last
|
|
DDim output_shape = slice_ddim(dout.dims(), 1, dout.dims().size());
|
|
|
|
// x matrix size: (i_c, h * w) or (i_c, d * h * w) for channel_first
|
|
// x matrix size: (h * w, i_c) or (d * h * w, i_c) for channel_last
|
|
DDim x_matrix_shape;
|
|
if (data_layout != DataLayout::NHWC) {
|
|
x_matrix_shape = {x_dims[1], col_matrix_shape[1]};
|
|
} else {
|
|
x_matrix_shape = {col_matrix_shape[1], x_dims[x_dims.size() - 1]};
|
|
}
|
|
|
|
// filter size: (i_c, o_c/g * k_h * k_w) or (i_c, o_c/g * k_d * k_h * k_w)
|
|
DDim filter_matrix_shape;
|
|
if (data_layout != DataLayout::NHWC) {
|
|
filter_matrix_shape = {x_dims[1], col_matrix_shape[0] / groups};
|
|
} else {
|
|
filter_matrix_shape = {x_dims[x_dims.size() - 1],
|
|
col_matrix_shape[0] / groups};
|
|
}
|
|
filter_.Resize(filter_matrix_shape);
|
|
|
|
int in_step = (data_layout != DataLayout::NHWC
|
|
? static_cast<int>(x_dims[1]) / groups
|
|
: static_cast<int>(x_dims[x_dims.size() - 1]) / groups);
|
|
int col_step = static_cast<int>(col_matrix_shape[0]) / groups;
|
|
|
|
// convolution transpose grad on x:
|
|
// im2col + gemm (similar to conv-forward)
|
|
// x need to compute gradient
|
|
auto blas = funcs::GetBlas<Context, T>(dev_ctx);
|
|
if (dx || dfilter) {
|
|
DenseTensor col;
|
|
col.Resize(col_shape);
|
|
dev_ctx.template Alloc<T>(&col);
|
|
// col_matrix shares the same piece of data with col,
|
|
// but will be reshaped into a two-dimensional matrix shape
|
|
// to call the matrix multiplication interface.
|
|
DenseTensor col_matrix;
|
|
col_matrix.ShareDataWith(col);
|
|
col_matrix.Resize(col_matrix_shape);
|
|
|
|
DenseTensor dfilter_;
|
|
funcs::SetConstant<Context, T> set_zero;
|
|
|
|
funcs::Im2ColFunctor<funcs::ColFormat::CFO, Context, T> im2col;
|
|
funcs::Vol2ColFunctor<Context, T> vol2col;
|
|
funcs::ConcatFunctor<Context, T> concat_functor;
|
|
|
|
if (dx) {
|
|
dev_ctx.template Alloc<T>(dx);
|
|
set_zero(dev_ctx, dx, static_cast<T>(0));
|
|
}
|
|
if (dfilter) { // dfilter_ size (i_c, o_c/g, k_h, k_w)
|
|
dev_ctx.template Alloc<T>(dfilter);
|
|
set_zero(dev_ctx, dfilter, static_cast<T>(0));
|
|
dfilter_ = *dfilter;
|
|
dfilter_.Resize(filter_matrix_shape);
|
|
}
|
|
|
|
size_t D = x.dims().size();
|
|
for (int i = 0; i < batch_size; i++) {
|
|
// batch with size (o_c, o_h, o_w) or (o_c, o_d, o_h, o_w) for
|
|
// channel_first
|
|
// batch with size (o_h, o_w, o_c) or (o_d, o_h, o_w, o_c) for
|
|
// channel_last
|
|
DenseTensor dout_batch = dout.Slice(i, i + 1).Resize(output_shape);
|
|
|
|
if (data_dim == 2U) {
|
|
// im2col: dy -> col matrix
|
|
// from (o_c, o_h, o_w) to (o_c * k_h * k_w, i_h * i_w) for
|
|
// channel_first
|
|
// from (o_h, o_w, o_c) to (o_c * k_h * k_w, i_h * i_w) for
|
|
// channel_last
|
|
im2col(dev_ctx,
|
|
dout_batch,
|
|
dilations_,
|
|
strides,
|
|
std::vector<int>{
|
|
paddings_[0], paddings_[2], paddings_[1], paddings_[3]},
|
|
&col,
|
|
data_layout);
|
|
} else if (data_dim == 3U) {
|
|
// vol2col: dy -> col_matrix
|
|
// from (o_c, o_d, o_h, o_w) to (o_c * k_d * k_h * k_w, i_d * i_h *
|
|
// i_w) for channel_first
|
|
// from (o_d, o_h, o_w, o_c) to (i_d * i_h * i_w, o_c * k_d * k_h *
|
|
// k_w) for channel_last
|
|
vol2col(dev_ctx,
|
|
dout_batch,
|
|
dilations_,
|
|
strides,
|
|
paddings_,
|
|
&col,
|
|
data_layout);
|
|
}
|
|
if (dx) {
|
|
// batch with size (i_c, i_h, i_w) or (i_h, i_w, i_c)
|
|
DenseTensor dx_batch = dx->Slice(i, i + 1).Resize(x_matrix_shape);
|
|
|
|
// gemm: dx = filter * dy
|
|
// (i_c, o_c * k_h * k_w) * (o_c * k_h * k_w, i_h * i_w) -> (i_c, i_h
|
|
// * i_w)
|
|
// or
|
|
// (i_c, o_c * k_d * k_h * k_w) * (o_c * k_d * k_h * k_w, i_d * i_h *
|
|
// i_w) -> (i_c,
|
|
// i_d, i_h, i_w)
|
|
// gemm: dx = dy^T * filter^T for channel_last
|
|
|
|
std::vector<DenseTensor> dx_batch_vec;
|
|
for (int g = 0; g < groups; g++) {
|
|
// dx_slice: (i_c/g, i_h * i_w) or (i_c/g, i_d * i_h * i_w)
|
|
// for channel_first
|
|
// dx_slice: (i_h * i_w, i_c/g) or (i_d * i_h * i_w, i_c/g)
|
|
// for channel_last
|
|
// filter_slice: (i_c/g, o_c/g * k_h * k_w)
|
|
DenseTensor filter_slice =
|
|
filter_.Slice(g * in_step, (g + 1) * in_step);
|
|
// col_matrix_slice: (o_c/g * k_h * k_w, h * w) or (o_c/g * k_d *
|
|
// k_h * k_w, d * h * w)
|
|
DenseTensor col_matrix_slice =
|
|
col_matrix.Slice(g * col_step, (g + 1) * col_step);
|
|
if (data_layout != DataLayout::NHWC) {
|
|
DenseTensor dx_slice =
|
|
dx_batch.Slice(g * in_step, (g + 1) * in_step);
|
|
blas.MatMul(filter_slice,
|
|
false,
|
|
col_matrix_slice,
|
|
false,
|
|
static_cast<T>(1.0),
|
|
&dx_slice,
|
|
static_cast<T>(0.0));
|
|
} else {
|
|
DenseTensor dx_slice;
|
|
funcs::Slice<Context, T, 2>(dev_ctx,
|
|
&dx_batch,
|
|
&dx_slice,
|
|
g * in_step,
|
|
(g + 1) * in_step,
|
|
1);
|
|
blas.MatMul(col_matrix_slice,
|
|
true,
|
|
filter_slice,
|
|
true,
|
|
static_cast<T>(1.0),
|
|
&dx_slice,
|
|
static_cast<T>(0.0));
|
|
DDim dx_slice_shape;
|
|
if (data_dim == 2U) {
|
|
dx_slice_shape = {x_dims[1], x_dims[2], in_step};
|
|
} else {
|
|
dx_slice_shape = {x_dims[1], x_dims[2], x_dims[3], in_step};
|
|
}
|
|
dx_slice = dx_slice.Resize(dx_slice_shape);
|
|
dx_batch_vec.push_back(dx_slice);
|
|
}
|
|
}
|
|
if (data_layout == DataLayout::NHWC) {
|
|
concat_functor(
|
|
dev_ctx, dx_batch_vec, static_cast<int>(D - 2), &dx_batch);
|
|
}
|
|
}
|
|
if (dfilter) {
|
|
// x batch: (i_c, i_h * i_w) or (i_h, i_w * i_c)
|
|
DenseTensor in_batch = x.Slice(i, i + 1).Resize(x_matrix_shape);
|
|
// gemm: d_filter = x * dy^T
|
|
// (i_c, i_h * i_w) * (i_h * i_w, o_c * k_h * k_w) -> (i_c, o_c * k_h
|
|
// * k_w)
|
|
// or
|
|
// (i_c, i_d * i_h * i_w) * (i_d * i_h * i_w, o_c * k_d * k_h * k_w)
|
|
// -> (i_c, o_c * k_d *
|
|
// k_h * k_w)
|
|
// gemm: d_filter = x^T * dy^T for channel_last
|
|
|
|
for (int g = 0; g < groups; g++) {
|
|
DenseTensor dfilter_slice =
|
|
dfilter_.Slice(g * in_step, (g + 1) * in_step);
|
|
DenseTensor col_matrix_slice =
|
|
col_matrix.Slice(g * col_step, (g + 1) * col_step);
|
|
if (data_layout != DataLayout::NHWC) {
|
|
DenseTensor in_batch_slice =
|
|
in_batch.Slice(g * in_step, (g + 1) * in_step);
|
|
blas.MatMul(in_batch_slice,
|
|
false,
|
|
col_matrix_slice,
|
|
true,
|
|
static_cast<T>(1.0),
|
|
&dfilter_slice,
|
|
static_cast<T>(1.0));
|
|
} else {
|
|
DenseTensor in_batch_slice;
|
|
funcs::Slice<Context, T, 2>(dev_ctx,
|
|
&in_batch,
|
|
&in_batch_slice,
|
|
g * in_step,
|
|
(g + 1) * in_step,
|
|
1);
|
|
blas.MatMul(in_batch_slice,
|
|
true,
|
|
col_matrix_slice,
|
|
true,
|
|
static_cast<T>(1.0),
|
|
&dfilter_slice,
|
|
static_cast<T>(1.0));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void Conv2dTransposeGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& filter,
|
|
const DenseTensor& dout,
|
|
const std::vector<int>& strides,
|
|
const std::vector<int>& paddings,
|
|
const std::vector<int>& output_padding UNUSED,
|
|
const IntArray& output_size UNUSED,
|
|
const std::string& padding_algorithm,
|
|
int groups,
|
|
const std::vector<int>& dilations,
|
|
const std::string& data_format,
|
|
DenseTensor* dx,
|
|
DenseTensor* dfilter) {
|
|
ConvTransposeGradRawKernel<T, Context>(dev_ctx,
|
|
x,
|
|
filter,
|
|
dout,
|
|
strides,
|
|
paddings,
|
|
padding_algorithm,
|
|
groups,
|
|
dilations,
|
|
data_format,
|
|
dx,
|
|
dfilter);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void Conv3dTransposeGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& filter,
|
|
const DenseTensor& dout,
|
|
const std::vector<int>& strides,
|
|
const std::vector<int>& paddings,
|
|
const std::vector<int>& output_padding UNUSED,
|
|
const std::vector<int>& output_size UNUSED,
|
|
const std::string& padding_algorithm,
|
|
int groups,
|
|
const std::vector<int>& dilations,
|
|
const std::string& data_format,
|
|
DenseTensor* dx,
|
|
DenseTensor* dfilter) {
|
|
ConvTransposeGradRawKernel<T, Context>(dev_ctx,
|
|
x,
|
|
filter,
|
|
dout,
|
|
strides,
|
|
paddings,
|
|
padding_algorithm,
|
|
groups,
|
|
dilations,
|
|
data_format,
|
|
dx,
|
|
dfilter);
|
|
}
|
|
|
|
} // namespace phi
|