Function FreezeGraphDef

tensorflow/cc/tools/freeze_saved_model.cc:175–219

Freezes the subgraph of all nodes needed by `outputs`.

// Freezes the subgraph of all nodes needed by `outputs`.
Status FreezeGraphDef(const SavedModelBundle& saved_model_bundle,
                      const std::unordered_set<string>& outputs,
                      GraphDef* frozen_graph_def) {
  GraphDef graph_def = saved_model_bundle.meta_graph_def.graph_def();
  // Copy versions and library as-is from original graph.
  *frozen_graph_def->mutable_versions() = graph_def.versions();
  *frozen_graph_def->mutable_library() = graph_def.library();
  // If the graph is empty there is nothing left to do.
  if (graph_def.node_size() == 0) {
    return Status::OK();
  }
  // name_to_node_map is needed to get the inputs from the NodeDef corresponding
  // the a string node name. These inputs are used when doing our backwards
  // traversal.
  std::unordered_map<string, NodeDef*> name_to_node_map;
  GetNodeNameToNodeDefMap(&graph_def, &name_to_node_map);
  std::unordered_set<string> reachable_node_names;
  std::unordered_set<string> variable_node_names;
  GetReachableNodesAndVariables(&graph_def, outputs, name_to_node_map,
                                &reachable_node_names, &variable_node_names);
  std::unordered_map<string, Tensor> variable_to_value_map;
  TF_RETURN_IF_ERROR(GetVariableNameToTensorMap(
      saved_model_bundle.session.get(), name_to_node_map, variable_node_names,
      &variable_to_value_map));
  // We copy the nodes in the same order they were in the original graph_def.
  for (const NodeDef& node : graph_def.node()) {
    if (reachable_node_names.find(node.name()) == reachable_node_names.end()) {
      continue;
    }
    if (variable_node_names.find(node.name()) != variable_node_names.end()) {
      ConvertVariableToConstant(node, variable_to_value_map[node.name()],
                                frozen_graph_def->add_node());
    } else if (node.op() == "ReadVariableOp" &&
               variable_node_names.find(node.input(0)) !=
                   variable_node_names.end()) {
      // If the node is a ReadVariableOp, its input VarHandleOp will be
      // converted to a Constant, so we will need to convert it to an Identity.
      ConvertReadVariableOpToIdentity(node, frozen_graph_def->add_node());
    } else {
      // If the node isn't a variable, just copy the node as-is.
      *frozen_graph_def->add_node() = node;
    }
  }
  return Status::OK();
}

}  // namespace