1a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar/* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 3a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarLicensed under the Apache License, Version 2.0 (the "License"); 4a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumaryou may not use this file except in compliance with the License. 5a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarYou may obtain a copy of the License at 6a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 7a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar http://www.apache.org/licenses/LICENSE-2.0 8a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 9a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarUnless required by applicable law or agreed to in writing, software 10a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumardistributed under the License is distributed on an "AS IS" BASIS, 11a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarSee the License for the specific language governing permissions and 13a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarlimitations under the License. 14a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar==============================================================================*/ 15a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 16a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/cc/tools/freeze_saved_model.h" 17a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 18a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include <queue> 19a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 20a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/framework/attr_value.pb.h" 21a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/framework/function.pb.h" 22a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/framework/graph.pb.h" 23a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/framework/node_def.pb.h" 24a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/framework/versions.pb.h" 25a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/lib/core/errors.h" 26a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/lib/strings/str_util.h" 27a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar#include "tensorflow/core/protobuf/meta_graph.pb.h" 28a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 29a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarnamespace tensorflow { 30a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 31a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarnamespace { 32a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 33a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Gets tensor names from tensor_info and inserts them into the set of tensor 34a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// names. 35a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarvoid GetTensorNamesFromTensorInfo(const TensorInfo& tensor_info, 36a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string>* tensor_names) { 37a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (tensor_info.has_coo_sparse()) { 38a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // If the tensor is sparse we have to add all three tensors of the sparse 39a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // representations. 40a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const TensorInfo_CooSparse& coo_sparse = tensor_info.coo_sparse(); 41a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar tensor_names->insert(coo_sparse.values_tensor_name()); 42a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar tensor_names->insert(coo_sparse.indices_tensor_name()); 43a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar tensor_names->insert(coo_sparse.dense_shape_tensor_name()); 44a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } else { 45a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar tensor_names->insert(tensor_info.name()); 46a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 47a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 48a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 49a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Gets the union of all inputs and outputs of all SignatureDefs in the bundle 50a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarvoid GetSignatureDefsInputsAndOutputs( 51a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const SavedModelBundle& saved_model_bundle, 52a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string>* inputs, std::unordered_set<string>* outputs) { 53a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (auto& sigdef_elem : saved_model_bundle.meta_graph_def.signature_def()) { 54a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const SignatureDef& signature_def = sigdef_elem.second; 55a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (auto& input_elem : signature_def.inputs()) { 56a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GetTensorNamesFromTensorInfo(input_elem.second, inputs); 57a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 58a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (auto& output_elem : signature_def.outputs()) { 59a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GetTensorNamesFromTensorInfo(output_elem.second, outputs); 60a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 61a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 62a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 63a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 64a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Gets a map from string node name to NodeDef. 65a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarvoid GetNodeNameToNodeDefMap( 66a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GraphDef* graph_def, 67a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_map<string, NodeDef*>* name_to_node_map) { 68a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (size_t i = 0; i < graph_def->node_size(); i++) { 69a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar NodeDef* node = graph_def->mutable_node(i); 70a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar (*name_to_node_map)[node->name()] = node; 71a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 72a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 73a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 74a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Gets the set of node names needed by `outputs` and the corresponding set of 75a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// variable nodes to convert. 76a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarvoid GetReachableNodesAndVariables( 77a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GraphDef* graph_def, const std::unordered_set<string>& outputs, 78a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string>* reachable_node_names, 79a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string>* variable_node_names) { 80a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // TODO(suharshs): Add support for ResourceVariables. 81a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar static const std::unordered_set<string>* kVariableTypes = 82a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar new std::unordered_set<string>({"Variable", "VariableV2"}); 83a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // name_to_node_map is needed to get the inputs from the NodeDef corresponding 84a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // the a string node name. These inputs are used when doing our backwards 85a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // traversal. 86a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_map<string, NodeDef*> name_to_node_map; 87a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GetNodeNameToNodeDefMap(graph_def, &name_to_node_map); 88a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::queue<string> nodes_to_visit; 89a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (const string& tensor_name : outputs) { 90a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // We need to strip off the tensor part to get the node name. 91a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::vector<string> tensor_name_parts = str_util::Split(tensor_name, ':'); 92a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar nodes_to_visit.push(tensor_name_parts[0]); 93a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 94a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // We do a traversal backwards from the outputs specified in the MetaGraphDef. 95a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar while (!nodes_to_visit.empty()) { 96a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const string node_name = nodes_to_visit.front(); 97a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar nodes_to_visit.pop(); 98a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (reachable_node_names->find(node_name) != reachable_node_names->end()) { 99a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar continue; 100a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 101a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar reachable_node_names->insert(node_name); 102a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar NodeDef* node = name_to_node_map[node_name]; 103a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (kVariableTypes->find(node->op()) != kVariableTypes->end()) { 104a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar variable_node_names->insert(node->name()); 105a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 106a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (const string& input : node->input()) { 107a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar nodes_to_visit.push(input); 108a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 109a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 110a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 111a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 112a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Gets a map from variable name to variable value. 113a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarStatus GetVariableNameToTensorMap( 114a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar Session* session, std::unordered_set<string> variable_names_set, 115a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_map<string, Tensor>* variable_name_to_value_map) { 116a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (variable_names_set.empty()) { 117a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar return Status::OK(); 118a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 119a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::vector<string> variable_names; 120a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::vector<string> tensor_names; 121a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (const string& node_name : variable_names_set) { 122a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar variable_names.push_back(node_name); 123a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // We need to run tensors, so append ":0". 124a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar tensor_names.push_back(node_name + ":0"); 125a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 126a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::vector<Tensor> outputs; 127a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar TF_RETURN_IF_ERROR( 128a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar session->Run(/* inputs */ {}, tensor_names, /* targets */ {}, &outputs)); 129a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (size_t i = 0; i < variable_names.size(); i++) { 130a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar (*variable_name_to_value_map)[variable_names[i]] = outputs[i]; 131a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 132a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar return Status::OK(); 133a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 134a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 135a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Converts a Variable NodeDef into a Constant NodeDef. 136a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumarvoid ConvertVariableToConstant(const NodeDef& variable_node, 137a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const Tensor& variable_value, 138a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar NodeDef* const_node) { 139a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const_node->set_name(variable_node.name()); 140a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const_node->set_op("Const"); 141a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar (*const_node->mutable_attr())["dtype"] = variable_node.attr().at("dtype"); 142a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar variable_value.AsProtoTensorContent( 143a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar (*const_node->mutable_attr())["value"].mutable_tensor()); 144a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 145a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 146a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar// Freezes the subgraph of all nodes needed by `outputs`. 147a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarStatus FreezeGraphDef(const SavedModelBundle& saved_model_bundle, 148a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar const std::unordered_set<string>& outputs, 149a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GraphDef* frozen_graph_def) { 150a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GraphDef graph_def = saved_model_bundle.meta_graph_def.graph_def(); 151a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // Copy versions and library as-is from original graph. 152a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar *frozen_graph_def->mutable_versions() = graph_def.versions(); 153a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar *frozen_graph_def->mutable_library() = graph_def.library(); 154a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // If the graph is empty there is nothing left to do. 155a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (graph_def.node_size() == 0) { 156a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar return Status::OK(); 157a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 158a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string> reachable_node_names; 159a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string> variable_node_names; 160a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GetReachableNodesAndVariables(&graph_def, outputs, &reachable_node_names, 161a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar &variable_node_names); 162a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_map<string, Tensor> variable_to_value_map; 163a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar TF_RETURN_IF_ERROR( 164a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GetVariableNameToTensorMap(saved_model_bundle.session.get(), 165a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar variable_node_names, &variable_to_value_map)); 166a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // We copy the nodes in the same order they were in the original graph_def. 167a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar for (const NodeDef& node : graph_def.node()) { 168a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (reachable_node_names.find(node.name()) == reachable_node_names.end()) { 169a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar continue; 170a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 171a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar if (variable_node_names.find(node.name()) != variable_node_names.end()) { 172a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar ConvertVariableToConstant(node, variable_to_value_map[node.name()], 173a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar frozen_graph_def->add_node()); 174a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } else { 175a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar // If the node isn't a variable, just copy the node as-is. 176a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar *frozen_graph_def->add_node() = node; 177a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 178a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar } 179a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar return Status::OK(); 180a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 181a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 182a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} // namespace 183a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 184a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh SivakumarStatus FreezeSavedModel(const SavedModelBundle& saved_model_bundle, 185a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GraphDef* frozen_graph_def, 186a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string>* inputs, 187a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar std::unordered_set<string>* outputs) { 188a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar GetSignatureDefsInputsAndOutputs(saved_model_bundle, inputs, outputs); 189a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar TF_RETURN_IF_ERROR( 190a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar FreezeGraphDef(saved_model_bundle, *outputs, frozen_graph_def)); 191a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar return Status::OK(); 192a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} 193a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar 194a79e97be8460ce3e1a7de2ddbc78b76151e0035aSuharsh Sivakumar} // namespace tensorflow 195