reshape_op.cc 24.7 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Y
Yibing Liu 已提交
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
Y
Yibing Liu 已提交
6

L
Luo Tao 已提交
7
    http://www.apache.org/licenses/LICENSE-2.0
Y
Yibing Liu 已提交
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. */
Y
Yibing Liu 已提交
14

Y
Yi Wang 已提交
15 16
#include <string>
#include <vector>
Y
yuyang18 已提交
17
#include "paddle/fluid/framework/op_registry.h"
Y
Yi Wang 已提交
18

Y
Yibing Liu 已提交
19 20 21
namespace paddle {
namespace operators {

22 23 24 25 26 27 28 29
using Tensor = framework::Tensor;

inline std::vector<int> get_new_shape(
    const std::vector<const Tensor *> &list_new_shape_tensor) {
  // get tensor from
  std::vector<int> vec_new_shape;
  for (size_t i = 0; i < list_new_shape_tensor.size(); ++i) {
    auto tensor = list_new_shape_tensor[i];
30 31 32 33 34 35
    PADDLE_ENFORCE_EQ(
        tensor->dims(), framework::make_ddim({1}),
        "ShapeError: If the element type of 'shape' in ReshapeOp is Tensor, "
        "the element's shape must be [1]. But received the element's shape "
        "is [%s]",
        tensor->dims());
36 37 38 39 40 41 42 43 44 45 46 47 48
    if (platform::is_gpu_place(tensor->place())) {
      framework::Tensor temp;
      TensorCopySync(*tensor, platform::CPUPlace(), &temp);

      vec_new_shape.push_back(static_cast<int32_t>(*temp.data<int32_t>()));
    } else {
      vec_new_shape.push_back(static_cast<int32_t>(*tensor->data<int32_t>()));
    }
  }

  return vec_new_shape;
}

Y
yuyang18 已提交
49 50 51 52 53 54 55 56
class ReshapeOp : public framework::OperatorWithKernel {
 public:
  ReshapeOp(const std::string &type, const framework::VariableNameMap &inputs,
            const framework::VariableNameMap &outputs,
            const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
57 58 59 60
    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true,
                      "Input(X) of ReshapeOp should not be null.");
    PADDLE_ENFORCE_EQ(ctx->HasOutput("Out"), true,
                      "Output(Out) of ReshapeOp should not be null.");
Y
yuyang18 已提交
61

62 63
    if (ctx->HasInputs("ShapeTensor")) {
      // top prority shape
64
      auto ShapeTensor = ctx->Inputs("ShapeTensor");
65 66 67 68 69 70
      PADDLE_ENFORCE_GT(
          ShapeTensor.size(), 0,
          "ShapeError: When `shape` in ReshapeOp is a list or tuple "
          "which contains Tensor, the shape's size can't be zero. "
          "But received shape's size is %d.",
          ShapeTensor.size());
71 72 73 74 75 76 77
      auto infer_shape = ctx->Attrs().Get<std::vector<int>>("shape");
      const int64_t copy_dim_val = 0;
      auto in_dims = ctx->GetInputDim("X");
      for (size_t i = 0; i < infer_shape.size(); ++i) {
        if (infer_shape[i] == copy_dim_val) {
          PADDLE_ENFORCE_LT(
              static_cast<int>(i), in_dims.size(),
78 79 80 81
              "ShapeError: The index of 0 in `shape` must be less than "
              "the input tensor X's dimensions. But received shape[%d] "
              "= 0, X's dimensions = %d, X's shape = [%s].",
              i, in_dims.size(), in_dims);
82 83 84 85 86 87 88
          infer_shape[i] = in_dims[i];
        }
      }
      auto infer_out_dims = framework::make_ddim(infer_shape);
      ctx->SetOutputDim("Out", infer_out_dims);
      return;
    }
Y
yuyang18 已提交
89

90 91 92 93 94 95 96 97 98 99 100
    const std::vector<int> &shape = ctx->Attrs().Get<std::vector<int>>("shape");
    if (ctx->HasInput("Shape") && shape.empty()) {
      auto shape_dims = ctx->GetInputDim("Shape");
      int num_ele = 1;
      for (int i = 0; i < shape_dims.size(); ++i) {
        num_ele *= shape_dims[i];
      }
      auto vec_dims = std::vector<int>(num_ele, -1);
      auto out_dims = framework::make_ddim(vec_dims);
      ctx->SetOutputDim("Out", out_dims);
      ctx->ShareLoD("X", /*->*/ "Out");
101 102
      return;
    }
103 104

    if (ctx->HasInput("Shape") && !shape.empty() && ctx->IsRuntime()) {
Y
yuyang18 已提交
105 106 107 108 109
      // If true, set the shape of Output(Out) according to Input(Shape) in
      // ReshapeKernel with ExecutionContext. Also check LoD in ReshapeKernel.
      ctx->ShareLoD("X", /*->*/ "Out");
      return;
    }
110

111 112 113 114
    PADDLE_ENFORCE_EQ(
        !shape.empty(), true,
        "ShapeError: The parameter 'shape' in ReshapeOp must be set. "
        "But received 'shape' is empty.");
Y
yuyang18 已提交
115 116 117 118 119 120 121 122 123 124 125 126 127
    auto x_dims = ctx->GetInputDim("X");
    auto out_dims = ValidateShape(shape, x_dims);
    ctx->SetOutputDim("Out", out_dims);
    if (x_dims[0] == out_dims[0]) {
      // Only pass LoD when the first dimension of output and Input(X)
      // are the same.
      ctx->ShareLoD("X", /*->*/ "Out");
    }
  }

  static framework::DDim ValidateShape(const std::vector<int> shape,
                                       const framework::DDim &in_dims) {
    const int64_t in_size = framework::product(in_dims);
C
chengduo 已提交
128 129 130
    auto in_dims_vec = framework::vectorize(in_dims);
    bool all_positive = std::all_of(in_dims_vec.cbegin(), in_dims_vec.cend(),
                                    [](int64_t i) { return i > 0; });
Y
yuyang18 已提交
131 132 133 134 135 136 137 138 139 140
    // only one dimension can be set to -1, whose size will be automatically
    // infered.
    const int64_t unk_dim_val = -1;
    const int64_t copy_dim_val = 0;

    std::vector<int64_t> output_shape(shape.size(), 0);
    int64_t capacity = 1;
    int unk_dim_idx = -1;
    for (size_t i = 0; i < shape.size(); ++i) {
      if (shape[i] == unk_dim_val) {
141 142
        PADDLE_ENFORCE_EQ(
            unk_dim_idx, -1,
143 144 145
            "ShapeError: Only one dimension value of 'shape' in ReshapeOp can "
            "be -1. But received shape = [%s], shape[%d] is also -1.",
            framework::make_ddim(shape), i);
Y
yuyang18 已提交
146 147
        unk_dim_idx = i;
      } else if (shape[i] == copy_dim_val) {
148 149
        PADDLE_ENFORCE_LT(
            static_cast<int>(i), in_dims.size(),
150 151 152 153 154
            "ShapeError: The index of 0 in `shape` must be less than "
            "the input tensor X's dimensions. "
            "But received shape = [%s], shape[%d] = 0, X's shape = [%s], "
            "X's dimensions = %d.",
            framework::make_ddim(shape), i, in_dims, in_dims.size());
Y
yuyang18 已提交
155
      } else {
156 157
        PADDLE_ENFORCE_GT(
            shape[i], 0,
158 159 160 161
            "ShapeError: Each dimension value of 'shape' in ReshapeOp must not "
            "be negtive except one unknown dimension. "
            "But received  shape = [%s], shape[%d] = %d.",
            framework::make_ddim(shape), i, shape[i]);
Y
yuyang18 已提交
162 163 164 165 166 167 168 169
      }

      capacity *= (shape[i] ? shape[i] : in_dims[i]);
      output_shape[i] =
          (shape[i] ? static_cast<int64_t>(shape[i]) : in_dims[i]);
    }

    if (unk_dim_idx != -1) {
C
chengduo 已提交
170
      if (all_positive) {
Y
yuyang18 已提交
171 172 173 174 175
        // in_size < 0 and is un-determinate in compile time, skip the check,
        // for example, in_dims = [-1, 8, 1, 1], shape = [-1, 3, 8],
        // capacity = -24, in_size = -8, output_shape[0] = 0
        // the following check will fail.
        output_shape[unk_dim_idx] = -in_size / capacity;
176 177 178 179 180 181 182 183 184 185
        PADDLE_ENFORCE_EQ(
            output_shape[unk_dim_idx] * capacity, -in_size,
            platform::errors::InvalidArgument(
                "The 'shape' attribute in ReshapeOp is invalid. "
                "The input tensor X'size must be divisible by known "
                "capacity of 'shape'. "
                "But received X's shape = [%s], X's size = %d, "
                "'shape' is [%s], known "
                "capacity of 'shape' is %d.",
                in_dims, in_size, framework::make_ddim(shape), capacity));
Y
yuyang18 已提交
186 187 188 189
      } else {
        output_shape[unk_dim_idx] = -1;
      }
    } else {
Y
Yamei-Lee 已提交
190 191 192 193 194 195 196 197 198 199
      if (all_positive) {
        PADDLE_ENFORCE_EQ(
            capacity, in_size,
            "ShapeError: The 'shape' in ReshapeOp is invalid. "
            "The input tensor X'size must be equal to the capacity of 'shape'. "
            "But received X's shape = [%s], X's size = %d, 'shape' is [%s], "
            "the "
            "capacity of 'shape' is %d.",
            in_dims, in_size, framework::make_ddim(shape), capacity);
      }
Y
yuyang18 已提交
200 201 202 203 204 205 206
    }
    return framework::make_ddim(output_shape);
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
207 208 209
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
Y
yuyang18 已提交
210
  }
211 212 213 214 215 216 217 218 219 220

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "ShapeTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
Y
yuyang18 已提交
221 222
};

Y
Yibing Liu 已提交
223 224
class ReshapeOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
225
  void Make() override {
226 227
    AddInput("X", "(Tensor). The input tensor of reshape operator.");
    AddInput("Shape",
228 229 230
             "(Tensor<int32>, optional). Target shape of reshape operator. "
             "It has a higher priority than Attr(shape) but a lower priority "
             "than Input(ShapeTensor). The Attr(shape) still should be "
231 232
             "set correctly to gurantee shape inference in compile time.")
        .AsDispensable();
233 234
    AddInput(
        "ShapeTensor",
235 236 237 238
        "(vector<Tensor<int32>>, optional). Target shape of reshape operator. "
        "It has the highest priority compare with Input(Shape) and "
        "Attr(shape)."
        "The shape of the element in vector must be [1].")
239 240
        .AsDuplicable()
        .AsDispensable();
241
    AddOutput("Out", "(Tensor). The output tensor of reshape operator.");
C
caoying03 已提交
242
    AddAttr<std::vector<int>>(
243 244 245 246
        "shape",
        "(std::vector<int>) Target shape of reshape operator."
        "It has the lowest priority compare with Input(Shape) and "
        " Input(ShapeTensor).")
247
        .SetDefault({});
K
kexinzhao 已提交
248 249
    AddComment(R"DOC(
Reshape Operator.
Y
Yibing Liu 已提交
250

251 252
Reshape Input(X) into the shape specified by Attr(shape) or Input(Shape). The
data in Input(X) are unchanged.
Y
Yibing Liu 已提交
253

C
caoying03 已提交
254
Examples:
Y
Yibing Liu 已提交
255

C
caoying03 已提交
256 257 258 259
1. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
specified by Attr(shape) is [6, 8], the reshape operator will transform Input(X)
into a 2-D tensor with shape [6, 8] and leaving Input(X)'s data unchanged.

260
2. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
C
caoying03 已提交
261 262 263 264 265 266
specified by Attr(shape) is [2, 3, -1, 2], the reshape operator will transform
Input(X) into a 4-D tensor with shape [2, 3, 4, 2] and leaving Input(X)'s data
unchanged. In this case, one and only dimension of Attr(shape) can be set to -1,
the value of this dimension is inferred from the total element number of
Input(X) and remaining dimensions.

267
3. Given a 3-D tensor Input(X) with a shape [2, 4, 6], and the target shape
C
caoying03 已提交
268 269 270 271
specified by Attr(shape) is [-1, 0, 3, 2], the reshape operator will transform
Input(X) into a 4-D tensor with shape [2, 4, 3, 2] and leaving Input(X)'s data
unchanged. In this case, besides -1, 0 means the actual dimension value is going
to be copied from the corresponding dimension of Input(X).
Y
Yibing Liu 已提交
272

C
caoying03 已提交
273
Note:
Y
Yibing Liu 已提交
274

C
caoying03 已提交
275 276 277
1. One and only one dimension in Attr(shape) can be set -1. In this case,
the actual dimension value will be infered from the total element number of
Input(X) and remaining dimensions.
278 279

2. More than one dimensions in Attr(shape) can be set to 0, which means the real
C
caoying03 已提交
280
dimension value will be copied from Input(X) at runtime. Note that the index of
G
guosheng 已提交
281
0 can not exceed Rank(X). For example, Input(X) is a 3-D tensor with shape
C
caoying03 已提交
282
[2, 3, 4], Attr(shape) = [2, 3, 2, 0] is an invalid input.
283 284

3. Input(Shape) has a higher priority than Attr(shape) if it is provided, while
M
minqiyang 已提交
285
Attr(shape) still should be set correctly to gurantee shape inference in
286
compile-time.
Y
Yibing Liu 已提交
287

Y
Yibing Liu 已提交
288 289 290 291 292 293 294 295 296 297 298 299
)DOC");
  }
};

class ReshapeGradOp : public framework::OperatorWithKernel {
 public:
  ReshapeGradOp(const std::string &type,
                const framework::VariableNameMap &inputs,
                const framework::VariableNameMap &outputs,
                const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

300
  void InferShape(framework::InferShapeContext *ctx) const override {
301 302 303
    PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true, "Input(X) shouldn't be null.");
    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
                      "Input(Out@GRAD) shouldn't be null.");
Q
Qiao Longfei 已提交
304
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
Y
Yibing Liu 已提交
305
  }
306 307 308 309

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
310 311 312
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "X"),
        ctx.device_context());
313
  }
Y
Yibing Liu 已提交
314 315
};

Y
yuyang18 已提交
316 317 318 319 320
class ReshapeKernel {
 public:
  void operator()(const framework::ExecutionContext &ctx) const {
    auto *out = ctx.Output<framework::LoDTensor>("Out");
    auto *in = ctx.Input<framework::LoDTensor>("X");
Y
yuyang18 已提交
321

Y
yuyang18 已提交
322
    framework::DDim out_dims = out->dims();
Y
yuyang18 已提交
323

324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346
    auto list_new_shape_tensor =
        ctx.MultiInput<framework::Tensor>("ShapeTensor");
    if (list_new_shape_tensor.size() > 0) {
      // have shape tensor
      auto new_shape = get_new_shape(list_new_shape_tensor);
      out_dims = ReshapeOp::ValidateShape(new_shape, in->dims());

    } else {
      auto *shape_tensor = ctx.HasInput("Shape")
                               ? ctx.Input<framework::LoDTensor>("Shape")
                               : nullptr;

      if (shape_tensor) {
        auto *shape_data = shape_tensor->data<int>();
        framework::Tensor cpu_shape_tensor;
        if (platform::is_gpu_place(shape_tensor->place())) {
          TensorCopySync(*shape_tensor, platform::CPUPlace(),
                         &cpu_shape_tensor);
          shape_data = cpu_shape_tensor.data<int>();
        }
        auto shape =
            std::vector<int>(shape_data, shape_data + shape_tensor->numel());
        out_dims = ReshapeOp::ValidateShape(shape, in->dims());
Y
yuyang18 已提交
347 348
      }
    }
Y
yuyang18 已提交
349

350
    out->Resize(out_dims);
351
    out->mutable_data(ctx.GetPlace(), in->type());
Y
Yiqun Liu 已提交
352 353 354
    framework::TensorCopy(
        *in, ctx.GetPlace(),
        ctx.template device_context<platform::DeviceContext>(), out);
Y
yuyang18 已提交
355 356
    out->Resize(out_dims);
  }
Y
yuyang18 已提交
357 358 359 360 361 362 363
};

class ReshapeGradKernel {
 public:
  void operator()(const framework::ExecutionContext &ctx) const {
    auto *d_out = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
    auto *d_x = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
D
dzhwinter 已提交
364
    auto in_dims = d_x->dims();
Y
yuyang18 已提交
365

366 367
    d_x->mutable_data(ctx.GetPlace(), d_out->type());
    framework::TensorCopySync(*d_out, ctx.GetPlace(), d_x);
D
dzhwinter 已提交
368
    d_x->Resize(in_dims);
Y
yuyang18 已提交
369
  }
Y
yuyang18 已提交
370 371
};

372 373 374 375 376 377 378 379 380 381 382 383 384 385
class ReshapeDoubleGradKernel {
 public:
  void operator()(const framework::ExecutionContext &ctx) const {
    auto *dd_x = ctx.Input<framework::Tensor>("DDX");
    auto *dd_out = ctx.Output<framework::Tensor>("DDOut");

    auto out_dims = dd_out->dims();

    dd_out->mutable_data(ctx.GetPlace(), dd_x->type());
    framework::TensorCopySync(*dd_x, ctx.GetPlace(), dd_out);
    dd_out->Resize(out_dims);
  }
};

386 387 388 389 390 391 392 393 394 395 396 397 398
// FIXME(zcd): reshape2 adds an intermediate output(XShape) based on reshape,
// the XShape is used to carry the shape and lod of X which will be used in
// reshape_grad, in this way, the framework can reuse the memory of X
// immediately the reshape_op is finished.
// Considering compatibility issues, we could not fix reshape_op
class Reshape2Op : public ReshapeOp {
 public:
  Reshape2Op(const std::string &type, const framework::VariableNameMap &inputs,
             const framework::VariableNameMap &outputs,
             const framework::AttributeMap &attrs)
      : ReshapeOp(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
399 400
    PADDLE_ENFORCE_EQ(ctx->HasOutput("XShape"), true,
                      "Output(XShape) of ReshapeOp should not be null.");
401 402 403 404 405 406 407 408
    const auto &x_dims = ctx->GetInputDim("X");
    std::vector<int64_t> xshape_dims(x_dims.size() + 1);
    xshape_dims[0] = 0;
    for (int i = 0; i < x_dims.size(); ++i) {
      xshape_dims[i + 1] = x_dims[i];
    }
    ctx->SetOutputDim("XShape", framework::make_ddim(xshape_dims));
    ctx->ShareLoD("X", /*->*/ "XShape");
M
minqiyang 已提交
409 410

    ReshapeOp::InferShape(ctx);
411 412 413 414 415 416 417 418 419 420 421 422 423 424
  }
};

class Reshape2OpMaker : public ReshapeOpMaker {
 public:
  void Make() override {
    ReshapeOpMaker::Make();
    AddOutput("XShape",
              "XShape is just used to store the shape and lod of X, which will "
              "be used in FlattenGradOp.")
        .AsIntermediate();
  }
};

H
hong 已提交
425 426
template <typename T>
class Reshape2GradMaker : public framework::SingleGradOpMaker<T> {
427
 public:
H
hong 已提交
428
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
429

H
hong 已提交
430 431
  std::unique_ptr<T> Apply() const override {
    auto *grad_op = new T();
432
    grad_op->SetType("reshape2_grad");
H
hong 已提交
433
    grad_op->SetInput("XShape", this->Output("XShape"));
H
hong 已提交
434 435 436
    if (this->HasInput("ShapeTensor")) {
      grad_op->SetInput("ShapeTensor", this->Input("ShapeTensor"));
    }
H
hong 已提交
437 438 439 440
    grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    grad_op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
    grad_op->SetAttrMap(this->Attrs());
    return std::unique_ptr<T>(grad_op);
441 442 443
  }
};

H
hong 已提交
444 445
template <typename T>
class Reshape2DoubleGradMaker : public framework::SingleGradOpMaker<T> {
446
 public:
H
hong 已提交
447
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
448

H
hong 已提交
449 450
  std::unique_ptr<T> Apply() const override {
    auto *grad_op = new T();
451 452
    grad_op->SetType("reshape2_grad_grad");

H
hong 已提交
453 454 455
    grad_op->SetInput("ShapeTensor", this->Input("ShapeTensor"));
    grad_op->SetInput("DOut", this->Input(framework::GradVarName("Out")));
    grad_op->SetInput("DDX", this->OutputGrad(framework::GradVarName("X")));
456

H
hong 已提交
457 458 459
    grad_op->SetOutput("DDOut", this->InputGrad(framework::GradVarName("Out")));
    grad_op->SetAttrMap(this->Attrs());
    return std::unique_ptr<T>(grad_op);
460 461 462
  }
};

463 464 465 466 467 468 469 470 471
class Reshape2GradOp : public framework::OperatorWithKernel {
 public:
  Reshape2GradOp(const std::string &type,
                 const framework::VariableNameMap &inputs,
                 const framework::VariableNameMap &outputs,
                 const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
472 473 474 475
    PADDLE_ENFORCE_EQ(ctx->HasInput("XShape"), true,
                      "Input(XShape) shouldn't be null.");
    PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true,
                      "Input(Out@GRAD) shouldn't be null.");
476 477 478 479 480 481 482 483 484
    auto xshape_dims = ctx->GetInputDim("XShape");
    auto x_dims = framework::slice_ddim(xshape_dims, 1, xshape_dims.size());
    ctx->SetOutputDim(framework::GradVarName("X"), x_dims);
    ctx->ShareLoD("XShape", framework::GradVarName("X"));
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
485 486 487
    return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
                                       ctx, framework::GradVarName("Out")),
                                   ctx.device_context());
488
  }
489 490 491 492 493 494 495 496 497 498

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "ShapeTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
499 500
};

501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520
class Reshape2DoubleGradOp : public framework::OperatorWithKernel {
 public:
  Reshape2DoubleGradOp(const std::string &type,
                       const framework::VariableNameMap &inputs,
                       const framework::VariableNameMap &outputs,
                       const framework::AttributeMap &attrs)
      : OperatorWithKernel(type, inputs, outputs, attrs) {}

  void InferShape(framework::InferShapeContext *ctx) const override {
    PADDLE_ENFORCE_EQ(ctx->HasInput("DDX"), true,
                      "Input(X@GRAD_GRAD) shouldn't be null.");

    if (ctx->HasOutput("DDOut") && ctx->HasInput("DDX")) {
      ctx->ShareDim("DOut", "DDOut");
    }
  }

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext &ctx) const override {
521 522 523
    return framework::OpKernelType(
        OperatorWithKernel::IndicateVarDataType(ctx, "DDX"),
        ctx.device_context());
524 525 526 527 528 529 530 531 532 533 534 535 536
  }

  framework::OpKernelType GetKernelTypeForVar(
      const std::string &var_name, const Tensor &tensor,
      const framework::OpKernelType &expected_kernel_type) const override {
    if (var_name == "ShapeTensor") {
      return expected_kernel_type;
    }
    return framework::OpKernelType(expected_kernel_type.data_type_,
                                   tensor.place(), tensor.layout());
  }
};

537 538 539 540
DECLARE_INPLACE_OP_INFERER(ReshapeOpInplaceInToOut, {"X", "Out"});
DECLARE_INPLACE_OP_INFERER(ReshapeGradInplaceInToOut,
                           {framework::GradVarName("Out"),
                            framework::GradVarName("X")});
541
DECLARE_INPLACE_OP_INFERER(ReshapeDoubleGradInplaceInToOut, {"DDX", "DDOut"});
D
dzhwinter 已提交
542

Y
Yibing Liu 已提交
543 544 545
}  // namespace operators
}  // namespace paddle
namespace ops = paddle::operators;
546
namespace plat = paddle::platform;
Y
Yibing Liu 已提交
547

H
hong 已提交
548 549 550 551 552
REGISTER_OPERATOR(
    reshape, ops::ReshapeOp, ops::ReshapeOpMaker,
    paddle::framework::DefaultGradOpMaker<paddle::framework::OpDesc, true>,
    paddle::framework::DefaultGradOpMaker<paddle::imperative::OpBase, true>,
    ops::ReshapeOpInplaceInToOut);
D
dzhwinter 已提交
553 554
REGISTER_OPERATOR(reshape_grad, ops::ReshapeGradOp,
                  ops::ReshapeGradInplaceInToOut);
555

556 557 558 559 560 561 562
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape, float, ops::ReshapeKernel, double,
                               ops::ReshapeKernel, int, ops::ReshapeKernel,
                               int64_t, ops::ReshapeKernel);
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape_grad, float, ops::ReshapeGradKernel,
                               double, ops::ReshapeGradKernel, int,
                               ops::ReshapeGradKernel, int64_t,
                               ops::ReshapeGradKernel);
563
REGISTER_OPERATOR(reshape2, ops::Reshape2Op, ops::Reshape2OpMaker,
H
hong 已提交
564 565 566
                  ops::Reshape2GradMaker<paddle::framework::OpDesc>,
                  ops::Reshape2GradMaker<paddle::imperative::OpBase>,
                  ops::ReshapeOpInplaceInToOut);
D
dzhwinter 已提交
567
REGISTER_OPERATOR(reshape2_grad, ops::Reshape2GradOp,
H
hong 已提交
568 569 570
                  ops::Reshape2DoubleGradMaker<paddle::framework::OpDesc>,
                  ops::Reshape2DoubleGradMaker<paddle::imperative::OpBase>,
                  ops::ReshapeGradInplaceInToOut);
571 572 573
REGISTER_OPERATOR(reshape2_grad_grad, ops::Reshape2DoubleGradOp,
                  ops::ReshapeDoubleGradInplaceInToOut);

574 575 576 577 578 579 580
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2, float, ops::ReshapeKernel, double,
                               ops::ReshapeKernel, int, ops::ReshapeKernel,
                               int64_t, ops::ReshapeKernel);
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2_grad, float, ops::ReshapeGradKernel,
                               double, ops::ReshapeGradKernel, int,
                               ops::ReshapeGradKernel, int64_t,
                               ops::ReshapeGradKernel);
581 582 583 584 585
REGISTER_OP_CPU_KERNEL_FUNCTOR(reshape2_grad_grad, float,
                               ops::ReshapeDoubleGradKernel, double,
                               ops::ReshapeDoubleGradKernel, int,
                               ops::ReshapeDoubleGradKernel, int64_t,
                               ops::ReshapeDoubleGradKernel);
586

Y
yuyang18 已提交
587
#ifdef PADDLE_WITH_CUDA
588 589
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape, float, ops::ReshapeKernel, double,
                                ops::ReshapeKernel, int, ops::ReshapeKernel,
590 591
                                int64_t, ops::ReshapeKernel, plat::float16,
                                ops::ReshapeKernel);
592 593 594
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape_grad, float, ops::ReshapeGradKernel,
                                double, ops::ReshapeGradKernel, int,
                                ops::ReshapeGradKernel, int64_t,
595
                                ops::ReshapeGradKernel, plat::float16,
596 597 598
                                ops::ReshapeGradKernel);
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2, float, ops::ReshapeKernel, double,
                                ops::ReshapeKernel, int, ops::ReshapeKernel,
599 600
                                int64_t, ops::ReshapeKernel, plat::float16,
                                ops::ReshapeKernel);
601 602 603
REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2_grad, float, ops::ReshapeGradKernel,
                                double, ops::ReshapeGradKernel, int,
                                ops::ReshapeGradKernel, int64_t,
604
                                ops::ReshapeGradKernel, plat::float16,
605
                                ops::ReshapeGradKernel);
606 607 608 609 610 611 612

REGISTER_OP_CUDA_KERNEL_FUNCTOR(reshape2_grad_grad, float,
                                ops::ReshapeDoubleGradKernel, double,
                                ops::ReshapeDoubleGradKernel, int,
                                ops::ReshapeDoubleGradKernel, int64_t,
                                ops::ReshapeDoubleGradKernel, plat::float16,
                                ops::ReshapeDoubleGradKernel);
Y
yuyang18 已提交
613
#endif