5 #ifndef CAFFE2_CORE_REGISTRY_H_ 6 #define CAFFE2_CORE_REGISTRY_H_ 15 #include "caffe2/core/common.h" 16 #include "caffe2/core/typeid.h" 20 template <
typename KeyType>
21 inline void PrintOffendingKey(
const KeyType& key) {
22 printf(
"[key type printing not supported]\n");
26 inline void PrintOffendingKey(
const string& key) {
27 printf(
"Offending key: %s.\n", key.c_str());
40 template <
class SrcType,
class ObjectPtrType,
class... Args>
43 typedef std::function<ObjectPtrType(Args...)> Creator;
47 void Register(
const SrcType& key, Creator creator) {
54 std::lock_guard<std::mutex> lock(register_mutex_);
55 if (registry_.count(key) != 0) {
56 printf(
"Key already registered.\n");
57 PrintOffendingKey(key);
60 registry_[key] = creator;
63 void Register(
const SrcType& key, Creator creator,
const string& help_msg) {
64 Register(key, creator);
65 help_message_[key] = help_msg;
68 inline bool Has(
const SrcType& key) {
return (registry_.count(key) != 0); }
70 ObjectPtrType Create(
const SrcType& key, Args... args) {
71 if (registry_.count(key) == 0) {
75 return registry_[key](args...);
83 for (
const auto& it : registry_) {
84 keys.push_back(it.first);
89 const CaffeMap<SrcType, string>& HelpMessage()
const {
93 const char* HelpMessage(
const SrcType& key)
const {
94 auto it = help_message_.find(key);
95 if (it == help_message_.end()) {
98 return it->second.c_str();
102 CaffeMap<SrcType, Creator> registry_;
103 CaffeMap<SrcType, string> help_message_;
104 std::mutex register_mutex_;
109 template <
class SrcType,
class ObjectPtrType,
class... Args>
115 typename Registry<SrcType, ObjectPtrType, Args...>::Creator creator,
116 const string& help_msg =
"") {
117 registry->Register(key, creator, help_msg);
120 template <
class DerivedType>
121 static ObjectPtrType DefaultCreator(Args... args) {
126 return ObjectPtrType(
new DerivedType(args...));
135 #define CAFFE_CONCATENATE_IMPL(s1, s2) s1##s2 136 #define CAFFE_CONCATENATE(s1, s2) CAFFE_CONCATENATE_IMPL(s1, s2) 138 #define CAFFE_ANONYMOUS_VARIABLE(str) CAFFE_CONCATENATE(str, __COUNTER__) 140 #define CAFFE_ANONYMOUS_VARIABLE(str) CAFFE_CONCATENATE(str, __LINE__) 148 #define CAFFE_DECLARE_TYPED_REGISTRY( \ 149 RegistryName, SrcType, ObjectType, PtrType, ...) \ 150 Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>* RegistryName(); \ 151 typedef Registerer<SrcType, PtrType<ObjectType>, ##__VA_ARGS__> \ 152 Registerer##RegistryName; 154 #define CAFFE_DEFINE_TYPED_REGISTRY( \ 155 RegistryName, SrcType, ObjectType, PtrType, ...) \ 156 Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>* RegistryName() { \ 157 static Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>* registry = \ 158 new Registry<SrcType, PtrType<ObjectType>, ##__VA_ARGS__>(); \ 164 #define CAFFE_REGISTER_TYPED_CREATOR(RegistryName, key, ...) \ 166 static Registerer##RegistryName CAFFE_ANONYMOUS_VARIABLE(g_##RegistryName)( \ 167 key, RegistryName(), __VA_ARGS__); \ 170 #define CAFFE_REGISTER_TYPED_CLASS(RegistryName, key, ...) \ 172 static Registerer##RegistryName CAFFE_ANONYMOUS_VARIABLE(g_##RegistryName)( \ 175 Registerer##RegistryName::DefaultCreator<__VA_ARGS__>, \ 176 DemangleType<__VA_ARGS__>()); \ 182 #define CAFFE_DECLARE_REGISTRY(RegistryName, ObjectType, ...) \ 183 CAFFE_DECLARE_TYPED_REGISTRY( \ 184 RegistryName, std::string, ObjectType, std::unique_ptr, ##__VA_ARGS__) 186 #define CAFFE_DEFINE_REGISTRY(RegistryName, ObjectType, ...) \ 187 CAFFE_DEFINE_TYPED_REGISTRY( \ 188 RegistryName, std::string, ObjectType, std::unique_ptr, ##__VA_ARGS__) 190 #define CAFFE_DECLARE_SHARED_REGISTRY(RegistryName, ObjectType, ...) \ 191 CAFFE_DECLARE_TYPED_REGISTRY( \ 192 RegistryName, std::string, ObjectType, std::shared_ptr, ##__VA_ARGS__) 194 #define CAFFE_DEFINE_SHARED_REGISTRY(RegistryName, ObjectType, ...) \ 195 CAFFE_DEFINE_TYPED_REGISTRY( \ 196 RegistryName, std::string, ObjectType, std::shared_ptr, ##__VA_ARGS__) 201 #define CAFFE_REGISTER_CREATOR(RegistryName, key, ...) \ 202 CAFFE_REGISTER_TYPED_CREATOR(RegistryName, #key, __VA_ARGS__) 204 #define CAFFE_REGISTER_CLASS(RegistryName, key, ...) \ 205 CAFFE_REGISTER_TYPED_CLASS(RegistryName, #key, __VA_ARGS__) 208 #endif // CAFFE2_CORE_REGISTRY_H_ vector< SrcType > Keys()
Returns the keys currently registered as a vector.
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
A template class that allows one to register classes by keys.