Caffe2 - C++ API
A deep learning, cross platform ML framework
registry.h
1 
21 #ifndef CAFFE2_CORE_REGISTRY_H_
22 #define CAFFE2_CORE_REGISTRY_H_
23 
24 #include <algorithm>
25 #include <cstdio>
26 #include <cstdlib>
27 #include <functional>
28 #include <memory>
29 #include <mutex>
30 
31 #include "caffe2/core/common.h"
32 #include "caffe2/core/typeid.h"
33 
34 namespace caffe2 {
35 
36 template <typename KeyType>
37 inline void PrintOffendingKey(const KeyType& key) {
38  printf("[key type printing not supported]\n");
39 }
40 
41 template <>
42 inline void PrintOffendingKey(const string& key) {
43  printf("Offending key: %s.\n", key.c_str());
44 }
45 
56 template <class SrcType, class ObjectPtrType, class... Args>
57 class Registry {
58  public:
59  typedef std::function<ObjectPtrType(Args...)> Creator;
60 
61  Registry() : registry_() {}
62 
63  void Register(const SrcType& key, Creator creator) {
64  // The if statement below is essentially the same as the following line:
65  // CHECK_EQ(registry_.count(key), 0) << "Key " << key
66  // << " registered twice.";
67  // However, CHECK_EQ depends on google logging, and since registration is
68  // carried out at static initialization time, we do not want to have an
69  // explicit dependency on glog's initialization function.
70  std::lock_guard<std::mutex> lock(register_mutex_);
71  if (registry_.count(key) != 0) {
72  printf("Key already registered.\n");
73  PrintOffendingKey(key);
74  std::exit(1);
75  }
76  registry_[key] = creator;
77  }
78 
79  void Register(const SrcType& key, Creator creator, const string& help_msg) {
80  Register(key, creator);
81  help_message_[key] = help_msg;
82  }
83 
84  inline bool Has(const SrcType& key) { return (registry_.count(key) != 0); }
85 
86  ObjectPtrType Create(const SrcType& key, Args... args) {
87  if (registry_.count(key) == 0) {
88  // Returns nullptr if the key is not registered.
89  return nullptr;
90  }
91  return registry_[key](args...);
92  }
93 
97  vector<SrcType> Keys() {
98  vector<SrcType> keys;
99  for (const auto& it : registry_) {
100  keys.push_back(it.first);
101  }
102  return keys;
103  }
104 
105  const CaffeMap<SrcType, string>& HelpMessage() const {
106  return help_message_;
107  }
108 
109  const char* HelpMessage(const SrcType& key) const {
110  auto it = help_message_.find(key);
111  if (it == help_message_.end()) {
112  return nullptr;
113  }
114  return it->second.c_str();
115  }
116 
117  private:
118  CaffeMap<SrcType, Creator> registry_;
119  CaffeMap<SrcType, string> help_message_;
120  std::mutex register_mutex_;
121 
122  DISABLE_COPY_AND_ASSIGN(Registry);
123 };
124 
125 template <class SrcType, class ObjectPtrType, class... Args>
126 class Registerer {
127  public:
128  Registerer(
129  const SrcType& key,
131  typename Registry<SrcType, ObjectPtrType, Args...>::Creator creator,
132  const string& help_msg = "") {
133  registry->Register(key, creator, help_msg);
134  }
135 
136  template <class DerivedType>
137  static ObjectPtrType DefaultCreator(Args... args) {
138  // TODO(jiayq): old versions of NVCC does not handle make_unique well
139  // so we are forced to use a unique_ptr constructor here. Check if it is
140  // fine to use make_unique in the future.
141  // return make_unique<DerivedType>(args...);
142  return ObjectPtrType(new DerivedType(args...));
143  }
144 };
145 
151 #define CAFFE_CONCATENATE_IMPL(s1, s2) s1##s2
152 #define CAFFE_CONCATENATE(s1, s2) CAFFE_CONCATENATE_IMPL(s1, s2)
153 #ifdef __COUNTER__
154 #define CAFFE_ANONYMOUS_VARIABLE(str) CAFFE_CONCATENATE(str, __COUNTER__)
155 #else
156 #define CAFFE_ANONYMOUS_VARIABLE(str) CAFFE_CONCATENATE(str, __LINE__)
157 #endif
158 
164 #define CAFFE_DECLARE_TYPED_REGISTRY( \
165  RegistryName, SrcType, ObjectType, PtrType, ...) \
166  Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>* RegistryName(); \
167  typedef Registerer<SrcType, PtrType<ObjectType>, ##__VA_ARGS__> \
168  Registerer##RegistryName;
169 
170 #define CAFFE_DEFINE_TYPED_REGISTRY( \
171  RegistryName, SrcType, ObjectType, PtrType, ...) \
172  Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>* RegistryName() { \
173  static Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>* registry = \
174  new Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>(); \
175  return registry; \
176  }
177 
178 // Note(Yangqing): The __VA_ARGS__ below allows one to specify a templated
179 // creator with comma in its templated arguments.
180 #define CAFFE_REGISTER_TYPED_CREATOR(RegistryName, key, ...) \
181  namespace { \
182  static Registerer##RegistryName CAFFE_ANONYMOUS_VARIABLE(g_##RegistryName)( \
183  key, RegistryName(), __VA_ARGS__); \
184  }
185 
186 #define CAFFE_REGISTER_TYPED_CLASS(RegistryName, key, ...) \
187  namespace { \
188  static Registerer##RegistryName CAFFE_ANONYMOUS_VARIABLE(g_##RegistryName)( \
189  key, \
190  RegistryName(), \
191  Registerer##RegistryName::DefaultCreator<__VA_ARGS__>, \
192  DemangleType<__VA_ARGS__>()); \
193  }
194 
195 // CAFFE_DECLARE_REGISTRY and CAFFE_DEFINE_REGISTRY are hard-wired to use string
196 // as the key
197 // type, because that is the most commonly used cases.
198 #define CAFFE_DECLARE_REGISTRY(RegistryName, ObjectType, ...) \
199  CAFFE_DECLARE_TYPED_REGISTRY( \
200  RegistryName, std::string, ObjectType, std::unique_ptr, ##__VA_ARGS__)
201 
202 #define CAFFE_DEFINE_REGISTRY(RegistryName, ObjectType, ...) \
203  CAFFE_DEFINE_TYPED_REGISTRY( \
204  RegistryName, std::string, ObjectType, std::unique_ptr, ##__VA_ARGS__)
205 
206 #define CAFFE_DECLARE_SHARED_REGISTRY(RegistryName, ObjectType, ...) \
207  CAFFE_DECLARE_TYPED_REGISTRY( \
208  RegistryName, std::string, ObjectType, std::shared_ptr, ##__VA_ARGS__)
209 
210 #define CAFFE_DEFINE_SHARED_REGISTRY(RegistryName, ObjectType, ...) \
211  CAFFE_DEFINE_TYPED_REGISTRY( \
212  RegistryName, std::string, ObjectType, std::shared_ptr, ##__VA_ARGS__)
213 
214 // CAFFE_REGISTER_CREATOR and CAFFE_REGISTER_CLASS are hard-wired to use string
215 // as the key
216 // type, because that is the most commonly used cases.
217 #define CAFFE_REGISTER_CREATOR(RegistryName, key, ...) \
218  CAFFE_REGISTER_TYPED_CREATOR(RegistryName, #key, __VA_ARGS__)
219 
220 #define CAFFE_REGISTER_CLASS(RegistryName, key, ...) \
221  CAFFE_REGISTER_TYPED_CLASS(RegistryName, #key, __VA_ARGS__)
222 
223 } // namespace caffe2
224 #endif // CAFFE2_CORE_REGISTRY_H_
vector< SrcType > Keys()
Returns the keys currently registered as a vector.
Definition: registry.h:97
Copyright (c) 2016-present, Facebook, Inc.
A template class that allows one to register classes by keys.
Definition: registry.h:57