|
| 1 | +// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +// you may not use this file except in compliance with the License. |
| 5 | +// You may obtain a copy of the License at |
| 6 | +// |
| 7 | +// http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +// |
| 9 | +// Unless required by applicable law or agreed to in writing, software |
| 10 | +// distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +// See the License for the specific language governing permissions and |
| 13 | +// limitations under the License. |
| 14 | + |
| 15 | +#include "paddle/fluid/framework/executor.h" |
| 16 | +#include "paddle/fluid/operators/detail/safe_ref.h" |
| 17 | +#include "paddle/fluid/operators/reader/reader_op_registry.h" |
| 18 | + |
| 19 | +namespace paddle { |
| 20 | +namespace operators { |
| 21 | +namespace reader { |
| 22 | + |
| 23 | +class CustomReader : public framework::DecoratedReader { |
| 24 | + public: |
| 25 | + CustomReader(ReaderBase* reader, const framework::BlockDesc& sub_block, |
| 26 | + const platform::Place& dev_place, |
| 27 | + const std::vector<std::string>& source_var_names, |
| 28 | + const std::vector<std::string>& sink_var_names) |
| 29 | + : DecoratedReader(reader), |
| 30 | + program_(*sub_block.Program()), |
| 31 | + sub_block_id_(sub_block.ID()), |
| 32 | + exe_(framework::Executor(dev_place)), |
| 33 | + source_var_names_(source_var_names), |
| 34 | + sink_var_names_(sink_var_names) {} |
| 35 | + |
| 36 | + void ReadNext(std::vector<framework::LoDTensor>* out) override; |
| 37 | + |
| 38 | + private: |
| 39 | + const framework::ProgramDesc program_; |
| 40 | + int sub_block_id_; |
| 41 | + framework::Executor exe_; |
| 42 | + |
| 43 | + std::vector<std::string> source_var_names_; |
| 44 | + std::vector<std::string> sink_var_names_; |
| 45 | +}; |
| 46 | + |
| 47 | +class CreateCustomReaderOp : public framework::OperatorBase { |
| 48 | + public: |
| 49 | + using framework::OperatorBase::OperatorBase; |
| 50 | + |
| 51 | + private: |
| 52 | + void RunImpl(const framework::Scope& scope, |
| 53 | + const platform::Place& dev_place) const override { |
| 54 | + auto* out = scope.FindVar(Output("Out")) |
| 55 | + ->template GetMutable<framework::ReaderHolder>(); |
| 56 | + auto* sub_block = Attr<framework::BlockDesc*>("sub_block"); |
| 57 | + if (out->Get() != nullptr) { |
| 58 | + return; |
| 59 | + } |
| 60 | + const auto& underlying_reader = scope.FindVar(Input("UnderlyingReader")) |
| 61 | + ->Get<framework::ReaderHolder>(); |
| 62 | + out->Reset( |
| 63 | + new CustomReader(underlying_reader.Get(), *sub_block, dev_place, |
| 64 | + Attr<std::vector<std::string>>("source_var_names"), |
| 65 | + Attr<std::vector<std::string>>("sink_var_names"))); |
| 66 | + } |
| 67 | +}; |
| 68 | + |
| 69 | +class CreateCustomReaderOpMaker : public DecoratedReaderMakerBase { |
| 70 | + protected: |
| 71 | + void Apply() override { |
| 72 | + AddAttr<framework::BlockDesc*>( |
| 73 | + "sub_block", "The block to hold all preprocessing operators."); |
| 74 | + AddAttr<std::vector<std::string>>( |
| 75 | + "source_var_names", |
| 76 | + "Source variables are starting points of data preprocessing. They hold " |
| 77 | + "preprocessing's input tensors. Each source variable corresponds to " |
| 78 | + "one of underlying reader's output datas."); |
| 79 | + AddAttr<std::vector<std::string>>( |
| 80 | + "sink_var_names", |
| 81 | + "Sink variables are ending points of data preprocessing. They hold " |
| 82 | + "preprocessing's output tensors. Each sink variable corresponds to " |
| 83 | + "one of custom reader's output datas."); |
| 84 | + AddComment(R"DOC( |
| 85 | + CreateCustomReader Operator |
| 86 | +
|
| 87 | + A custom reader can be used for input data preprocessing. |
| 88 | + A custom reader holds its own sub-block, which will be executed in its |
| 89 | + 'ReadNext()' function. Users can configurate their own preprocessing |
| 90 | + pipelines by inserting operators into custom reader's sub-block. |
| 91 | + )DOC"); |
| 92 | + } |
| 93 | +}; |
| 94 | + |
| 95 | +class CustomReaderInferShape : public framework::InferShapeBase { |
| 96 | + public: |
| 97 | + void operator()(framework::InferShapeContext* ctx) const override { |
| 98 | + PADDLE_ENFORCE(!ctx->IsRuntime(), |
| 99 | + "'CustomReaderInferShape' should only be invoked during " |
| 100 | + "compile time."); |
| 101 | + PADDLE_ENFORCE(ctx->HasOutput("Out"), |
| 102 | + "The output decorated reader should not be null."); |
| 103 | + const auto* sub_block = |
| 104 | + ctx->Attrs().Get<framework::BlockDesc*>("sub_block"); |
| 105 | + const auto sink_var_names = |
| 106 | + ctx->Attrs().Get<std::vector<std::string>>("sink_var_names"); |
| 107 | + std::vector<std::vector<int64_t>> res_dims; |
| 108 | + std::vector<int32_t> res_lod_levels; |
| 109 | + for (const std::string& var_name : sink_var_names) { |
| 110 | + auto* sink_var = sub_block->FindVar(var_name); |
| 111 | + PADDLE_ENFORCE_NOT_NULL(sink_var); |
| 112 | + res_dims.emplace_back(sink_var->GetShape()); |
| 113 | + res_lod_levels.push_back(sink_var->GetLoDLevel()); |
| 114 | + } |
| 115 | + auto* out_reader = |
| 116 | + boost::get<framework::VarDesc*>(ctx->GetOutputVarPtrs("Out")[0]); |
| 117 | + out_reader->SetShapes(res_dims); |
| 118 | + out_reader->SetLoDLevels(res_lod_levels); |
| 119 | + } |
| 120 | +}; |
| 121 | + |
| 122 | +class CustomReaderInferVarType : public framework::VarTypeInference { |
| 123 | + public: |
| 124 | + void operator()(const framework::OpDesc& op_desc, |
| 125 | + framework::BlockDesc* block) const override { |
| 126 | + framework::VarDesc* out_reader = block->FindVar(op_desc.Output("Out")[0]); |
| 127 | + PADDLE_ENFORCE_NOT_NULL(out_reader); |
| 128 | + out_reader->SetType(framework::proto::VarType::READER); |
| 129 | + |
| 130 | + auto sink_var_names = |
| 131 | + boost::get<std::vector<std::string>>(op_desc.GetAttr("sink_var_names")); |
| 132 | + const auto* sub_block = |
| 133 | + boost::get<framework::BlockDesc*>(op_desc.GetAttr("sub_block")); |
| 134 | + std::vector<framework::proto::VarType::Type> res_data_types; |
| 135 | + for (const std::string& var_name : sink_var_names) { |
| 136 | + framework::VarDesc* var = sub_block->FindVar(var_name); |
| 137 | + PADDLE_ENFORCE_NOT_NULL(var); |
| 138 | + res_data_types.emplace_back(var->GetDataType()); |
| 139 | + } |
| 140 | + out_reader->SetDataTypes(res_data_types); |
| 141 | + } |
| 142 | +}; |
| 143 | + |
| 144 | +void CustomReader::ReadNext(std::vector<framework::LoDTensor>* out) { |
| 145 | + out->clear(); |
| 146 | + std::vector<framework::LoDTensor> underlying_outs; |
| 147 | + reader_->ReadNext(&underlying_outs); |
| 148 | + if (underlying_outs.empty()) { |
| 149 | + // There is not next data. |
| 150 | + return; |
| 151 | + } |
| 152 | + PADDLE_ENFORCE(source_var_names_.size() == underlying_outs.size(), |
| 153 | + "The size of source_var_names(%d) and the size of " |
| 154 | + "underlying_outs(%d) are not consistent. Each feeding element " |
| 155 | + "must have its own source variable.", |
| 156 | + source_var_names_.size(), underlying_outs.size()); |
| 157 | + // The scope for CustomReader's sub-block should be independent and shouldn't |
| 158 | + // be any other computation scope's child. Otherwise, data preprocessing and |
| 159 | + // compution cannot be concurrent. |
| 160 | + framework::Scope scope; |
| 161 | + // 1. Copy LoDTensors from underlying reader's output to source variables. |
| 162 | + for (size_t i = 0; i < source_var_names_.size(); ++i) { |
| 163 | + framework::Variable* var = scope.Var(source_var_names_[i]); |
| 164 | + framework::LoDTensor* tensor = var->GetMutable<framework::LoDTensor>(); |
| 165 | + tensor->ShareDataWith(underlying_outs[i]); |
| 166 | + tensor->set_lod(underlying_outs[i].lod()); |
| 167 | + } |
| 168 | + // 2. Run the sub-block. |
| 169 | + exe_.Run(program_, &scope, sub_block_id_, false, true); |
| 170 | + // 3. Copy LoDTensors from sink variables to out. |
| 171 | + out->resize(sink_var_names_.size()); |
| 172 | + for (size_t i = 0; i < sink_var_names_.size(); ++i) { |
| 173 | + const auto& tensor = detail::Ref(scope.FindVar(sink_var_names_[i])) |
| 174 | + .Get<framework::LoDTensor>(); |
| 175 | + framework::TensorCopySync(tensor, platform::CPUPlace(), &(*out)[i]); |
| 176 | + } |
| 177 | +} |
| 178 | + |
| 179 | +} // namespace reader |
| 180 | +} // namespace operators |
| 181 | +} // namespace paddle |
| 182 | + |
| 183 | +namespace ops = paddle::operators::reader; |
| 184 | +REGISTER_OPERATOR(create_custom_reader, ops::CreateCustomReaderOp, |
| 185 | + ops::CreateCustomReaderOpMaker, ops::CustomReaderInferShape, |
| 186 | + ops::CustomReaderInferVarType, |
| 187 | + paddle::framework::EmptyGradOpMaker) |
0 commit comments