| 2223 | } |
| 2224 | |
| 2225 | Node* SavedModelOptimizer::FindRestoreShardNode() { |
| 2226 | const std::string restore_op_name = |
| 2227 | meta_graph_def_->saver_def().restore_op_name(); |
| 2228 | Node* restore_op = nullptr; |
| 2229 | for (Node* node : graph_.nodes()) { |
| 2230 | if (node->name() == restore_op_name) { |
| 2231 | restore_op = node; |
| 2232 | break; |
| 2233 | } |
| 2234 | } |
| 2235 | |
| 2236 | if (!restore_op) { |
| 2237 | LOG(FATAL) << "Can not find restore op: " |
| 2238 | << restore_op_name; |
| 2239 | } |
| 2240 | |
| 2241 | if (restore_op->in_edges().size() != 1) { |
| 2242 | LOG(FATAL) << "Restore all op only allow 1 input: " |
| 2243 | << restore_op->DebugString(); |
| 2244 | } |
| 2245 | |
| 2246 | Node* restore_shard = nullptr; |
| 2247 | for (const Edge* edge : restore_op->in_edges()) { |
| 2248 | restore_shard = edge->src(); |
| 2249 | if (restore_shard->type_string() != "NoOp") { |
| 2250 | LOG(FATAL) << "Restore op " << restore_op_name |
| 2251 | << " has input node which type is not NoOp." |
| 2252 | << restore_shard->DebugString(); |
| 2253 | } |
| 2254 | } |
| 2255 | |
| 2256 | return restore_shard; |
| 2257 | } |
| 2258 | |
| 2259 | Status SavedModelOptimizer::GetIncrRestoreOpInputs( |
| 2260 | const Node* restore_op, |
nothing calls this directly
no test coverage detected