1# Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Functions to convert SavedModel to frozen GraphDefs.""" 16 17from tensorflow.core.framework import types_pb2 18from tensorflow.lite.python import util 19from tensorflow.lite.python.convert_phase import Component 20from tensorflow.lite.python.convert_phase import convert_phase 21from tensorflow.lite.python.convert_phase import SubComponent 22from tensorflow.python.client import session 23from tensorflow.python.framework import ops 24from tensorflow.python.platform import tf_logging as logging 25from tensorflow.python.saved_model import constants 26from tensorflow.python.saved_model import loader 27 28 29def _log_tensor_details(tensor_info): 30 """Log tensor details: name, shape, and type.""" 31 for key in tensor_info: 32 val = tensor_info[key] 33 dtype = types_pb2.DataType.Name(val.dtype) 34 if val.tensor_shape.unknown_rank: 35 shape = "unknown_rank" 36 else: 37 dims = [str(dim.size) for dim in val.tensor_shape.dim] 38 shape = "({})".format(", ".join(dims)) 39 40 logging.info("Tensor's key in saved_model's tensor_map: %s", key) 41 logging.info(" tensor name: %s, shape: %s, type: %s", val.name, shape, 42 dtype) 43 44 45def get_meta_graph_def(saved_model_dir, tag_set): 46 """Validate saved_model and extract MetaGraphDef. 47 48 Args: 49 saved_model_dir: saved_model path to convert. 50 tag_set: Set of tag(s) of the MetaGraphDef to load. 51 52 Returns: 53 The meta_graph_def used for tflite conversion. 54 55 Raises: 56 ValueError: No valid MetaGraphDef for given tag_set. 57 """ 58 with session.Session(graph=ops.Graph()) as sess: 59 return loader.load(sess, tag_set, saved_model_dir) 60 61 62def get_signature_def(meta_graph, signature_key): 63 """Get the signature def from meta_graph with given signature_key. 64 65 Args: 66 meta_graph: meta_graph_def. 67 signature_key: signature_def in the meta_graph_def. 68 69 Returns: 70 The signature_def used for tflite conversion. 71 72 Raises: 73 ValueError: Given signature_key is not valid for this meta_graph. 74 """ 75 signature_def_map = meta_graph.signature_def 76 signature_def_keys = set(signature_def_map.keys()) 77 logging.info( 78 "The given SavedModel MetaGraphDef contains SignatureDefs with the " 79 "following keys: %s", signature_def_keys) 80 if signature_key not in signature_def_keys: 81 raise ValueError("No '{}' in the SavedModel\'s SignatureDefs. Possible " 82 "values are '{}'.".format(signature_key, 83 ",".join(signature_def_keys))) 84 return signature_def_map[signature_key] 85 86 87def get_inputs_outputs(signature_def): 88 """Get inputs and outputs from SignatureDef. 89 90 Args: 91 signature_def: SignatureDef in the meta_graph_def for conversion. 92 93 Returns: 94 The inputs and outputs in the graph for conversion. 95 """ 96 inputs_tensor_info = signature_def.inputs 97 outputs_tensor_info = signature_def.outputs 98 logging.info("input tensors info: ") 99 _log_tensor_details(inputs_tensor_info) 100 logging.info("output tensors info: ") 101 _log_tensor_details(outputs_tensor_info) 102 103 def gather_names(tensor_info): 104 return [tensor_info[key].name for key in tensor_info] 105 106 inputs = gather_names(inputs_tensor_info) 107 outputs = gather_names(outputs_tensor_info) 108 return inputs, outputs 109 110 111def _get_tensors(graph, signature_def_tensor_names=None, 112 user_tensor_names=None): 113 """Gets the tensors associated with the tensor names. 114 115 Either signature_def_tensor_names or user_tensor_names should be provided. If 116 the user provides tensors, the tensors associated with the user provided 117 tensor names are provided. Otherwise, the tensors associated with the names in 118 the SignatureDef are provided. 119 120 Args: 121 graph: GraphDef representing graph. 122 signature_def_tensor_names: Tensor names stored in either the inputs or 123 outputs of a SignatureDef. (default None) 124 user_tensor_names: Tensor names provided by the user. (default None) 125 126 Returns: 127 List of tensors. 128 129 Raises: 130 ValueError: 131 signature_def_tensors and user_tensor_names are undefined or empty. 132 user_tensor_names are not valid. 133 """ 134 tensors = [] 135 if user_tensor_names: 136 # Sort the tensor names. 137 user_tensor_names = sorted(user_tensor_names) 138 139 tensors = util.get_tensors_from_tensor_names(graph, user_tensor_names) 140 elif signature_def_tensor_names: 141 tensors = [ 142 graph.get_tensor_by_name(name) 143 for name in sorted(signature_def_tensor_names) 144 ] 145 else: 146 # Throw ValueError if signature_def_tensors and user_tensor_names are both 147 # either undefined or empty. 148 raise ValueError( 149 "Specify either signature_def_tensor_names or user_tensor_names") 150 151 return tensors 152 153 154@convert_phase(Component.PREPARE_TF_MODEL, SubComponent.FREEZE_SAVED_MODEL) 155def freeze_saved_model(saved_model_dir, input_arrays, input_shapes, 156 output_arrays, tag_set, signature_key): 157 """Converts a SavedModel to a frozen graph. 158 159 Args: 160 saved_model_dir: SavedModel directory to convert. 161 input_arrays: List of input tensors to freeze graph with. Uses input arrays 162 from SignatureDef when none are provided. 163 input_shapes: Dict of strings representing input tensor names to list of 164 integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}). 165 Automatically determined when input shapes is None (e.g., {"foo" : None}). 166 output_arrays: List of output tensors to freeze graph with. Uses output 167 arrays from SignatureDef when none are provided. 168 tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to 169 analyze. All tags in the tag set must be present. 170 signature_key: Key identifying SignatureDef containing inputs and outputs. 171 172 Returns: 173 frozen_graph_def: Frozen GraphDef. 174 in_tensors: List of input tensors for the graph. 175 out_tensors: List of output tensors for the graph. 176 graph: `Graph` object. 177 178 Raises: 179 ValueError: 180 SavedModel doesn't contain a MetaGraphDef identified by tag_set. 181 signature_key is not in the MetaGraphDef. 182 assets/ directory is in the MetaGraphDef. 183 input_shapes does not match the length of input_arrays. 184 input_arrays or output_arrays are not valid. 185 """ 186 # Read SignatureDef. 187 meta_graph = get_meta_graph_def(saved_model_dir, tag_set) 188 signature_def = get_signature_def(meta_graph, signature_key) 189 inputs, outputs = get_inputs_outputs(signature_def) 190 191 # Check SavedModel for assets directory. 192 collection_def = meta_graph.collection_def 193 if constants.ASSETS_KEY in collection_def: 194 raise ValueError("SavedModels with assets/ directory are not supported.") 195 196 graph = ops.Graph() 197 with session.Session(graph=graph) as sess: 198 loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir) 199 200 # Gets input and output tensors. 201 # TODO(zhixianyan): Use TFLite supported Op list to filter outputs. 202 in_tensors = _get_tensors(graph, inputs, input_arrays) 203 out_tensors = _get_tensors(graph, outputs, output_arrays) 204 util.set_tensor_shapes(in_tensors, input_shapes) 205 206 frozen_graph_def = util.freeze_graph(sess, in_tensors, out_tensors) 207 return frozen_graph_def, in_tensors, out_tensors, sess.graph 208