// 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. #include "paddle/fluid/framework/executor.h" #include "paddle/fluid/operators/reader/reader_op_registry.h" namespace paddle { namespace operators { namespace reader { class CustomReader : public framework::DecoratedReader { public: CustomReader(ReaderBase* reader, const framework::BlockDesc& sub_block, const framework::Scope& scope, const platform::Place& dev_place, const std::vector& source_var_names, const std::vector& sink_var_names) : DecoratedReader(reader), sub_block_(sub_block), scope_(scope), dev_place_(dev_place), source_var_names_(source_var_names), sink_var_names_(sink_var_names) {} void ReadNext(std::vector* out) override; private: const framework::BlockDesc& sub_block_; const framework::Scope& scope_; platform::Place dev_place_; std::vector source_var_names_; std::vector sink_var_names_; }; class CreateCustomReaderOp : public framework::OperatorBase { public: using framework::OperatorBase::OperatorBase; private: void RunImpl(const framework::Scope& scope, const platform::Place& dev_place) const override { auto* out = scope.FindVar(Output("Out")) ->template GetMutable(); if (out->Get() != nullptr) { return; } const auto& underlying_reader = scope.FindVar(Input("UnderlyingReader")) ->Get(); out->Reset(new CustomReader( underlying_reader.Get(), *Attr("sub_block"), scope, dev_place, Attr>("source_var_names"), Attr>("sink_var_names"))); } }; class CreateCustomReaderOpMaker : public DecoratedReaderMakerBase { public: CreateCustomReaderOpMaker(OpProto* op_proto, OpAttrChecker* op_checker) : DecoratedReaderMakerBase(op_proto, op_checker) { AddAttr("sub_block", ""); AddAttr>("source_var_names", ""); AddAttr>("sink_var_names", ""); AddComment(R"DOC( CreateCustomReader Operator )DOC"); } }; class CustomReaderInferShape : public framework::InferShapeBase { public: void operator()(framework::InferShapeContext* ctx) const override { PADDLE_ENFORCE(!ctx->IsRuntime(), "'CustomReaderInferShape' should only be invoked during " "compile time."); PADDLE_ENFORCE(ctx->HasOutput("Out"), "The output decorated reader should not be null."); const auto sink_var_names = ctx->Attrs().Get>("sink_var_names"); std::vector> res_dims; std::vector res_lod_levels; for (const std::string& var_name : sink_var_names) { auto* sink_var = boost::get(ctx->GetVarPtr(var_name)); PADDLE_ENFORCE_NOT_NULL(sink_var); res_dims.emplace_back(sink_var->GetShape()); res_lod_levels.push_back(sink_var->GetLoDLevel()); } auto* out_reader = boost::get(ctx->GetOutputVarPtrs("Out")[0]); out_reader->SetShapes(res_dims); out_reader->SetLoDLevels(res_lod_levels); } }; class CustomReaderInferVarType : public framework::VarTypeInference { public: void operator()(const framework::OpDesc& op_desc, framework::BlockDesc* block) const override { framework::VarDesc* out_reader = block->FindVar(op_desc.Output("Out")[0]); PADDLE_ENFORCE_NOT_NULL(out_reader); out_reader->SetType(framework::proto::VarType::READER); auto sink_var_names = boost::get>(op_desc.GetAttr("sink_var_names")); std::vector res_data_types; for (const std::string& var_name : sink_var_names) { framework::VarDesc* var = block->FindVar(var_name); PADDLE_ENFORCE_NOT_NULL(var); res_data_types.emplace_back(var->GetDataType()); } out_reader->SetDataTypes(res_data_types); } }; void CustomReader::ReadNext(std::vector* out) { out->clear(); std::vector underlying_outs; reader_->ReadNext(&underlying_outs); if (underlying_outs.empty()) { // There is not next data. return; } PADDLE_ENFORCE( source_var_names_.size() == underlying_outs.size() && sink_var_names_.size() == underlying_outs.size(), "The size of source_var_names(%d), the size of sink_var_names(%d) and " "the size of underlying_outs(%d) are not consistent. Each feeding " "element must have its own source and sink variable.", source_var_names_.size(), sink_var_names_.size(), underlying_outs.size()); // 1. Copy LoDTensors from underlying reader's output to source variables. for (size_t i = 0; i < source_var_names_.size(); ++i) { framework::Variable* var = scope_.FindVar(source_var_names_[i]); PADDLE_ENFORCE_NOT_NULL( var, "CustomReader's source variable '%s' doesn't exist."); framework::LoDTensor* tensor = var->GetMutable(); tensor->ShareDataWith(underlying_outs[i]); tensor->set_lod(underlying_outs[i].lod()); } // 2. Run the sub-block. framework::Executor executor(dev_place_); framework::ProgramDesc* program = sub_block_.Program(); framework::Scope* exe_scope = &scope_.NewScope(); executor.Run(*program, exe_scope, sub_block_.ID(), false /*create_local_scope*/, true); scope_.DeleteScope(exe_scope); // 3. Copy LoDTensors from sink variables to out. out->resize(sink_var_names_.size()); for (size_t i = 0; i < sink_var_names_.size(); ++i) { framework::Variable* var = scope_.FindVar(sink_var_names_[i]); PADDLE_ENFORCE_NOT_NULL(var, "CustomReader's sink variable '%s' doesn't exist."); const framework::LoDTensor& tensor = var->Get(); (*out)[i].ShareDataWith(tensor); (*out)[i].set_lod(tensor.lod()); } } } // namespace reader } // namespace operators } // namespace paddle namespace ops = paddle::operators::reader; REGISTER_OPERATOR(create_custom_reader, ops::CreateCustomReaderOp, ops::CreateCustomReaderOpMaker, ops::CustomReaderInferShape, ops::CustomReaderInferVarType, paddle::framework::EmptyGradOpMaker)