Caffe2 - C++ API
A deep learning, cross platform ML framework
init.h
1 #ifndef CAFFE2_CORE_INIT_H_
2 #define CAFFE2_CORE_INIT_H_
3 
4 #include "caffe2/core/common.h"
5 #include "caffe2/core/flags.h"
6 #include "caffe2/core/logging.h"
7 
8 namespace caffe2 {
9 
10 namespace internal {
12  public:
13  typedef bool (*InitFunction)(int*, char***);
14  // Registry() is defined in .cpp file to make registration work across
15  // multiple shared libraries loaded with RTLD_LOCAL
17 
18  void Register(InitFunction function, bool run_early,
19  const char* description) {
20  if (run_early) {
21  early_init_functions_.emplace_back(function, description);
22  } else {
23  init_functions_.emplace_back(function, description);
24  }
25  }
26 
27  bool RunRegisteredEarlyInitFunctions(int* pargc, char*** pargv) {
28  return RunRegisteredInitFunctionsInternal(
29  early_init_functions_, pargc, pargv);
30  }
31 
32  bool RunRegisteredInitFunctions(int* pargc, char*** pargv) {
33  return RunRegisteredInitFunctionsInternal(init_functions_, pargc, pargv);
34  }
35 
36  private:
37  // Run all registered initialization functions. This has to be called AFTER
38  // all static initialization are finished and main() has started, since we are
39  // using logging.
40  bool RunRegisteredInitFunctionsInternal(
41  vector<std::pair<InitFunction, const char*>>& functions,
42  int* pargc, char*** pargv) {
43  for (const auto& init_pair : functions) {
44  VLOG(1) << "Running init function: " << init_pair.second;
45  if (!(*init_pair.first)(pargc, pargv)) {
46  LOG(ERROR) << "Initialization function failed.";
47  return false;
48  }
49  }
50  return true;
51  }
52 
54  vector<std::pair<InitFunction, const char*> > early_init_functions_;
55  vector<std::pair<InitFunction, const char*> > init_functions_;
56 };
57 } // namespace internal
58 
60  public:
61  InitRegisterer(internal::Caffe2InitializeRegistry::InitFunction function,
62  bool run_early, const char* description) {
63  internal::Caffe2InitializeRegistry::Registry()
64  ->Register(function, run_early, description);
65  }
66 };
67 
68 #define REGISTER_CAFFE2_INIT_FUNCTION(name, function, description) \
69  namespace { \
70  ::caffe2::InitRegisterer g_caffe2_initregisterer_##name( \
71  function, false, description); \
72  } // namespace
73 
74 #define REGISTER_CAFFE2_EARLY_INIT_FUNCTION(name, function, description) \
75  namespace { \
76  ::caffe2::InitRegisterer g_caffe2_initregisterer_##name( \
77  function, true, description); \
78  } // namespace
79 
97 bool GlobalInit(int* pargc, char*** argv);
98 
99 } // namespace caffe2
100 #endif // CAFFE2_CORE_INIT_H_
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
Definition: init.cc:15
A template class that allows one to register classes by keys.
Definition: registry.h:31
Simple registry implementation in Caffe2 that uses static variables to register object creators durin...
Commandline flags support for Caffe2.