| 427 | } |
| 428 | |
| 429 | Status SegmentGraph(const Graph* tf_graph, |
| 430 | const std::function<Status(const Node*)>& candidate_fn, |
| 431 | const std::function<bool(const Edge*)>& input_candidate_fn, |
| 432 | const std::function<bool(const Edge*)>& output_candidate_fn, |
| 433 | const SegmentOptions& options, |
| 434 | SegmentNodesVector* segments) { |
| 435 | // Steps: |
| 436 | // 1. run the segmentation algorithm to find all the segments, which uses |
| 437 | // candidate_fn to determine the candidates segment nodes; |
| 438 | // 2. for each segments, remove the nodes that are inputs/outputs of the |
| 439 | // segment but are not eligible, using input/output_candidate_fn to |
| 440 | // determine the eligibilities; |
| 441 | // 3. convert the segment into expected return format and return the result. |
| 442 | |
| 443 | // --------------------------------- Step 1 --------------------------------- |
| 444 | auto graph = std::unique_ptr<SimpleGraph>(new SimpleGraph(tf_graph)); |
| 445 | // Use a union-find to collect the nodes that belong to the same |
| 446 | // segment. A node value of nullptr indicates that the node is not a candidate |
| 447 | // for TRT. |
| 448 | std::unordered_set<string> unsupported_ops; |
| 449 | int num_unsupported_ops = 0; |
| 450 | |
| 451 | // Getting the nodes blacklisted for conversion |
| 452 | string tftrt_node_blacklist_str; |
| 453 | TF_CHECK_OK(ReadStringFromEnvVar( |
| 454 | "TF_TRT_OP_BLACKLIST", "", &tftrt_node_blacklist_str |
| 455 | )); |
| 456 | |
| 457 | auto tftrt_node_blacklist = gtl::FlatSet<string>{}; |
| 458 | |
| 459 | for (const auto& x : str_util::Split(tftrt_node_blacklist_str, ",")) { |
| 460 | tftrt_node_blacklist.insert(x); |
| 461 | } |
| 462 | |
| 463 | // User defined special subgraphs which can not be convert to trt graph. |
| 464 | // e.g. some sparse lookup subgraphs. |
| 465 | auto labeled_node_blacklist = gtl::FlatSet<string>{}; |
| 466 | GetLabeledNodes(&labeled_node_blacklist, const_cast<Graph*>(tf_graph)); |
| 467 | |
| 468 | // Parsing each node of the graph |
| 469 | std::vector<UnionFind<SimpleNode*>> node_segments; |
| 470 | for (int i = 0; i < graph->num_node_ids(); ++i) { |
| 471 | SimpleNode* node = graph->FindNodeId(i); |
| 472 | if (options.exclude_node_list.count(node->name()) != 0) { |
| 473 | VLOG(1) << "Not a TF-TRT candidate, " |
| 474 | << "(Op type: " << node->tf_node()->type_string() << "), " |
| 475 | << "(Op name: " << node->name() << "), " |
| 476 | << "(Reason: excluded by segmenter option)"; |
| 477 | unsupported_ops.emplace(node->tf_node()->type_string()); |
| 478 | num_unsupported_ops++; |
| 479 | node = nullptr; |
| 480 | } else { |
| 481 | const Status status = candidate_fn(node->tf_node()); |
| 482 | if (!status.ok()) { |
| 483 | VLOG(1) << "Not a TF-TRT candidate, " |
| 484 | << "(Op type: " << node->tf_node()->type_string() << "), " |
| 485 | << "(Op name: " << node->name() << "), " |
| 486 | << "(Reason: " << status << ")"; |