Function ConnectVariablesToSaveOp

tensorflow/core/graph/quantize_training.cc:185–243

Add the added_variables as an inputs to the Save op. We change the inputs of the SaveV2 op to include the names of the added variables. We also add the variables as inputs to the save op.

// We change the inputs of the SaveV2 op to include the names of the added
// variables. We also add the variables as inputs to the save op.
Status ConnectVariablesToSaveOp(Graph* graph, Node* save_op,
                                const std::vector<const Edge*>& in_edges,
                                const std::vector<Node*>& added_variables) {
  Node* tensor_names_op = in_edges[1]->src();
  Node* shape_and_slices_op = in_edges[2]->src();

  // Get the tensor_names and shape_and_slices tensors from the const op.
  Tensor tensor_names;
  Tensor shape_and_slices;
  TF_RETURN_IF_ERROR(
      GetNodeAttr(tensor_names_op->attrs(), "value", &tensor_names));
  TF_RETURN_IF_ERROR(
      GetNodeAttr(shape_and_slices_op->attrs(), "value", &shape_and_slices));

  int tn_size = tensor_names.NumElements();
  int var_size = added_variables.size();

  // Create a new save_op that has inputs to all the new variables.
  NodeBuilder save_op_builder =
      NodeBuilder(save_op->name(), save_op->type_string());
  // The first three inputs are prefix, tensor_names, and shapes_and_slices.
  for (int i = 0; i < 3; i++) {
    save_op_builder = save_op_builder.Input(in_edges[i]->src());
  }
  std::vector<NodeBuilder::NodeOut> var_nodeouts;
  var_nodeouts.reserve(tn_size + var_size);
  // The rest of the inputs need to be used the construct the tensor list arg.
  for (int i = 3; i < in_edges.size(); i++) {
    var_nodeouts.emplace_back(in_edges[i]->src());
  }

  // Add the new values to the tensors and the op input.
  Tensor new_tensor_names(DT_STRING, TensorShape({tn_size + var_size}));
  Tensor new_shape_and_slices(DT_STRING, TensorShape({tn_size + var_size}));
  FillStringTensor(&new_tensor_names, tensor_names);
  FillStringTensor(&new_shape_and_slices, shape_and_slices);
  for (int i = 0; i < var_size; i++) {
    Node* var = added_variables[i];
    new_tensor_names.flat<tstring>()(tn_size + i) = var->name();
    new_shape_and_slices.flat<tstring>()(tn_size + i) = "";
    var_nodeouts.emplace_back(var);
  }
  save_op_builder = save_op_builder.Input(var_nodeouts);

  // Update the attrs.
  tensor_names_op->AddAttr("value", new_tensor_names);
  shape_and_slices_op->AddAttr("value", new_shape_and_slices);

  // Remove the old save_op and add the new one.
  Node* new_save_op;
  TF_RETURN_IF_ERROR(save_op_builder.Finalize(graph, &new_save_op));
  // Add outputs to the new_save_op, all outputs are control edges.
  for (const Edge* edge : save_op->out_edges()) {
    graph->AddControlEdge(new_save_op, edge->dst());
  }
  graph->RemoveNode(save_op);

  return Status::OK();