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