21 #include "caffe2/core/db.h" 22 #include "caffe2/core/init.h" 23 #include "caffe2/core/timer.h" 24 #include "caffe2/core/logging.h" 26 CAFFE2_DEFINE_string(input_db,
"",
"The input db.");
27 CAFFE2_DEFINE_string(input_db_type,
"",
"The input db type.");
28 CAFFE2_DEFINE_int(report_interval, 1000,
"The report interval.");
29 CAFFE2_DEFINE_int(repeat, 10,
"The number to repeat the throughput test.");
30 CAFFE2_DEFINE_bool(use_reader,
false,
"If true, use the reader interface.");
31 CAFFE2_DEFINE_int(num_read_threads, 1,
32 "The number of concurrent reading threads.");
39 void TestThroughputWithDB() {
40 std::unique_ptr<DB> in_db(caffe2::db::CreateDB(
41 caffe2::FLAGS_input_db_type, caffe2::FLAGS_input_db, caffe2::db::READ));
42 std::unique_ptr<Cursor> cursor(in_db->NewCursor());
43 for (
int iter_id = 0; iter_id < caffe2::FLAGS_repeat; ++iter_id) {
45 for (
int i = 0; i < caffe2::FLAGS_report_interval; ++i) {
46 string key = cursor->key();
47 string value = cursor->value();
50 if (!cursor->Valid()) {
51 cursor->SeekToFirst();
54 double elapsed_seconds = timer.
Seconds();
55 printf(
"Iteration %03d, took %4.5f seconds, throughput %f items/sec.\n",
56 iter_id, elapsed_seconds,
57 caffe2::FLAGS_report_interval / elapsed_seconds);
61 void TestThroughputWithReaderWorker(
const DBReader* reader,
int thread_id) {
63 for (
int iter_id = 0; iter_id < caffe2::FLAGS_repeat; ++iter_id) {
65 for (
int i = 0; i < caffe2::FLAGS_report_interval; ++i) {
66 reader->
Read(&key, &value);
68 double elapsed_seconds = timer.
Seconds();
69 printf(
"Thread %03d iteration %03d, took %4.5f seconds, " 70 "throughput %f items/sec.\n",
71 thread_id, iter_id, elapsed_seconds,
72 caffe2::FLAGS_report_interval / elapsed_seconds);
76 void TestThroughputWithReader() {
78 caffe2::FLAGS_input_db_type, caffe2::FLAGS_input_db);
79 std::vector<std::unique_ptr<std::thread>> reading_threads(
80 caffe2::FLAGS_num_read_threads);
81 for (
int i = 0; i < reading_threads.size(); ++i) {
82 reading_threads[i].reset(
new std::thread(
83 TestThroughputWithReaderWorker, &reader, i));
85 for (
int i = 0; i < reading_threads.size(); ++i) {
86 reading_threads[i]->join();
90 int main(
int argc,
char** argv) {
92 if (caffe2::FLAGS_use_reader) {
93 TestThroughputWithReader();
95 TestThroughputWithDB();
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
void Read(string *key, string *value) const
Read a set of key and value from the db and move to next.
An abstract class for the cursor of the database while reading.
A reader wrapper for DB that also allows us to serialize it.
float Seconds()
Returns the elapsed time in seconds.
An abstract class for accessing a database of key-value pairs.
A simple timer object for measuring time.