read_op.cc 3.9 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"
Y
Yu Yang 已提交
17
#include "paddle/fluid/operators/detail/safe_ref.h"
F
fengjiayi 已提交
18 19 20 21 22 23 24 25 26 27 28

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 已提交
29
    std::vector<framework::DDim> reader_dims = ctx->GetReaderDims("Reader");
F
fengjiayi 已提交
30 31 32 33 34 35 36 37 38 39 40 41 42 43
    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 已提交
44 45
    framework::VarDesc* reader = block->FindVarRecursive(reader_name);
    auto dtypes = reader->GetDataTypes();
F
fengjiayi 已提交
46 47
    PADDLE_ENFORCE_EQ(dtypes.size(), out_names.size());
    for (size_t i = 0; i < dtypes.size(); ++i) {
F
fengjiayi 已提交
48
      framework::VarDesc& out = block->FindRecursiveOrCreateVar(out_names[i]);
49
      out.SetType(framework::proto::VarType::LOD_TENSOR);
F
fengjiayi 已提交
50 51 52 53 54 55 56 57
      out.SetDataType(dtypes[i]);
    }
  }
};

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

 private:
  void RunImpl(const framework::Scope& scope,
               const platform::Place& dev_place) const override {
F
fengjiayi 已提交
62
    framework::ReaderHolder* reader =
Y
Yu Yang 已提交
63 64 65
        detail::Ref(scope.FindVar(Input("Reader")),
                    "Cannot find reader variable %s", Input("Reader"))
            .GetMutable<framework::ReaderHolder>();
F
fengjiayi 已提交
66 67 68 69
    std::vector<std::string> out_arg_names = Outputs("Out");
    std::vector<framework::LoDTensor> ins;
    reader->ReadNext(&ins);
    if (ins.empty()) {
F
fengjiayi 已提交
70
      reader->ReInit();
F
fengjiayi 已提交
71
      reader->ReadNext(&ins);
F
fengjiayi 已提交
72
      PADDLE_ENFORCE(
F
fengjiayi 已提交
73
          !ins.empty(),
F
fengjiayi 已提交
74
          "Reader can not read the next data even it has been re-initialized.");
F
fengjiayi 已提交
75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
    }
    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 已提交
96
    )DOC");
F
fengjiayi 已提交
97 98 99 100
  }
};

}  // namespace operators
F
fengjiayi 已提交
101 102 103 104 105
}  // namespace paddle

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