Caffe2 - C++ API
A deep learning, cross platform ML framework
predictor_config.cc
1 #include "predictor_config.h"
2 
3 #include <atomic>
4 
5 #include "caffe2/core/init.h"
6 #include "caffe2/utils/proto_utils.h"
7 #ifdef CAFFE2_OPTIMIZER
8 #include "caffe2/opt/optimizer.h"
9 #endif
10 
11 namespace caffe2 {
12 
13 namespace {
14 
15 // We don't use the getNet() from predictor_utils.cc here because that file
16 // has additional dependencies that we want to avoid bringing in, to keep the
17 // binary size as small as possible.
18 const NetDef& getNet(const MetaNetDef& def, const std::string& name) {
19  for (const auto& n : def.nets()) {
20  if (n.key() == name) {
21  return n.value();
22  }
23  }
24  CAFFE_THROW("Net not found: ", name);
25 }
26 
27 const ::google::protobuf::RepeatedPtrField<::std::string>& getBlobs(
28  const MetaNetDef& def,
29  const std::string& name) {
30  for (const auto& b : def.blobs()) {
31  if (b.key() == name) {
32  return b.value();
33  }
34  }
35  CAFFE_THROW("Blob not found: ", name);
36 }
37 
38 } // namespace
39 
40 PredictorConfig
41 makePredictorConfig(const MetaNetDef& def, Workspace* parent, bool run_init) {
42  const auto& init_net =
43  getNet(def, PredictorConsts::default_instance().global_init_net_type());
44  const auto& run_net =
45  getNet(def, PredictorConsts::default_instance().predict_net_type());
46  auto config = makePredictorConfig(init_net, run_net, parent, run_init);
47  const auto& inputs =
48  getBlobs(def, PredictorConsts::default_instance().inputs_blob_type());
49  for (const auto& input : inputs) {
50  config.input_names.emplace_back(input);
51  }
52 
53  const auto& outputs =
54  getBlobs(def, PredictorConsts::default_instance().outputs_blob_type());
55  for (const auto& output : outputs) {
56  config.output_names.emplace_back(output);
57  }
58  return config;
59 }
60 
61 PredictorConfig makePredictorConfig(
62  const NetDef& init_net,
63  const NetDef& run_net,
64  Workspace* parent,
65  bool run_init,
66  int optimization) {
67  PredictorConfig config;
68  config.ws = std::make_shared<Workspace>(parent);
69  auto& ws = *config.ws;
70  config.predict_net = std::make_shared<NetDef>(run_net);
71  if (run_init) {
72  CAFFE_ENFORCE(ws.RunNetOnce(init_net));
73  }
74 #ifdef C10_MOBILE
75  GlobalInit();
76 #endif
77  if (optimization &&
78  !ArgumentHelper::HasArgument(*config.predict_net, "disable_nomnigraph")) {
79 #ifdef CAFFE2_OPTIMIZER
80  try {
81  *config.predict_net =
82  opt::optimize(*config.predict_net, &ws, optimization);
83  } catch (const std::exception& e) {
84  LOG(WARNING) << "Optimization pass failed: " << e.what();
85  }
86 #else
87  static std::atomic<bool> warningEmitted{};
88  // Emit the log only once.
89  if (!warningEmitted.exchange(true)) {
90  LOG(WARNING) << "Caffe2 is compiled without optimization passes.";
91  }
92 #endif
93  }
94  return config;
95 }
96 
97 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13
bool GlobalInit(int *pargc, char ***pargv)
Initialize the global environment of caffe2.
Definition: init.cc:44