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