read_op.cc 3.7 KB
Newer Older
F
fengjiayi 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
//   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
//
// 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
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// 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
Yi Wang 已提交
15 16
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/reader.h"
F
fengjiayi 已提交
17 18 19 20 21 22 23 24 25 26 27

namespace paddle {
namespace operators {

class ReadInferShape : public framework::InferShapeBase {
 public:
  void operator()(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("Reader"),
                   "The ReadOp must take a reader as input.");
    PADDLE_ENFORCE(ctx->HasOutputs("Out"),
                   "The ReadOp should be assigned with output.");
F
fengjiayi 已提交
28
    std::vector<framework::DDim> reader_dims = ctx->GetReaderDims("Reader");
F
fengjiayi 已提交
29 30 31 32 33 34 35 36 37 38 39 40 41 42
    std::vector<std::string> out_names = ctx->Outputs("Out");
    PADDLE_ENFORCE_EQ(
        reader_dims.size(), out_names.size(),
        "The reader's dim number doesn't match the output number.");
    ctx->SetOutputsDim("Out", reader_dims);
  }
};

class ReadInferVarType : public framework::VarTypeInference {
 public:
  void operator()(const framework::OpDesc& op_desc,
                  framework::BlockDesc* block) const override {
    std::string reader_name = op_desc.Input("Reader")[0];
    std::vector<std::string> out_names = op_desc.Output("Out");
F
fengjiayi 已提交
43 44
    framework::VarDesc* reader = block->FindVarRecursive(reader_name);
    auto dtypes = reader->GetDataTypes();
F
fengjiayi 已提交
45 46
    PADDLE_ENFORCE_EQ(dtypes.size(), out_names.size());
    for (size_t i = 0; i < dtypes.size(); ++i) {
F
fengjiayi 已提交
47 48
      framework::VarDesc& out = block->FindRecursiveOrCreateVar(out_names[i]);
      out.SetType(framework::proto::VarDesc::LOD_TENSOR);
F
fengjiayi 已提交
49 50 51 52 53 54 55 56
      out.SetDataType(dtypes[i]);
    }
  }
};

class ReadOp : public framework::OperatorBase {
 public:
  using framework::OperatorBase::OperatorBase;
57 58 59 60

 private:
  void RunImpl(const framework::Scope& scope,
               const platform::Place& dev_place) const override {
F
fengjiayi 已提交
61 62 63
    framework::ReaderHolder* reader =
        scope.FindVar(Input("Reader"))->GetMutable<framework::ReaderHolder>();
    if (!reader->HasNext()) {
F
fengjiayi 已提交
64 65 66 67
      reader->ReInit();
      PADDLE_ENFORCE(
          reader->HasNext(),
          "Reader can not read the next data even it has been re-initialized.");
F
fengjiayi 已提交
68 69 70
    }
    std::vector<std::string> out_arg_names = Outputs("Out");
    std::vector<framework::LoDTensor> ins;
F
fengjiayi 已提交
71
    reader->ReadNext(&ins);
F
fengjiayi 已提交
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91
    PADDLE_ENFORCE_EQ(ins.size(), out_arg_names.size());
    for (size_t i = 0; i < ins.size(); ++i) {
      auto* out =
          scope.FindVar(out_arg_names[i])->GetMutable<framework::LoDTensor>();
      out->ShareDataWith(ins[i]);
      out->set_lod(ins[i].lod());
    }
  }
};

class ReadOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
  ReadOpMaker(OpProto* op_proto, OpAttrChecker* op_checker)
      : OpProtoAndCheckerMaker(op_proto, op_checker) {
    AddInput("Reader", "(ReaderHolder) The executed reader.");
    AddOutput("Out", "(LoDTensor) The output data.").AsDuplicable();
    AddComment(R"DOC(
      Read Operator

      Execute a given reader once and output data.
F
fengjiayi 已提交
92
    )DOC");
F
fengjiayi 已提交
93 94 95 96
  }
};

}  // namespace operators
F
fengjiayi 已提交
97 98 99 100 101
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OPERATOR(read, ops::ReadOp, ops::ReadInferShape, ops::ReadOpMaker,
                  paddle::framework::EmptyGradOpMaker, ops::ReadInferVarType);