create_threaded_reader_op.cc 3.2 KB
Newer Older
F
fengjiayi 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23
//   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/operators/detail/safe_ref.h"
#include "paddle/fluid/operators/reader/reader_op_registry.h"

namespace paddle {
namespace operators {
namespace reader {

class ThreadedReader : public framework::DecoratedReader {
 public:
F
fengjiayi 已提交
24 25
  ThreadedReader(ReaderBase* reader, bool safe_mode)
      : DecoratedReader(reader), safe_mode_(safe_mode) {}
F
fengjiayi 已提交
26 27 28

  void ReadNext(std::vector<framework::LoDTensor>* out) override {
    std::lock_guard<std::mutex> lock(mutex_);
F
fengjiayi 已提交
29
    reader_->ReadNext(out);
F
fengjiayi 已提交
30 31
  }

F
fengjiayi 已提交
32
  void ReInit() override {
F
fengjiayi 已提交
33
    if (safe_mode_) {
F
fengjiayi 已提交
34
      PADDLE_THROW(
F
fengjiayi 已提交
35
          "ThreadedReader::ReInit() is disabled when 'safe_mode' is true.");
F
fengjiayi 已提交
36 37 38 39 40
    }
    VLOG(5) << "ThreadedReader::ReInit() is invoked! It might be buggy in "
               "multi-thread environment.";
    reader_->ReInit();
  }
F
fengjiayi 已提交
41 42

 private:
F
fengjiayi 已提交
43 44
  bool safe_mode_;
  std::mutex mutex_;
F
fengjiayi 已提交
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
};

class CreateThreadedReaderOp : public framework::OperatorBase {
 public:
  using framework::OperatorBase::OperatorBase;

 private:
  void RunImpl(const framework::Scope& scope,
               const platform::Place& dev_place) const override {
    auto* out = detail::Ref(scope.FindVar(Output("Out")))
                    .GetMutable<framework::ReaderHolder>();
    if (out->Get() != nullptr) {
      return;
    }
    const auto& underlying_reader = scope.FindVar(Input("UnderlyingReader"))
                                        ->Get<framework::ReaderHolder>();
F
fengjiayi 已提交
61 62
    bool safe_mode = Attr<bool>("safe_mode");
    out->Reset(new ThreadedReader(underlying_reader.Get(), safe_mode));
F
fengjiayi 已提交
63 64 65 66 67 68 69
  }
};

class CreateThreadedReaderOpMaker : public DecoratedReaderMakerBase {
 public:
  CreateThreadedReaderOpMaker(OpProto* op_proto, OpAttrChecker* op_checker)
      : DecoratedReaderMakerBase(op_proto, op_checker) {
F
fengjiayi 已提交
70 71 72
    AddAttr<bool>("safe_mode",
                  "When 'safe_mode' is true, 'ReInit()' is disabled to avoid "
                  "unexpected bugs in multi-thread environment.")
J
JiayiFeng 已提交
73
        .SetDefault(true);
F
fengjiayi 已提交
74 75 76 77 78 79
    AddComment(R"DOC(
      CreateThreadedReader Operator

      This operator creates a threaded reader. A threaded reader's 
      'ReadNext()' can be invoked by several threads at the same 
      time. 
F
fengjiayi 已提交
80 81 82
      When the attribute 'safe_mode' is true, the threaded reader's 
      'ReInit()' is disabled to avoid unexpected bugs in multi-thread 
      environment.
J
JiayiFeng 已提交
83
    )DOC");
F
fengjiayi 已提交
84 85 86 87 88 89
  }
};

}  // namespace reader
}  // namespace operators
}  // namespace paddle
F
fengjiayi 已提交
90 91

namespace reader = paddle::operators::reader;
J
JiayiFeng 已提交
92 93 94
REGISTER_DECORATED_READER_OPERATOR(create_threaded_reader,
                                   reader::CreateThreadedReaderOp,
                                   reader::CreateThreadedReaderOpMaker);