ctr_reader.h 3.7 KB
Newer Older
Q
Qiao Longfei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
// 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.

#pragma once

Q
Qiao Longfei 已提交
17 18 19 20 21 22
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <unordered_map>
Q
Qiao Longfei 已提交
23
#include <vector>
Q
Qiao Longfei 已提交
24

Q
Qiao Longfei 已提交
25
#include "paddle/fluid/framework/reader.h"
Q
Qiao Longfei 已提交
26
#include "paddle/fluid/framework/threadpool.h"
Q
Qiao Longfei 已提交
27 28 29 30 31 32
#include "paddle/fluid/operators/reader/lod_tensor_blocking_queue.h"

namespace paddle {
namespace operators {
namespace reader {

Q
Qiao Longfei 已提交
33 34
enum ReaderThreadStatus { Running, Stopped };

Q
Qiao Longfei 已提交
35 36
void ReadThread(const std::vector<std::string>& file_list,
                const std::vector<std::string>& slots, int batch_size,
Q
Qiao Longfei 已提交
37
                int thread_id, std::vector<ReaderThreadStatus>* thread_status,
Q
Qiao Longfei 已提交
38 39
                std::shared_ptr<LoDTensorBlockingQueue> queue);

Q
Qiao Longfei 已提交
40 41
class CTRReader : public framework::FileReader {
 public:
Q
Qiao Longfei 已提交
42 43 44 45
  explicit CTRReader(const std::shared_ptr<LoDTensorBlockingQueue>& queue,
                     int batch_size, int thread_num,
                     const std::vector<std::string>& slots,
                     const std::vector<std::string>& file_list)
Q
Qiao Longfei 已提交
46
      : batch_size_(batch_size), slots_(slots), file_list_(file_list) {
Q
Qiao Longfei 已提交
47
    PADDLE_ENFORCE(queue != nullptr, "LoDTensorBlockingQueue must not be null");
Q
Qiao Longfei 已提交
48 49
    PADDLE_ENFORCE_GT(file_list.size(), 0, "file list should not be empty");
    thread_num_ =
Q
Qiao Longfei 已提交
50
        file_list_.size() > thread_num ? thread_num : file_list_.size();
Q
Qiao Longfei 已提交
51
    queue_ = queue;
Q
Qiao Longfei 已提交
52
    SplitFiles();
Q
Qiao Longfei 已提交
53
    for (int i = 0; i < thread_num_; ++i) {
Q
Qiao Longfei 已提交
54 55
      read_thread_status_.push_back(Stopped);
    }
Q
Qiao Longfei 已提交
56 57
  }

Q
Qiao Longfei 已提交
58
  ~CTRReader() { Shutdown(); }
Q
Qiao Longfei 已提交
59

Q
Qiao Longfei 已提交
60 61 62 63 64 65
  void ReadNext(std::vector<framework::LoDTensor>* out) override {
    bool success;
    *out = queue_->Pop(&success);
    if (!success) out->clear();
  }

Q
Qiao Longfei 已提交
66 67 68 69 70 71 72 73
  void Shutdown() override {
    VLOG(3) << "Shutdown reader";
    for (auto& read_thread : read_threads_) {
      read_thread->join();
    }
    read_threads_.clear();
    queue_->Close();
  }
Q
Qiao Longfei 已提交
74

Q
Qiao Longfei 已提交
75
  void Start() override {
Q
Qiao Longfei 已提交
76
    VLOG(3) << "Start reader";
Q
Qiao Longfei 已提交
77
    PADDLE_ENFORCE_EQ(read_threads_.size(), 0, "read thread should be empty!");
Q
Qiao Longfei 已提交
78
    queue_->ReOpen();
Q
Qiao Longfei 已提交
79 80 81
    VLOG(3) << "reopen success";
    VLOG(3) << "thread_num " << thread_num_;
    for (int thread_id = 0; thread_id < thread_num_; thread_id++) {
Q
Qiao Longfei 已提交
82 83 84
      read_threads_.emplace_back(new std::thread(
          std::bind(&ReadThread, file_groups_[thread_id], slots_, batch_size_,
                    thread_id, &read_thread_status_, queue_)));
Q
Qiao Longfei 已提交
85
    }
Q
Qiao Longfei 已提交
86 87 88
  }

 private:
Q
Qiao Longfei 已提交
89
  void SplitFiles() {
Q
Qiao Longfei 已提交
90
    file_groups_.resize(thread_num_);
Q
Qiao Longfei 已提交
91
    for (int i = 0; i < file_list_.size(); ++i) {
Q
Qiao Longfei 已提交
92 93 94 95
      auto& file_name = file_list_[i];
      std::ifstream f(file_name.c_str());
      PADDLE_ENFORCE(f.good(), "file %s not exist!", file_name);
      file_groups_[i % thread_num_].push_back(file_name);
Q
Qiao Longfei 已提交
96 97
    }
  }
Q
Qiao Longfei 已提交
98 99

 private:
Q
Qiao Longfei 已提交
100
  int thread_num_;
Q
Qiao Longfei 已提交
101 102 103
  const int batch_size_;
  const std::vector<std::string> slots_;
  const std::vector<std::string> file_list_;
Q
Qiao Longfei 已提交
104
  std::shared_ptr<LoDTensorBlockingQueue> queue_;
Q
Qiao Longfei 已提交
105
  std::vector<std::unique_ptr<std::thread>> read_threads_;
Q
Qiao Longfei 已提交
106
  std::vector<ReaderThreadStatus> read_thread_status_;
Q
Qiao Longfei 已提交
107
  std::vector<std::vector<std::string>> file_groups_;
Q
Qiao Longfei 已提交
108 109 110 111 112
};

}  // namespace reader
}  // namespace operators
}  // namespace paddle