Function ProcessTargetEdges

tensorflow/core/graph/quantize_training.cc:568–594

Insert conversion op, connect it to the graph and remove the old edge.

// Insert conversion op, connect it to the graph and remove the old edge.
Status ProcessTargetEdges(Graph* graph, const string& quant_op_type,
                          const std::vector<EdgeToConvert>& target_edges) {
  // Remember previously converted ops to avoid duplicated conversion on the
  // same input.
  std::unordered_map<string, Node*, StringPieceHasher> name_index;
  std::vector<Node*> added_variables;
  for (const EdgeToConvert edge : target_edges) {
    Node* convert_node;
    string name_prefix = edge.edge->src()->name();

    auto iter = name_index.find(name_prefix);
    if (iter == name_index.end()) {
      TF_RETURN_IF_ERROR(MakeQuantizeOp(graph, name_prefix, quant_op_type, edge,
                                        &added_variables, &convert_node));
      name_index[name_prefix] = convert_node;
    } else {
      convert_node = iter->second;
    }

    graph->AddEdge(convert_node, 0, edge.edge->dst(), edge.edge->dst_input());
    graph->RemoveEdge(edge.edge);
  }

  TF_RETURN_IF_ERROR(AddSaveAndRestore(graph, added_variables));

  return Status::OK();
}

}  // namespace