未验证 提交 5161f71a 编写于 作者: J jiangcheng 提交者: GitHub

optimize slice op and slice grad op (#32266)

* optimize slice op and slice grad op, test=develop

* optimize variable name and annotation information, test=develop
上级 ab3d2bf0
......@@ -259,7 +259,20 @@ class SliceKernel : public framework::OpKernel<T> {
auto out_t =
framework::EigenTensor<T, D, Eigen::RowMajor, Eigen::DenseIndex>::From(
*out, new_out_dims);
if (in->numel() <= Eigen::NumTraits<int>::highest()) {
// similar to tf.slice:
// if element number less than INT_MAX, change the type of index to int
Eigen::DSizes<int, D> offsets_32bit, extents_32bit;
for (size_t i = 0; i < D; i++) {
offsets_32bit[i] = offsets[i];
extents_32bit[i] = extents[i];
}
framework::To32BitIndex(out_t).device(place) =
framework::To32BitIndex(in_t).slice(offsets_32bit, extents_32bit);
} else {
out_t.device(place) = in_t.slice(offsets, extents);
}
out->Resize(out_dims);
}
......@@ -300,8 +313,6 @@ class SliceGradKernel : public framework::OpKernel<T> {
private:
template <size_t D>
void SliceCompute(const framework::ExecutionContext& context) const {
auto& place =
*context.template device_context<DeviceContext>().eigen_device();
auto axes = context.Attr<std::vector<int>>("axes");
auto starts_int = context.Attr<std::vector<int>>("starts");
......@@ -435,14 +446,190 @@ class SliceGradKernel : public framework::OpKernel<T> {
paddings[i].first = offsets[i];
paddings[i].second = (in_dims[i] - out_dims[i]) - offsets[i];
}
EigenPaddingCompute(context, d_input, in_dims, d_out, out_dims, paddings);
}
template <size_t D>
void EigenPaddingCompute(
const framework::ExecutionContext& context, framework::Tensor* d_input,
const framework::DDim& in_dims, const framework::Tensor* d_out,
const framework::DDim& out_dims,
const Eigen::array<std::pair<int64_t, int64_t>, D>& paddings) const {
if (D <= 3) {
// if dimension less than 3, cannot reduce dimension
LaunchEigenPadding(context, d_input, in_dims, d_out, out_dims, paddings);
} else { // else we can reduce dimension
// count not-zero padding number, and record the dimension
int need_pad_num = 0, pad_dim = -1;
for (size_t i = 0; i < D; i++) {
if (paddings[i].first != 0 || paddings[i].second != 0) {
need_pad_num++;
pad_dim = i;
}
}
if (need_pad_num == 0) {
// do not need padding, pass if data address same, else copy
if (d_input->mutable_data<T>(context.GetPlace()) == d_out->data<T>()) {
// inplace, do not any operator, pass
} else {
framework::TensorCopy(
*d_out, context.GetPlace(),
context.template device_context<platform::DeviceContext>(),
d_input);
}
} else if (need_pad_num == 1) {
// only need padding one dimension, we can reduce dimension.
// only the padding dimension is available for us.
// How to reduce dimension(5 to 3 for example):
// before(D=5):
// in_dims: [x1, x2, x3, x4, x5]
// padding.first: [0, 0, a, 0, 0]
// padding.second: [0, 0, b, 0, 0]
// | |
// V V
// after(D=3):
// reshaped_in_dims: [x1*x2, x3, x4*x5]
// reshaped_padding.first: [0, a, 0]
// reshaped_padding.second: [0, b, 0]
if (pad_dim == D - 1) {
// only last dimension need padding,
// reshape the dimension of tensor in 2: [preceding, padding]
std::vector<int64_t> in_tore_shape(2, 1), out_tore_shape(2, 1);
Eigen::array<std::pair<int64_t, int64_t>, 2> reshaped_padding;
// first dimension is the accumulate of preceding dimension
for (int i = 0; i < pad_dim; i++) {
in_tore_shape[0] *= in_dims[i];
out_tore_shape[0] *= out_dims[i];
}
// second dimension is the padding dimension
in_tore_shape[1] = in_dims[pad_dim];
out_tore_shape[1] = out_dims[pad_dim];
// convert array from std::vector to DDim
framework::DDim reshaped_in_dims =
framework::make_ddim(in_tore_shape);
framework::DDim reshaped_out_dims =
framework::make_ddim(out_tore_shape);
// after reshape: the first dimension do not need padding,
// set padding[0] zero
reshaped_padding[0].first = reshaped_padding[0].second = 0;
// the second dimension is the previous padding dimension
reshaped_padding[1].first = paddings[pad_dim].first;
reshaped_padding[1].second = paddings[pad_dim].second;
LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out,
reshaped_out_dims, reshaped_padding);
} else if (pad_dim == 0) {
// only first dimension need padding,
// reshape the dimension of tensor in 2: [padding, succeeding]
// similar to (D - 1)
std::vector<int64_t> in_tore_shape(2, 1), out_tore_shape(2, 1);
Eigen::array<std::pair<int64_t, int64_t>, 2> reshaped_padding;
// first dimension is the padding dimension
in_tore_shape[0] = in_dims[pad_dim];
out_tore_shape[0] = out_dims[pad_dim];
// sencond dimension is the accumulate of succeeding dimension
for (size_t i = pad_dim + 1; i < D; i++) {
in_tore_shape[1] *= in_dims[i];
out_tore_shape[1] *= out_dims[i];
}
// convert array from std::vector to DDim
framework::DDim reshaped_in_dims =
framework::make_ddim(in_tore_shape);
framework::DDim reshaped_out_dims =
framework::make_ddim(out_tore_shape);
// after reshape:
// the first dimension is the previous padding dimension
reshaped_padding[0].first = paddings[pad_dim].first;
reshaped_padding[0].second = paddings[pad_dim].second;
// the second dimension do not need padding, set padding[1] zero
reshaped_padding[1].first = reshaped_padding[1].second = 0;
LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out,
reshaped_out_dims, reshaped_padding);
} else {
// other dimension need padding
// reshape the dimension of tensor in 3:
// [preceding, padding, succeeding]
std::vector<int64_t> in_tore_shape(3, 1), out_tore_shape(3, 1);
Eigen::array<std::pair<int64_t, int64_t>, 3> reshaped_padding;
// first dimension is the accumulate of preceding dimension
for (int i = 0; i < pad_dim; i++) {
in_tore_shape[0] *= in_dims[i];
out_tore_shape[0] *= out_dims[i];
}
// second dimension is the padding dimension
in_tore_shape[1] = in_dims[pad_dim];
out_tore_shape[1] = out_dims[pad_dim];
// third dimension is the accumulate of succeeding dimension
for (size_t i = pad_dim + 1; i < D; i++) {
in_tore_shape[2] *= in_dims[i];
out_tore_shape[2] *= out_dims[i];
}
// convert array from std::vector to DDim
framework::DDim reshaped_in_dims =
framework::make_ddim(in_tore_shape);
framework::DDim reshaped_out_dims =
framework::make_ddim(out_tore_shape);
// after reshape:
// the first dimension do not need padding, set padding[0] zero
reshaped_padding[0].first = reshaped_padding[2].second = 0;
// the second dimension is the previous padding dimension
reshaped_padding[1].first = paddings[pad_dim].first;
reshaped_padding[1].second = paddings[pad_dim].second;
// the third dimension do not need padding, set padding[2] zero
reshaped_padding[2].first = reshaped_padding[2].second = 0;
LaunchEigenPadding(context, d_input, reshaped_in_dims, d_out,
reshaped_out_dims, reshaped_padding);
}
} else {
// need padding at many dimension, cannot reduce dimension
LaunchEigenPadding(context, d_input, in_dims, d_out, out_dims,
paddings);
}
}
}
template <size_t D>
void LaunchEigenPadding(
const framework::ExecutionContext& context, framework::Tensor* d_input,
const framework::DDim& in_dims, const framework::Tensor* d_out,
const framework::DDim& out_dims,
const Eigen::array<std::pair<int64_t, int64_t>, D>& paddings) const {
auto& place =
*context.template device_context<DeviceContext>().eigen_device();
auto d_in_t =
framework::EigenTensor<T, D, Eigen::RowMajor, Eigen::DenseIndex>::From(
*d_input);
*d_input, in_dims);
auto d_out_t =
framework::EigenTensor<T, D, Eigen::RowMajor, Eigen::DenseIndex>::From(
*d_out, out_dims);
if (d_input->numel() <= Eigen::NumTraits<int>::highest()) {
// similar to tf.pad:
// if element number less than INT_MAX, change the type of index to int
Eigen::array<std::pair<int, int>, D> paddings_32bit;
for (size_t i = 0; i < D; i++) {
paddings_32bit[i] =
std::make_pair(paddings[i].first, paddings[i].second);
}
framework::To32BitIndex(d_in_t).device(place) =
framework::To32BitIndex(d_out_t).pad(paddings_32bit, T(0));
} else {
d_in_t.device(place) = d_out_t.pad(paddings, T(0));
}
}
};
} // namespace operators
} // namespace paddle
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册