|
|
|
@ -20,39 +20,8 @@
|
|
|
|
|
namespace paddle {
|
|
|
|
|
namespace operators {
|
|
|
|
|
|
|
|
|
|
template <typename Place, typename T>
|
|
|
|
|
void NaiveCpuTranspose(const framework::ExecutionContext& context,
|
|
|
|
|
const framework::Tensor& in, framework::Tensor& out,
|
|
|
|
|
std::vector<int> axis) {
|
|
|
|
|
auto in_data = in.data<T>();
|
|
|
|
|
auto out_data = out.mutable_data<T>(context.GetPlace());
|
|
|
|
|
auto in_dim = in.dims();
|
|
|
|
|
auto out_dim = out.dims();
|
|
|
|
|
size_t ndims = in_dim.size();
|
|
|
|
|
|
|
|
|
|
std::vector<int> in_offset(ndims, 1);
|
|
|
|
|
std::vector<int> out_offset(ndims, 1);
|
|
|
|
|
|
|
|
|
|
for (int i = ndims - 2; i >= 0; i--) {
|
|
|
|
|
in_offset[i] = in_offset[i + 1] * in_dim[i + 1];
|
|
|
|
|
out_offset[i] = out_offset[i + 1] * out_dim[i + 1];
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
size_t data_size = product(in_dim);
|
|
|
|
|
|
|
|
|
|
for (size_t to_index = 0; to_index < data_size; to_index++) {
|
|
|
|
|
int from_index = 0;
|
|
|
|
|
int temp = to_index;
|
|
|
|
|
for (size_t i = 0; i < ndims; i++) {
|
|
|
|
|
from_index += (temp / out_offset[i]) * in_offset[axis[i]];
|
|
|
|
|
temp = temp % out_offset[i];
|
|
|
|
|
}
|
|
|
|
|
out_data[to_index] = in_data[from_index];
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
template <typename Place, typename T, int Dims>
|
|
|
|
|
void DoTranspose(const framework::ExecutionContext& context,
|
|
|
|
|
void EigenTranspose(const framework::ExecutionContext& context,
|
|
|
|
|
const framework::Tensor& in, framework::Tensor& out,
|
|
|
|
|
std::vector<int> axis) {
|
|
|
|
|
Eigen::array<int, Dims> permute;
|
|
|
|
@ -72,28 +41,32 @@ template <typename Place, typename T>
|
|
|
|
|
class TransposeKernel : public framework::OpKernel {
|
|
|
|
|
public:
|
|
|
|
|
void Compute(const framework::ExecutionContext& context) const override {
|
|
|
|
|
auto* in = context.Input<framework::Tensor>("X");
|
|
|
|
|
auto* out = context.Output<framework::Tensor>("Out");
|
|
|
|
|
out->mutable_data<T>(context.GetPlace());
|
|
|
|
|
auto* input = context.Input<framework::Tensor>("Input");
|
|
|
|
|
auto* output = context.Output<framework::Tensor>("Output");
|
|
|
|
|
output->mutable_data<T>(context.GetPlace());
|
|
|
|
|
|
|
|
|
|
auto axis = context.Attr<std::vector<int>>("axis");
|
|
|
|
|
int ndims = axis.size();
|
|
|
|
|
switch (ndims) {
|
|
|
|
|
case 1:
|
|
|
|
|
break;
|
|
|
|
|
case 2:
|
|
|
|
|
DoTranspose<Place, T, 2>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 2>(context, *input, *output, axis);
|
|
|
|
|
break;
|
|
|
|
|
case 3:
|
|
|
|
|
DoTranspose<Place, T, 3>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 3>(context, *input, *output, axis);
|
|
|
|
|
break;
|
|
|
|
|
case 4:
|
|
|
|
|
DoTranspose<Place, T, 4>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 4>(context, *input, *output, axis);
|
|
|
|
|
break;
|
|
|
|
|
case 5:
|
|
|
|
|
DoTranspose<Place, T, 5>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 5>(context, *input, *output, axis);
|
|
|
|
|
break;
|
|
|
|
|
default:
|
|
|
|
|
NaiveCpuTranspose<Place, T>(context, *in, *out, axis);
|
|
|
|
|
case 6:
|
|
|
|
|
EigenTranspose<Place, T, 6>(context, *input, *output, axis);
|
|
|
|
|
break;
|
|
|
|
|
default:
|
|
|
|
|
PADDLE_THROW("Tensors with rank at most 6 are supported");
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
@ -102,9 +75,11 @@ template <typename Place, typename T>
|
|
|
|
|
class TransposeGradKernel : public framework::OpKernel {
|
|
|
|
|
public:
|
|
|
|
|
void Compute(const framework::ExecutionContext& context) const override {
|
|
|
|
|
auto* in = context.Input<framework::Tensor>(framework::GradVarName("Out"));
|
|
|
|
|
auto* out = context.Output<framework::Tensor>(framework::GradVarName("X"));
|
|
|
|
|
out->mutable_data<T>(context.GetPlace());
|
|
|
|
|
auto* output_grad =
|
|
|
|
|
context.Input<framework::Tensor>(framework::GradVarName("Output"));
|
|
|
|
|
auto* input_grad =
|
|
|
|
|
context.Output<framework::Tensor>(framework::GradVarName("Input"));
|
|
|
|
|
input_grad->mutable_data<T>(context.GetPlace());
|
|
|
|
|
|
|
|
|
|
auto axis_temp = context.Attr<std::vector<int>>("axis");
|
|
|
|
|
std::vector<int> axis(axis_temp);
|
|
|
|
@ -116,21 +91,25 @@ class TransposeGradKernel : public framework::OpKernel {
|
|
|
|
|
int ndims = axis.size();
|
|
|
|
|
|
|
|
|
|
switch (ndims) {
|
|
|
|
|
case 1:
|
|
|
|
|
break;
|
|
|
|
|
case 2:
|
|
|
|
|
DoTranspose<Place, T, 2>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 2>(context, *output_grad, *input_grad, axis);
|
|
|
|
|
break;
|
|
|
|
|
case 3:
|
|
|
|
|
DoTranspose<Place, T, 3>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 3>(context, *output_grad, *input_grad, axis);
|
|
|
|
|
break;
|
|
|
|
|
case 4:
|
|
|
|
|
DoTranspose<Place, T, 4>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 4>(context, *output_grad, *input_grad, axis);
|
|
|
|
|
break;
|
|
|
|
|
case 5:
|
|
|
|
|
DoTranspose<Place, T, 5>(context, *in, *out, axis);
|
|
|
|
|
EigenTranspose<Place, T, 5>(context, *output_grad, *input_grad, axis);
|
|
|
|
|
break;
|
|
|
|
|
default:
|
|
|
|
|
NaiveCpuTranspose<Place, T>(context, *in, *out, axis);
|
|
|
|
|
case 6:
|
|
|
|
|
EigenTranspose<Place, T, 6>(context, *output_grad, *input_grad, axis);
|
|
|
|
|
break;
|
|
|
|
|
default:
|
|
|
|
|
PADDLE_THROW("Tensors with rank at most 6 are supported");
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
};
|
|
|
|
|