mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
6706e9af19
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17481 Usually, feature macros are either defined or undefined and checked accordingly. C10_MOBILE was a weird special case that was always defined but either defined to 1 or to 0. This caused a lot of confusion for me when trying to disable something from mobile build and it also disabled it from the server build (because I was using ifdef). Also, I found a place in the existing code base that made that wrong assumption and used the macro wrongly, see https://fburl.com/y4icohts Reviewed By: dzhulgakov Differential Revision: D14214825 fbshipit-source-id: f3a155b6d43d334e8839e2b2e3c40ed2c773eab6
98 lines
2.6 KiB
C++
98 lines
2.6 KiB
C++
#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
|