Caffe2 - C++ API
A deep learning, cross platform ML framework
protodb.cc
1 #include <unordered_set>
2 
3 #include "caffe2/core/db.h"
4 #include "caffe2/utils/proto_utils.h"
5 #include "caffe2/core/logging.h"
6 
7 namespace caffe2 {
8 namespace db {
9 
10 class ProtoDBCursor : public Cursor {
11  public:
12  explicit ProtoDBCursor(const TensorProtos* proto)
13  : proto_(proto), iter_(0) {}
14  ~ProtoDBCursor() {}
15 
16  void Seek(const string& str) override {
17  CAFFE_THROW("ProtoDB is not designed to support seeking.");
18  }
19 
20  void SeekToFirst() override { iter_ = 0; }
21  void Next() override { ++iter_; }
22  string key() override { return proto_->protos(iter_).name(); }
23  string value() override { return proto_->protos(iter_).SerializeAsString(); }
24  bool Valid() override { return iter_ < proto_->protos_size(); }
25 
26  private:
27  const TensorProtos* proto_;
28  int iter_;
29 };
30 
32  public:
33  explicit ProtoDBTransaction(TensorProtos* proto)
34  : proto_(proto), existing_names_() {
35  for (const auto& tensor : proto_->protos()) {
36  existing_names_.insert(tensor.name());
37  }
38  }
39  ~ProtoDBTransaction() { Commit(); }
40  void Put(const string& key, const string& value) override {
41  if (existing_names_.count(key)) {
42  CAFFE_THROW("An item with key ", key, " already exists.");
43  }
44  auto* tensor = proto_->add_protos();
45  CAFFE_ENFORCE(
46  tensor->ParseFromString(value),
47  "Cannot parse content from the value string.");
48  CAFFE_ENFORCE(
49  tensor->name() == key,
50  "Passed in key ",
51  key,
52  " does not equal to the tensor name ",
53  tensor->name());
54  }
55  // Commit does nothing. The protocol buffer will be written at destruction
56  // of ProtoDB.
57  void Commit() override {}
58 
59  private:
60  TensorProtos* proto_;
61  std::unordered_set<string> existing_names_;
62 
63  DISABLE_COPY_AND_ASSIGN(ProtoDBTransaction);
64 };
65 
66 class ProtoDB : public DB {
67  public:
68  ProtoDB(const string& source, Mode mode)
69  : DB(source, mode), proto_(), source_(source) {
70  if (mode == READ || mode == WRITE) {
71  // Read the current protobuffer.
72  CAFFE_ENFORCE(
73  ReadProtoFromFile(source, &proto_), "Cannot read protobuffer.");
74  }
75  LOG(INFO) << "Opened protodb " << source;
76  }
77  ~ProtoDB() { Close(); }
78 
79  void Close() override {
80  if (mode_ == NEW || mode_ == WRITE) {
81  WriteProtoToBinaryFile(proto_, source_);
82  }
83  }
84 
85  unique_ptr<Cursor> NewCursor() override {
86  return make_unique<ProtoDBCursor>(&proto_);
87  }
88  unique_ptr<Transaction> NewTransaction() override {
89  return make_unique<ProtoDBTransaction>(&proto_);
90  }
91 
92  private:
93  TensorProtos proto_;
94  string source_;
95 };
96 
97 REGISTER_CAFFE2_DB(ProtoDB, ProtoDB);
98 // For lazy-minded, one can also call with lower-case name.
99 REGISTER_CAFFE2_DB(protodb, ProtoDB);
100 
101 } // namespace db
102 } // namespace caffe2
An abstract class for accessing a database of key-value pairs.
Definition: db.h:80
bool Valid() override
Returns whether the current location is valid - for example, if we have reached the end of the databa...
Definition: protodb.cc:24
string value() override
Returns the current value.
Definition: protodb.cc:23
An abstract class for the cursor of the database while reading.
Definition: db.h:22
void Commit() override
Commits the current writes.
Definition: protodb.cc:57
void Seek(const string &str) override
Seek to a specific key (or if the key does not exist, seek to the immediate next).
Definition: protodb.cc:16
unique_ptr< Transaction > NewTransaction() override
Returns a transaction to write data to the database.
Definition: protodb.cc:88
Simple registry implementation in Caffe2 that uses static variables to register object creators durin...
void Next() override
Go to the next location in the database.
Definition: protodb.cc:21
void SeekToFirst() override
Seek to the first key in the database.
Definition: protodb.cc:20
void Close() override
Closes the database.
Definition: protodb.cc:79
void Put(const string &key, const string &value) override
Puts the key value pair to the database.
Definition: protodb.cc:40
An abstract class for the current database transaction while writing.
Definition: db.h:61
unique_ptr< Cursor > NewCursor() override
Returns a cursor to read the database.
Definition: protodb.cc:85
string key() override
Returns the current key.
Definition: protodb.cc:22