data_set.cc 6.5 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/* 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. */

15
#include <random>
16
#include "paddle/fluid/framework/data_set.h"
17 18 19
#include "google/protobuf/io/zero_copy_stream_impl.h"
#include "google/protobuf/message.h"
#include "google/protobuf/text_format.h"
20 21 22 23 24
#include "paddle/fluid/framework/data_feed_factory.h"

namespace paddle {
namespace framework {

25 26
template <typename T>
DatasetImpl<T>::DatasetImpl() { thread_num_ = 1; }
27

28 29
template <typename T>
void DatasetImpl<T>::SetFileList(const std::vector<std::string>& filelist) {
30
  VLOG(3) << "filelist size: " << filelist.size();
31
  filelist_ = filelist;
32
  /*
33 34
  int file_cnt = filelist_.size();
  if (thread_num_ > file_cnt) {
D
dongdaxiang 已提交
35 36
    VLOG(1) << "DataSet thread num = " << thread_num_
            << ", file num = " << file_cnt
37 38
            << ". Changing DataSet thread num = " << file_cnt;
    thread_num_ = file_cnt;
39
  }*/
40 41
}

42 43
// buggy here, a user should set filelist first before this function
// not user friendly
44 45
template <typename T>
void DatasetImpl<T>::SetThreadNum(int thread_num) {
46 47
  int file_cnt = filelist_.size();
  if (file_cnt != 0 && thread_num > file_cnt) {
D
dongdaxiang 已提交
48 49
    VLOG(1) << "DataSet thread num = " << thread_num
            << ", file num = " << file_cnt
50 51 52 53 54 55
            << ". Changing DataSet thread num = " << file_cnt;
    thread_num = file_cnt;
  }
  thread_num_ = thread_num;
}

56
template <typename T>
X
xujiaqi01 已提交
57 58 59
void DatasetImpl<T>::SetTrainerNum(int trainer_num) {
  trainer_num_ = trainer_num;
}
60

61 62
template <typename T>
void DatasetImpl<T>::SetDataFeedDesc(const std::string& data_feed_desc_str) {
63 64
  google::protobuf::TextFormat::ParseFromString(data_feed_desc_str,
                                                &data_feed_desc_);
65 66
}

67 68 69
template <typename T>
std::vector<std::shared_ptr<paddle::framework::DataFeed>>&
    DatasetImpl<T>::GetReaders() {
70 71 72
  return readers_;
}

73 74 75
template <typename T>
void DatasetImpl<T>::LoadIntoMemory() {
  VLOG(3) << "DatasetImpl<T>::LoadIntoMemory() begin";
76 77 78 79 80
  if (readers_.size() == 0) {
    CreateReaders();
  }
  std::vector<std::thread> load_threads;
  for (int64_t i = 0; i < thread_num_; ++i) {
D
dongdaxiang 已提交
81 82
    load_threads.push_back(std::thread(
        &paddle::framework::DataFeed::LoadIntoMemory, readers_[i].get()));
83 84 85 86
  }
  for (std::thread& t : load_threads) {
    t.join();
  }
87
  VLOG(3) << "DatasetImpl<T>::LoadIntoMemory() end";
88 89
}

90 91 92
template <typename T>
void DatasetImpl<T>::LocalShuffle() {
  VLOG(3) << "DatasetImpl<T>::LocalShuffle() begin";
93 94 95
  if (readers_.size() == 0) {
    CreateReaders();
  }
96 97 98
  // if it is not InMemory, memory_data_ is empty
  std::random_shuffle(memory_data_.begin(), memory_data_.end());

99 100
  std::vector<std::thread> local_shuffle_threads;
  for (int64_t i = 0; i < thread_num_; ++i) {
D
dongdaxiang 已提交
101 102
    local_shuffle_threads.push_back(std::thread(
        &paddle::framework::DataFeed::LocalShuffle, readers_[i].get()));
103 104 105 106
  }
  for (std::thread& t : local_shuffle_threads) {
    t.join();
  }
107 108
  std::vector<T>().swap(memory_data_);
  VLOG(3) << "DatasetImpl<T>::LocalShuffle() end";
109 110
}

111 112 113 114 115 116 117 118
template <typename T>
void DatasetImpl<T>::GlobalShuffle() {
  VLOG(3) << "DatasetImpl<T>::GlobalShuffle() begin";
  if (readers_.size() == 0) {
      CreateReaders();
  }
  // if it is not InMemory, memory_data_ is empty
  std::random_shuffle(memory_data_.begin(), memory_data_.end());
119
  auto fleet_ptr = FleetWrapper::GetInstance();
120 121
  VLOG(3) << "RegisterClientToClientMsgHandler";
  fleet_ptr->RegisterClientToClientMsgHandler(0,
122 123 124
    [this](int msg_type, int client_id, const std::string& msg) -> int {
    return this->ReceiveFromClient(msg_type, client_id, msg);
  });
X
xujiaqi01 已提交
125
  VLOG(3) << "start global shuffle threads";
126
  std::vector<std::thread> global_shuffle_threads;
127 128 129 130
  for (int i = 0; i < thread_num_; ++i) {
    global_shuffle_threads.push_back(
        std::thread(&paddle::framework::DataFeed::GlobalShuffle,
        readers_[i].get()));
131 132 133
  }
  for (std::thread& t : global_shuffle_threads) {
    t.join();
134 135
  }
  VLOG(3) << "DatasetImpl<T>::GlobalShuffle() end";
136 137
}

138 139
template <typename T>
void DatasetImpl<T>::CreateReaders() {
140
  VLOG(3) << "Calling CreateReaders()";
141
  CHECK(thread_num_ > 0) << "thread_num should > 0";
142 143
  VLOG(3) << "thread_num in Readers: " << thread_num_;
  VLOG(3) << "readers size: " << readers_.size();
144 145 146
  if (readers_.size() != 0) {
    return;
  }
147
  VLOG(3) << "data feed class name: " << data_feed_desc_.name();
148
  for (int i = 0; i < thread_num_; ++i) {
149 150
    readers_.push_back(DataFeedFactory::CreateDataFeed(data_feed_desc_.name()));
    readers_.back()->Init(data_feed_desc_);
151 152 153 154 155
    readers_.back()->SetMemoryData(&memory_data_);
    readers_.back()->SetMemoryDataMutex(&mutex_for_update_memory_data_);
    readers_.back()->SetThreadId(i);
    readers_.back()->SetThreadNum(thread_num_);
    readers_.back()->SetTrainerNum(trainer_num_);
156
  }
157
  VLOG(3) << "Filelist size in readers: " << filelist_.size();
158 159 160
  readers_[0]->SetFileList(filelist_);
}

161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184
template <typename T>
void DatasetImpl<T>::DestroyReaders() {
  VLOG(3) << "Calling DestroyReaders()";
  // clear memory_data_ before fill it
  // because if LoadIntoMemory but no Shuffle,
  // memory_data_ has empty data which has been std::move to channel
  if (memory_data_.size() != 0) {
    std::vector<T>().swap(memory_data_);
  }
  std::vector<std::thread> fill_threads;
  for (int i = 0; i < thread_num_; ++i) {
    fill_threads.push_back(std::thread(
        &paddle::framework::DataFeed::FillChannelToMemoryData,
        readers_[i].get()));
  }
  for (std::thread& t : fill_threads) {
    t.join();
  }
  std::vector<std::string>().swap(filelist_);
  std::vector<std::shared_ptr<paddle::framework::DataFeed>>().swap(readers_);
}

template <typename T>
int DatasetImpl<T>::ReceiveFromClient(int msg_type, int client_id,
D
dongdaxiang 已提交
185
                               const std::string& msg) {
186
  // todo random
187 188 189 190 191 192
  // int64_t index = paddle::ps::local_random_engine()() % thread_num_;
  int64_t index = 0;
  readers_[index]->PutInsToChannel(msg);
  return 0;
}

193 194 195
// explicit instantiation
template class DatasetImpl<std::vector<MultiSlotType>>;

D
dongdaxiang 已提交
196 197
}  // end namespace framework
}  // end namespace paddle