transpose_op.cc 13.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
X
xzl 已提交
2

L
Luo Tao 已提交
3 4 5
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
X
xzl 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
X
xzl 已提交
8

L
Luo Tao 已提交
9 10 11 12 13
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. */
X
xzl 已提交
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/transpose_op.h"
16

17
#include <memory>
18
#include <string>
19
#include <vector>
X
xzl 已提交
20

21 22 23 24
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif

X
xzl 已提交
25 26 27 28 29 30 31
namespace paddle {
namespace operators {

class TransposeOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

32
  void InferShape(framework::InferShapeContext *ctx) const override {
33 34
    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Transpose");
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Transpose");
Q
Qiao Longfei 已提交
35 36
    auto x_dims = ctx->GetInputDim("X");
    std::vector<int> axis = ctx->Attrs().Get<std::vector<int>>("axis");
X
xzl 已提交
37
    size_t x_rank = x_dims.size();
X
xzl 已提交
38
    size_t axis_size = axis.size();
X
xzl 已提交
39

40 41
    PADDLE_ENFORCE_EQ(x_rank,
                      axis_size,
42 43 44 45 46
                      platform::errors::InvalidArgument(
                          "The input tensor's dimension "
                          "should be equal to the axis's size. "
                          "But received input tensor's dimension is %d, "
                          "axis's size is %d",
47 48
                          x_rank,
                          axis_size));
49 50 51

    std::vector<int> count(axis_size, 0);
    for (size_t i = 0; i < axis_size; i++) {
52 53
      PADDLE_ENFORCE_GE(axis[i],
                        0,
54 55 56
                        platform::errors::InvalidArgument(
                            "The axis should be greater than or equal to 0."
                            "But received %d of axis[%d]",
57 58
                            axis[i],
                            i));
59

60
      PADDLE_ENFORCE_EQ(
61 62
          axis[i] < static_cast<int>(axis_size) && ++count[axis[i]] == 1,
          true,
63 64 65 66 67 68 69
          platform::errors::InvalidArgument(
              "Each element of Attribute axis should "
              "be a unique value range from 0 to (dims - 1), "
              "where the dims is the axis's size, "
              "unique value means this axis value can appear only once. "
              "But received axis[%d] is %d, axis_size is %d, "
              "count[axis[%d]] is %d",
70 71 72 73 74
              i,
              axis[i],
              axis_size,
              i,
              count[axis[i]]));
X
xzl 已提交
75
    }
X
xzl 已提交
76

X
xzl 已提交
77
    framework::DDim out_dims(x_dims);
J
Jacek Czaja 已提交
78 79 80
#ifdef PADDLE_WITH_MKLDNN
    // Here we need to match dims to paddle layout
    // as we are producing non-oneDNN result
81
    if (ctx->IsRunMKLDNNKernel() && (x_dims.size() >= 3) &&
J
Jacek Czaja 已提交
82
        (paddle::platform::MKLDNNDeviceContext::tls()
83
             .get_cur_paddle_data_layout() == phi::DataLayout::kNHWC)) {
84
      auto dims = phi::vectorize<int>(x_dims);
J
Jacek Czaja 已提交
85 86 87 88 89 90
      std::rotate(dims.begin() + 1, dims.begin() + 2, dims.end());
      x_dims = x_dims.reshape(dims);
      VLOG(3)
          << "Rotating Shape in Transpose from: kMKLDNN to: kNHWC output_shape";
    }
#endif
91
    for (size_t i = 0; i < axis_size; i++) {
X
xzl 已提交
92
      out_dims[i] = x_dims[axis[i]];
X
xzl 已提交
93
    }
Q
Qiao Longfei 已提交
94
    ctx->SetOutputDim("Out", out_dims);
X
xzl 已提交
95
  }
96 97 98 99

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
100
    auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
J
jiahongyu 已提交
101
    auto &data_format = ctx.Attr<std::string>("data_format");
102
    phi::DataLayout layout_ = phi::StringToDataLayout(data_format);
J
jiahongyu 已提交
103
    return framework::OpKernelType(data_type, ctx.GetPlace(), layout_);
104
  }
X
xzl 已提交
105 106 107 108
};

class TransposeOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
109
  void Make() override {
110
    AddInput(
X
xzl 已提交
111
        "X",
112 113
        "(Tensor) The input tensor, tensors with rank up to 6 are supported.");
    AddOutput("Out", "(Tensor)The output tensor.");
X
xzl 已提交
114 115
    AddAttr<std::vector<int>>(
        "axis",
116 117 118
        "(vector<int>) A list of values, and the size of the list should be "
        "the same with the input tensor rank. This operator permutes the input "
        "tensor's axes according to the values given.");
119 120
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
121 122
        .SetDefault(false)
        .AsExtra();
123 124 125 126 127 128
    AddAttr<std::string>(
        "data_format",
        "(string, default NCHW) Only used in "
        "An optional string from: \"NHWC\", \"NCHW\". "
        "Defaults to \"NHWC\". Specify the data format of the output data, "
        "the input will be transformed automatically. ")
129 130
        .SetDefault("AnyLayout")
        .AsExtra();
131 132 133 134
    AddAttr<bool>(
        "use_quantizer",
        "(bool, default false) "
        "This parameter is no longer used. Use 'mkldnn_data_type' instead.")
135 136
        .SetDefault(false)
        .AsExtra();
137 138 139 140
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
141 142
        .InEnum({"float32", "int8", "bfloat16"})
        .AsExtra();
143
    /* int8 parameters */
X
xzl 已提交
144
    AddComment(R"DOC(
145 146
Transpose Operator.

147 148
The input tensor will be permuted according to the axes given.
The behavior of this operator is similar to how `numpy.transpose` works.
Y
ying 已提交
149

150 151 152 153 154 155
- suppose the input `X` is a 2-D tensor:
    $$
    X = \begin{pmatrix}
    0 &1 &2 \\
    3 &4 &5
    \end{pmatrix}$$
W
wanghaoshuang 已提交
156

157
    the given `axes` is: $[1, 0]$, and $Y$ = transpose($X$, axis)
W
wanghaoshuang 已提交
158

159
    then the output $Y$ is:
W
wanghaoshuang 已提交
160

161 162 163 164 165 166
    $$
    Y = \begin{pmatrix}
         0 &3 \\
         1 &4  \\
         2 &5
    \end{pmatrix}$$
W
wanghaoshuang 已提交
167

168
- Given a input tensor with shape $(N, C, H, W)$ and the `axes` is
169
$[0, 2, 3, 1]$, then shape of the output tensor will be: $(N, H, W, C)$.
170

X
xzl 已提交
171 172 173 174 175 176 177 178
)DOC");
  }
};

class TransposeOpGrad : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

179
  void InferShape(framework::InferShapeContext *ctx) const override {
180
    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "TransposeOpGrad");
181 182 183 184
    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")),
                   "Input",
                   framework::GradVarName("Out"),
                   "TransposeOpGrad");
Q
Qiao Longfei 已提交
185 186 187 188 189
    auto x_dims = ctx->GetInputDim("X");
    ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
    if (ctx->HasOutput(framework::GradVarName("X"))) {
      ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
    }
X
xzl 已提交
190
  }
191 192 193 194

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
195 196
    auto data_type = OperatorWithKernel::IndicateVarDataType(
        ctx, framework::GradVarName("Out"));
J
jiahongyu 已提交
197
    std::string data_format = ctx.Attr<std::string>("data_format");
198
    phi::DataLayout layout_ = phi::StringToDataLayout(data_format);
J
jiahongyu 已提交
199
    return framework::OpKernelType(data_type, ctx.GetPlace(), layout_);
200
  }
X
xzl 已提交
201 202
};

203 204 205 206 207 208 209 210 211 212 213 214 215 216 217
// FIXME(zcd): transpose2 adds an intermediate output(XShape) based on
// transpose, the XShape is used to carry the shape and lod of X which
// will be used in transpose_grad, in this way, the framework can reuse
// the memory of X immediately the transpose2_op is finished.
// Considering compatibility issues, we could not fix transpose2_op
class Transpose2Op : public TransposeOp {
 public:
  Transpose2Op(const std::string &type,
               const framework::VariableNameMap &inputs,
               const framework::VariableNameMap &outputs,
               const framework::AttributeMap &attrs)
      : TransposeOp(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
    TransposeOp::InferShape(ctx);
218
    if (!ctx->HasOutput("XShape")) return;
219 220 221 222 223 224
    const auto &in_dims = ctx->GetInputDim("X");
    std::vector<int64_t> x_shape_dim(in_dims.size() + 1);
    x_shape_dim[0] = 0;
    for (int i = 0; i < in_dims.size(); ++i) {
      x_shape_dim[i + 1] = in_dims[i];
    }
225
    ctx->SetOutputDim("XShape", phi::make_ddim(x_shape_dim));
226 227 228 229 230 231
    ctx->ShareLoD("X", /*->*/ "XShape");
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
232 233
    framework::proto::VarType::Type data_type =
        OperatorWithKernel::IndicateVarDataType(ctx, "X");
J
jiahongyu 已提交
234
    std::string data_format = ctx.Attr<std::string>("data_format");
235
    phi::DataLayout layout_ = phi::StringToDataLayout(data_format);
J
jiahongyu 已提交
236
    return framework::OpKernelType(data_type, ctx.GetPlace(), layout_);
237 238 239
  }
};

240
class Transpose2OpMaker : public framework::OpProtoAndCheckerMaker {
241 242
 public:
  void Make() override {
243 244 245 246 247 248 249 250 251
    AddInput(
        "X",
        "(Tensor) The input tensor, tensors with rank up to 6 are supported.");
    AddOutput("Out", "(Tensor)The output tensor.");
    AddAttr<std::vector<int>>(
        "axis",
        "(vector<int>) A list of values, and the size of the list should be "
        "the same with the input tensor rank. This operator permutes the input "
        "tensor's axes according to the values given.");
252 253 254
    AddOutput("XShape", "(Tensor)The output tensor.")
        .AsIntermediate()
        .AsExtra();
255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
    AddComment(R"DOC(
Transpose Operator.

The input tensor will be permuted according to the axes given.
The behavior of this operator is similar to how `numpy.transpose` works.

- suppose the input `X` is a 2-D tensor:
    $$
    X = \begin{pmatrix}
    0 &1 &2 \\
    3 &4 &5
    \end{pmatrix}$$

    the given `axes` is: $[1, 0]$, and $Y$ = transpose($X$, axis)

    then the output $Y$ is:

    $$
    Y = \begin{pmatrix}
         0 &3 \\
         1 &4  \\
         2 &5
    \end{pmatrix}$$

- Given a input tensor with shape $(N, C, H, W)$ and the `axes` is
$[0, 2, 3, 1]$, then shape of the output tensor will be: $(N, H, W, C)$.

)DOC");
283 284 285
  }
};

H
hong 已提交
286 287
template <typename T>
class Transpose2GradMaker : public framework::SingleGradOpMaker<T> {
288
 public:
H
hong 已提交
289
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
290

291
  void Apply(GradOpPtr<T> grad_op) const override {
292
    grad_op->SetType("transpose2_grad");
H
hong 已提交
293 294 295 296
    grad_op->SetInput("XShape", this->Output("XShape"));
    grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
    grad_op->SetAttrMap(this->Attrs());
297 298 299
  }
};

300 301 302 303 304 305 306 307 308 309 310 311 312 313
template <typename T>
class Transpose2DoubleGradMaker : public framework::SingleGradOpMaker<T> {
 public:
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;

  void Apply(GradOpPtr<T> grad_op) const override {
    grad_op->SetType("transpose2");
    grad_op->SetInput("X", this->OutputGrad(framework::GradVarName("X")));
    grad_op->SetOutput("Out", this->InputGrad(framework::GradVarName("Out")));
    grad_op->SetOutput("XShape", this->Input("XShape"));
    grad_op->SetAttrMap(this->Attrs());
  }
};

314 315 316 317 318
class Transpose2OpGrad : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext *ctx) const override {
319 320 321 322 323
    OP_INOUT_CHECK(
        ctx->HasInput("XShape"), "Input", "XShape", "Transpose2OpGrad");
    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")),
                   "Input",
                   framework::GradVarName("Out"),
324
                   "Transpose2OpGrad");
325 326
    if (ctx->HasOutput(framework::GradVarName("X"))) {
      auto xshape_dim = ctx->GetInputDim("XShape");
327
      auto x_shape_dim = phi::slice_ddim(xshape_dim, 1, xshape_dim.size());
328 329 330 331 332 333 334 335
      ctx->SetOutputDim(framework::GradVarName("X"), x_shape_dim);
      ctx->ShareLoD("XShape", framework::GradVarName("X"));
    }
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
336 337 338
    framework::proto::VarType::Type data_type =
        OperatorWithKernel::IndicateVarDataType(ctx,
                                                framework::GradVarName("Out"));
J
jiahongyu 已提交
339
    std::string data_format = ctx.Attr<std::string>("data_format");
340
    phi::DataLayout layout_ = phi::StringToDataLayout(data_format);
J
jiahongyu 已提交
341
    return framework::OpKernelType(data_type, ctx.GetPlace(), layout_);
342 343 344
  }
};

H
hong 已提交
345 346 347 348 349 350 351 352
class TransposeGradInferVarType : public framework::VarTypeInference {
 public:
  void operator()(framework::InferVarTypeContext *ctx) const override {
    ctx->SyncTypeAndDataType(framework::GradVarName("Out"),
                             framework::GradVarName("X"));
  }
};

X
xzl 已提交
353 354 355 356
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
H
hong 已提交
357
REGISTER_OPERATOR(
358 359 360
    transpose,
    ops::TransposeOp,
    ops::TransposeOpMaker,
H
hong 已提交
361 362
    paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
    paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>);
363 364
REGISTER_OPERATOR(transpose_grad,
                  ops::TransposeOpGrad,
H
hong 已提交
365
                  ops::TransposeGradInferVarType);
366

367 368 369
REGISTER_OPERATOR(transpose2,
                  ops::Transpose2Op,
                  ops::Transpose2OpMaker,
H
hong 已提交
370 371
                  ops::Transpose2GradMaker<paddle::framework::OpDesc>,
                  ops::Transpose2GradMaker<paddle::imperative::OpBase>);
372 373
REGISTER_OPERATOR(transpose2_grad,
                  ops::Transpose2OpGrad,
H
hong 已提交
374
                  ops::TransposeGradInferVarType,
375 376
                  ops::Transpose2DoubleGradMaker<paddle::framework::OpDesc>,
                  ops::Transpose2DoubleGradMaker<paddle::imperative::OpBase>);