| 2292 | } |
| 2293 | |
| 2294 | Status SavedModelOptimizer::AddIncrRestoreOps() { |
| 2295 | const std::string restore_op_name = |
| 2296 | meta_graph_def_->saver_def().restore_op_name(); |
| 2297 | Node* restore_op = nullptr; |
| 2298 | for (Node* node : graph_.nodes()) { |
| 2299 | if (node->name() == restore_op_name) { |
| 2300 | restore_op = node; |
| 2301 | break; |
| 2302 | } |
| 2303 | } |
| 2304 | |
| 2305 | if (!restore_op) { |
| 2306 | return errors::Internal( |
| 2307 | "Can not find restore op: " + restore_op_name); |
| 2308 | } |
| 2309 | |
| 2310 | std::vector<SrcInfo> input_nodes; |
| 2311 | Status s; |
| 2312 | if (restore_op->in_edges().size() == 1) { |
| 2313 | s = GetIncrRestoreOpInputs(restore_op, input_nodes); |
| 2314 | } else { |
| 2315 | for (const Edge* edge : restore_op->in_edges()) { |
| 2316 | const Node* sub_restore_op = edge->src(); |
| 2317 | if (sub_restore_op->in_edges().size() != 1) { |
| 2318 | LOG(FATAL) << "Sub restore all op only allow 1 input: " |
| 2319 | << sub_restore_op->DebugString(); |
| 2320 | } |
| 2321 | s = GetIncrRestoreOpInputs(sub_restore_op, input_nodes); |
| 2322 | } |
| 2323 | } |
| 2324 | |
| 2325 | TF_RETURN_IF_ERROR(s); |
| 2326 | |
| 2327 | return CreateRestoreAllNode( |
| 2328 | meta_graph_def_->saver_def().restore_op_name() + |
| 2329 | GetKvIncrRestoreAllNameSuffix(), |
| 2330 | "NoOp", input_nodes, &graph_); |
| 2331 | } |
| 2332 | |
| 2333 | Node* SavedModelOptimizer::UpdateRestoreShardNodeInputs( |
| 2334 | std::unordered_map<std::string, std::vector<Node*>>& origin_import_nodes, |
nothing calls this directly
no test coverage detected