reshape_op.h 1.7 KB
Newer Older
Y
Yibing Liu 已提交
1 2 3 4 5 6
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

   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

7
       http://www.apache.org/licenses/LICENSE-2.0
Y
Yibing Liu 已提交
8 9 10 11 12 13 14 15 16 17 18 19 20 21 22

   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/framework/eigen.h"
#include "paddle/framework/op_registry.h"

namespace paddle {
namespace operators {

Q
QI JUN 已提交
23
template <typename DeviceContext, typename T>
24
class ReshapeKernel : public framework::OpKernel<T> {
Y
Yibing Liu 已提交
25 26
 public:
  void Compute(const framework::ExecutionContext& ctx) const {
Y
Yibing Liu 已提交
27 28
    auto* out = ctx.Output<framework::Tensor>("Out");
    auto* in = ctx.Input<framework::Tensor>("X");
29
    auto out_dims = out->dims();
Y
Yibing Liu 已提交
30
    out->mutable_data<T>(ctx.GetPlace());
D
dzhwinter 已提交
31
    framework::CopyFrom(*in, ctx.GetPlace(), ctx.device_context(), out);
Y
Yibing Liu 已提交
32 33 34 35
    out->Resize(out_dims);
  }
};

Q
QI JUN 已提交
36
template <typename DeviceContext, typename T>
37
class ReshapeGradKernel : public framework::OpKernel<T> {
Y
Yibing Liu 已提交
38 39
 public:
  void Compute(const framework::ExecutionContext& ctx) const {
Y
Yibing Liu 已提交
40 41
    auto* d_out = ctx.Input<framework::Tensor>(framework::GradVarName("Out"));
    auto* d_x = ctx.Output<framework::Tensor>(framework::GradVarName("X"));
Y
Yibing Liu 已提交
42 43 44
    d_x->mutable_data<T>(ctx.GetPlace());

    auto in_dims = d_x->dims();
D
dzhwinter 已提交
45
    framework::CopyFrom(*d_out, ctx.GetPlace(), ctx.device_context(), d_x);
Y
Yibing Liu 已提交
46 47 48
    d_x->Resize(in_dims);
  }
};
H
Helin Wang 已提交
49 50
}  // namespace operators
}  // namespace paddle
新手
引导
客服 返回
顶部