Caffe2 - C++ API
A deep learning, cross platform ML framework
init.h
1 
17 #ifndef CAFFE2_CORE_INIT_H_
18 #define CAFFE2_CORE_INIT_H_
19 
20 #include "caffe2/core/common.h"
21 #include "caffe2/core/flags.h"
22 #include "caffe2/core/logging.h"
23 
24 namespace caffe2 {
25 
26 namespace internal {
28  public:
29  typedef bool (*InitFunction)(int*, char***);
30  // Registry() is defined in .cpp file to make registration work across
31  // multiple shared libraries loaded with RTLD_LOCAL
33 
34  void Register(InitFunction function, bool run_early,
35  const char* description) {
36  if (run_early) {
37  early_init_functions_.emplace_back(function, description);
38  } else {
39  init_functions_.emplace_back(function, description);
40  }
41  }
42 
43  bool RunRegisteredEarlyInitFunctions(int* pargc, char*** pargv) {
44  return RunRegisteredInitFunctionsInternal(
45  early_init_functions_, pargc, pargv);
46  }
47 
48  bool RunRegisteredInitFunctions(int* pargc, char*** pargv) {
49  return RunRegisteredInitFunctionsInternal(init_functions_, pargc, pargv);
50  }
51 
52  private:
53  // Run all registered initialization functions. This has to be called AFTER
54  // all static initialization are finished and main() has started, since we are
55  // using logging.
56  bool RunRegisteredInitFunctionsInternal(
57  vector<std::pair<InitFunction, const char*>>& functions,
58  int* pargc, char*** pargv) {
59  for (const auto& init_pair : functions) {
60  VLOG(1) << "Running init function: " << init_pair.second;
61  if (!(*init_pair.first)(pargc, pargv)) {
62  LOG(ERROR) << "Initialization function failed.";
63  return false;
64  }
65  }
66  return true;
67  }
68 
70  vector<std::pair<InitFunction, const char*> > early_init_functions_;
71  vector<std::pair<InitFunction, const char*> > init_functions_;
72 };
73 } // namespace internal
74 
76  public:
77  InitRegisterer(internal::Caffe2InitializeRegistry::InitFunction function,
78  bool run_early, const char* description) {
79  internal::Caffe2InitializeRegistry::Registry()
80  ->Register(function, run_early, description);
81  }
82 };
83 
84 #define REGISTER_CAFFE2_INIT_FUNCTION(name, function, description) \
85  namespace { \
86  ::caffe2::InitRegisterer g_caffe2_initregisterer_##name( \
87  function, false, description); \
88  } // namespace
89 
90 #define REGISTER_CAFFE2_EARLY_INIT_FUNCTION(name, function, description) \
91  namespace { \
92  ::caffe2::InitRegisterer g_caffe2_initregisterer_##name( \
93  function, true, description); \
94  } // namespace
95 
113 bool GlobalInit(int* pargc, char*** argv);
114 
115 } // namespace caffe2
116 #endif // CAFFE2_CORE_INIT_H_
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
Definition: init.cc:34
Copyright (c) 2016-present, Facebook, Inc.
Copyright (c) 2016-present, Facebook, Inc.
A template class that allows one to register classes by keys.
Definition: registry.h:57