| 344 | } |
| 345 | |
| 346 | Status LoadSavedModelInternal(const SessionOptions& session_options, |
| 347 | const RunOptions& run_options, |
| 348 | const string& export_dir, |
| 349 | const std::unordered_set<string>& tags, |
| 350 | SavedModelBundle* const bundle) { |
| 351 | const uint64 read_start_microseconds = Env::Default()->NowMicros(); |
| 352 | TF_RETURN_IF_ERROR(ReadMetaGraphDefFromSavedModel(export_dir, tags, |
| 353 | &bundle->meta_graph_def)); |
| 354 | |
| 355 | // If allocator starts with a '/' then it is being used to |
| 356 | // communicate the CPU/GPU that the graph runs on. |
| 357 | SessionOptions lsession_options = session_options; |
| 358 | const std::string& alloc_type = |
| 359 | lsession_options.config.gpu_options().allocator_type(); |
| 360 | if (!alloc_type.empty() && (alloc_type[0] == '/')) { |
| 361 | // Clear the device field from the graphdef so that the default device |
| 362 | // setting below will control which GPU the graph will run on. |
| 363 | for (tensorflow::NodeDef& node : |
| 364 | *bundle->meta_graph_def.mutable_graph_def()->mutable_node()) { |
| 365 | if (!tensorflow::grappler::NodeIsOnCpu(&node)) { |
| 366 | node.clear_device(); |
| 367 | } |
| 368 | } |
| 369 | graph::SetDefaultDevice(alloc_type, bundle->meta_graph_def.mutable_graph_def()); |
| 370 | lsession_options.config.mutable_gpu_options()->clear_allocator_type(); |
| 371 | } |
| 372 | |
| 373 | TF_RETURN_IF_ERROR(LoadMetaGraphIntoSession( |
| 374 | bundle->meta_graph_def, lsession_options, &bundle->session)); |
| 375 | |
| 376 | std::vector<AssetFileDef> asset_file_defs; |
| 377 | TF_RETURN_IF_ERROR( |
| 378 | GetAssetFileDefs(bundle->meta_graph_def, &asset_file_defs)); |
| 379 | TF_RETURN_IF_ERROR( |
| 380 | RunRestore(run_options, export_dir, |
| 381 | bundle->meta_graph_def.saver_def().restore_op_name(), |
| 382 | bundle->meta_graph_def.saver_def().filename_tensor_name(), |
| 383 | asset_file_defs, bundle->session.get())); |
| 384 | // Record walltime spent in restoring graph from disk, but postpone metric |
| 385 | // increments until graph init finishes. |
| 386 | const uint64 restore_graph_walltime = |
| 387 | GetLatencyMicroseconds(read_start_microseconds); |
| 388 | |
| 389 | const uint64 graph_init_start_microseconds = Env::Default()->NowMicros(); |
| 390 | string init_op_name; |
| 391 | TF_RETURN_IF_ERROR( |
| 392 | GetInitOp(export_dir, bundle->meta_graph_def, &init_op_name)); |
| 393 | TF_RETURN_IF_ERROR(RunInitOp(run_options, export_dir, bundle->meta_graph_def, |
| 394 | asset_file_defs, bundle->session.get(), |
| 395 | init_op_name)); |
| 396 | load_latency_by_stage->GetCell(export_dir, "restore_graph") |
| 397 | ->Add(restore_graph_walltime); |
| 398 | // Record wall time spent in init op. |
| 399 | load_latency_by_stage->GetCell(export_dir, "init_graph") |
| 400 | ->Add(GetLatencyMicroseconds(graph_init_start_microseconds)); |
| 401 | return Status::OK(); |
| 402 | } |
| 403 |
no test coverage detected