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