parallel data load

Author: js1010

cpp/include/ioutils.hpp

@@ -32,16 +32,18 @@ class IoUtils {
   ~IoUtils();
   bool Init(std::string opt_path);
   int LoadStreamFile(std::string filepath);
-  std::pair<int, bool> ReadStreamForVocab(int num_lines);
+  std::pair<int, int> ReadStreamForVocab(int num_lines, int num_threads);
   void GetWordVocab(int min_count);
  private:
   void ParseLine(std::string line, std::vector<std::string>& line_vec);
 
+  std::mutex global_lock_;
   std::ifstream stream_fin_;
   json11::Json opt_;
   std::shared_ptr<spdlog::logger> logger_;
   std::unordered_map<std::string, int> word_idmap_, word_count_;
   std::vector<std::string> word_list_;
+  int num_lines_, remain_lines_;
 };  // class IoUtils
 
 } // namespace cusim

cpp/src/ioutils.cc

@@ -56,28 +56,45 @@ int IoUtils::LoadStreamFile(std::string filepath) {
   word_idmap_.clear();
   word_list_.clear();
   word_count_.clear();
+  num_lines_ = count;
+  remain_lines_ = num_lines_;
   return count;
 }
 
-std::pair<int, bool> IoUtils::ReadStreamForVocab(int num_lines) {
-  int read_cnt = 0;
-  std::string line;
-  std::vector<std::string> line_vec;
-  while (not stream_fin_.eof() and read_cnt < num_lines) {
-    getline(stream_fin_, line);
-    ParseLine(line, line_vec);
-    for (auto& word: line_vec) {
-      if (not word_count_.count(word)) word_count_[word] = 0;
-      word_count_[word]++;
+std::pair<int, int> IoUtils::ReadStreamForVocab(int num_lines, int num_threads) {
+  int read_lines = std::min(num_lines, remain_lines_);
+  remain_lines_ -= read_lines;
+  #pragma omp parallel num_threads(num_threads)
+  {
+    std::string line;
+    std::vector<std::string> line_vec;
+    std::unordered_map<std::string, int> word_count;
+    #pragma omp for schedule(dynamic, 4)
+    for (int i = 0; i < read_lines; ++i) {
+      // get line thread-safely
+      {
+        std::unique_lock<std::mutex> lock(global_lock_);
+        getline(stream_fin_, line);
+      }
+
+      // seems to bottle-neck
+      ParseLine(line, line_vec);
+
+      // update private word count
+      for (auto& word: line_vec) {
+        word_count[word]++;
+      }
+    }
+
+    // update word count to class variable
+    {
+      std::unique_lock<std::mutex> lock(global_lock_);
+      for (auto& it: word_count) {
+        word_count_[it.first] += it.second;
+      }
     }
-    read_cnt++;
-  }
-  bool finished = false;
-  if (stream_fin_.eof()) {
-    stream_fin_.close();
-    finished = true;
   }
-  return {read_cnt, finished};
+  return {read_lines, remain_lines_};
 }
 
 void IoUtils::GetWordVocab(int min_count) {

cusim/ioutils/bindings.cc

@@ -27,8 +27,8 @@ class IoUtilsBind {
     return obj_.LoadStreamFile(filepath);
   }
 
-  std::pair<int, bool> ReadStreamForVocab(int num_lines) {
-    return obj_.ReadStreamForVocab(num_lines);
+  std::pair<int, int> ReadStreamForVocab(int num_lines, int num_threads) {
+    return obj_.ReadStreamForVocab(num_lines, num_threads);
   }
 
   void GetWordVocab(int min_count) {
@@ -46,7 +46,8 @@ PYBIND11_PLUGIN(ioutils_bind) {
   .def(py::init())
   .def("init", &IoUtilsBind::Init, py::arg("opt_path"))
   .def("load_stream_file", &IoUtilsBind::LoadStreamFile, py::arg("filepath"))
-  .def("read_stream_for_vocab", &IoUtilsBind::ReadStreamForVocab, py::arg("num_lines"))
+  .def("read_stream_for_vocab", &IoUtilsBind::ReadStreamForVocab,
+      py::arg("num_lines"), py::arg("num_threads"))
   .def("get_word_vocab", &IoUtilsBind::GetWordVocab, py::arg("min_count"))
   .def("__repr__",
   [](const IoUtilsBind &a) {

cusim/ioutils/pyioutils.py

@@ -29,13 +29,15 @@ def __init__(self, opt=None):
     assert self.obj.init(bytes(tmp.name, "utf8")), f"failed to load {tmp.name}"
     os.remove(tmp.name)
 
-  def load_stream_vocab(self, filepath, min_count, chunk_lines=100000):
+  def load_stream_vocab(self, filepath, min_count,
+                        chunk_lines=100000, num_threads=4):
     full_num_lines = self.obj.load_stream_file(filepath)
     pbar = tqdm.trange(full_num_lines)
     while True:
-      num_lines, finished = self.obj.read_stream_for_vocab(chunk_lines)
-      pbar.update(num_lines)
-      if finished:
+      read_lines, remain_lines = \
+        self.obj.read_stream_for_vocab(chunk_lines, num_threads)
+      pbar.update(read_lines)
+      if not remain_lines:
         break
     pbar.close()
     self.obj.get_word_vocab(min_count)

examples/example1.py

@@ -14,7 +14,8 @@
 
 LOGGER = aux.get_logger()
 DOWNLOAD_PATH = "./res"
-DATASET = "wiki-english-20171001"
+# DATASET = "wiki-english-20171001"
+DATASET = "fake-news"
 DATA_PATH = f"./res/{DATASET}.stream.txt"
 MIN_COUNT = 5
 
@@ -33,7 +34,7 @@ def download():
 def run():
   download()
   iou = IoUtils()
-  iou.load_stream_vocab(DATA_PATH, 5, 10000)
+  iou.load_stream_vocab(DATA_PATH, 5, 100000, 8)
 
 
 if __name__ == "__main__":