feed_op.cc 2.7 KB
Newer Older
Q
qijun 已提交
1 2
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.

Y
Yu Yang 已提交
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
Q
qijun 已提交
6

Y
Yu Yang 已提交
7
   http://www.apache.org/licenses/LICENSE-2.0
Q
qijun 已提交
8

Y
Yu Yang 已提交
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. */
Q
qijun 已提交
14

Y
Yu Yang 已提交
15 16 17
#include "paddle/framework/feed_fetch_type.h"
#include "paddle/framework/op_registry.h"
#include "paddle/framework/operator.h"
Q
qijun 已提交
18 19 20

namespace paddle {
namespace operators {
Y
Yu Yang 已提交
21
class FeedOp : public framework::OperatorBase {
Q
qijun 已提交
22
 public:
Y
Yu Yang 已提交
23 24 25 26 27 28
  FeedOp(const std::string &type, const framework::VariableNameMap &inputs,
         const framework::VariableNameMap &outputs,
         const framework::AttributeMap &attrs)
      : OperatorBase(type, inputs, outputs, attrs) {}
  void Run(const framework::Scope &scope,
           const platform::DeviceContext &dev_ctx) const override {
Y
Yu Yang 已提交
29
    auto feed_var_name = Input("X");
Y
Yu Yang 已提交
30
    auto *feed_var = scope.FindVar(feed_var_name);
Y
Yu Yang 已提交
31

Y
Yu Yang 已提交
32 33 34 35 36 37 38 39 40 41 42 43
    PADDLE_ENFORCE(feed_var != nullptr,
                   "Cannot find feed_var in scope, feed_var_name is %s",
                   feed_var_name);

    auto out_name = this->Output("Out");
    auto *out_var = scope.FindVar(out_name);
    PADDLE_ENFORCE(out_var != nullptr,
                   "Cannot find out_var in scope, out_var_name is %s",
                   out_name);

    auto col = Attr<int>("col");

Y
Yu Yang 已提交
44 45 46
    VLOG(3) << "Feed Var " << feed_var_name << "'s " << col << " column to var"
            << out_name;

Y
Yu Yang 已提交
47 48 49
    auto &feed_list = feed_var->Get<framework::FeedFetchList>();
    auto &feed_item = feed_list.at(static_cast<size_t>(col));
    auto *out_item = out_var->GetMutable<framework::FeedFetchType>();
50
    out_item->CopyFrom(feed_item, dev_ctx.GetPlace(), dev_ctx);
Y
Yu Yang 已提交
51
    out_item->set_lod(feed_item.lod());
Q
qijun 已提交
52 53 54
  }
};

Y
Yu Yang 已提交
55 56 57 58 59 60 61 62 63 64 65 66
class FeedOpInfoMaker : public framework::OpProtoAndCheckerMaker {
 public:
  FeedOpInfoMaker(framework::OpProto *proto,
                  framework::OpAttrChecker *op_checker)
      : OpProtoAndCheckerMaker(proto, op_checker) {
    AddInput("X", "The input of feed op");
    AddOutput("Out", "The output of feed op");
    AddComment("feed op, it should not be configured by users directly");
    AddAttr<int>("col", "column of feed");
  }
};

Q
qijun 已提交
67 68 69
}  // namespace operators
}  // namespace paddle

Y
Yu Yang 已提交
70
REGISTER_OPERATOR(feed, paddle::operators::FeedOp,
Y
Yu Yang 已提交
71 72
                  paddle::framework::EmptyGradOpMaker,
                  paddle::operators::FeedOpInfoMaker);