| 247 | } // namespace |
| 248 | |
| 249 | Status InitializeSession(int num_threads, const string& graph, |
| 250 | std::unique_ptr<Session>* session, |
| 251 | std::unique_ptr<GraphDef>* graph_def) { |
| 252 | LOG(INFO) << "Loading TensorFlow."; |
| 253 | |
| 254 | tensorflow::SessionOptions options; |
| 255 | tensorflow::ConfigProto& config = options.config; |
| 256 | if (num_threads > 0) { |
| 257 | config.set_intra_op_parallelism_threads(num_threads); |
| 258 | config.set_inter_op_parallelism_threads(num_threads); |
| 259 | } |
| 260 | LOG(INFO) << "Got config, " << config.device_count_size() << " devices"; |
| 261 | |
| 262 | session->reset(tensorflow::NewSession(options)); |
| 263 | graph_def->reset(new GraphDef()); |
| 264 | tensorflow::GraphDef tensorflow_graph; |
| 265 | Status s = ReadBinaryProto(Env::Default(), graph, graph_def->get()); |
| 266 | if (!s.ok()) { |
| 267 | s = ReadTextProto(Env::Default(), graph, graph_def->get()); |
| 268 | } |
| 269 | |
| 270 | if (!s.ok()) { |
| 271 | LOG(ERROR) << "Could not create TensorFlow Graph: " << s; |
| 272 | return s; |
| 273 | } |
| 274 | |
| 275 | s = (*session)->Create(*(graph_def->get())); |
| 276 | if (!s.ok()) { |
| 277 | LOG(ERROR) << "Could not create TensorFlow Session: " << s; |
| 278 | return s; |
| 279 | } |
| 280 | |
| 281 | return Status::OK(); |
| 282 | } |
| 283 | |
| 284 | Status RunBenchmark(const std::vector<InputLayerInfo>& inputs, |
| 285 | const std::vector<string>& outputs, |