forked from liulei01/DRBox
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdata_reader.cpp
More file actions
144 lines (128 loc) · 4.14 KB
/
data_reader.cpp
File metadata and controls
144 lines (128 loc) · 4.14 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
#include <boost/thread.hpp>
#include <map>
#include <string>
#include <vector>
#include "caffe/common.hpp"
#include "caffe/data_reader.hpp"
#include "caffe/layers/annotated_r_data_layer.hpp"
#include "caffe/layers/annotated_data_layer.hpp"
#include "caffe/layers/data_layer.hpp"
#include "caffe/proto/caffe.pb.h"
namespace caffe {
using boost::weak_ptr;
// It has to explicitly initialize the map<> in order to work. It seems to be a
// gcc bug.
// http://www.cplusplus.com/forum/beginner/31576/
template <>
map<const string, weak_ptr<DataReader<Datum>::Body> >
DataReader<Datum>::bodies_
= map<const string, weak_ptr<DataReader<Datum>::Body> >();
template <>
map<const string, weak_ptr<DataReader<AnnotatedDatum>::Body> >
DataReader<AnnotatedDatum>::bodies_
= map<const string, weak_ptr<DataReader<AnnotatedDatum>::Body> >();
template <>
map<const string, weak_ptr<DataReader<AnnotatedDatumR>::Body> >
DataReader<AnnotatedDatumR>::bodies_
= map<const string, weak_ptr<DataReader<AnnotatedDatumR>::Body> >();
static boost::mutex bodies_mutex_;
template <typename T>
DataReader<T>::DataReader(const LayerParameter& param)
: queue_pair_(new QueuePair( //
param.data_param().prefetch() * param.data_param().batch_size())) {
// Get or create a body
boost::mutex::scoped_lock lock(bodies_mutex_);
string key = source_key(param);
weak_ptr<Body>& weak = bodies_[key];
body_ = weak.lock();
if (!body_) {
body_.reset(new Body(param));
bodies_[key] = weak_ptr<Body>(body_);
}
body_->new_queue_pairs_.push(queue_pair_);
}
template <typename T>
DataReader<T>::~DataReader() {
string key = source_key(body_->param_);
body_.reset();
boost::mutex::scoped_lock lock(bodies_mutex_);
if (bodies_[key].expired()) {
bodies_.erase(key);
}
}
template <typename T>
DataReader<T>::QueuePair::QueuePair(int size) {
// Initialize the free queue with requested number of data
for (int i = 0; i < size; ++i) {
free_.push(new T());
}
}
template <typename T>
DataReader<T>::QueuePair::~QueuePair() {
T* t;
while (free_.try_pop(&t)) {
delete t;
}
while (full_.try_pop(&t)) {
delete t;
}
}
template <typename T>
DataReader<T>::Body::Body(const LayerParameter& param)
: param_(param),
new_queue_pairs_() {
StartInternalThread();
}
template <typename T>
DataReader<T>::Body::~Body() {
StopInternalThread();
}
template <typename T>
void DataReader<T>::Body::InternalThreadEntry() {
shared_ptr<db::DB> db(db::GetDB(param_.data_param().backend()));
db->Open(param_.data_param().source(), db::READ);
shared_ptr<db::Cursor> cursor(db->NewCursor());
vector<shared_ptr<QueuePair> > qps;
try {
int solver_count = param_.phase() == TRAIN ? Caffe::solver_count() : 1;
// To ensure deterministic runs, only start running once all solvers
// are ready. But solvers need to peek on one item during initialization,
// so read one item, then wait for the next solver.
for (int i = 0; i < solver_count; ++i) {
shared_ptr<QueuePair> qp(new_queue_pairs_.pop());
read_one(cursor.get(), qp.get());
qps.push_back(qp);
}
// Main loop
while (!must_stop()) {
for (int i = 0; i < solver_count; ++i) {
read_one(cursor.get(), qps[i].get());
}
// Check no additional readers have been created. This can happen if
// more than one net is trained at a time per process, whether single
// or multi solver. It might also happen if two data layers have same
// name and same source.
CHECK_EQ(new_queue_pairs_.size(), 0);
}
} catch (boost::thread_interrupted&) {
// Interrupted exception is expected on shutdown
}
}
template <typename T>
void DataReader<T>::Body::read_one(db::Cursor* cursor, QueuePair* qp) {
T* t = qp->free_.pop();
// TODO deserialize in-place instead of copy?
t->ParseFromString(cursor->value());
qp->full_.push(t);
// go to the next iter
cursor->Next();
if (!cursor->valid()) {
DLOG(INFO) << "Restarting data prefetching from start.";
cursor->SeekToFirst();
}
}
// Instance class
template class DataReader<Datum>;
template class DataReader<AnnotatedDatum>;
template class DataReader<AnnotatedDatumR>;
} // namespace caffe