Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/lite/python/convert_saved_model.py: 30%

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.lite.python import util 

18from tensorflow.lite.python.convert_phase import Component 

19from tensorflow.lite.python.convert_phase import convert_phase 

20from tensorflow.lite.python.convert_phase import SubComponent 

21from tensorflow.python.client import session 

22from tensorflow.python.framework import ops 

23from tensorflow.python.platform import tf_logging as logging 

24from tensorflow.python.saved_model import constants 

25from tensorflow.python.saved_model import loader 

26 

27 

28def get_meta_graph_def(saved_model_dir, tag_set): 

29 """Validate saved_model and extract MetaGraphDef. 

30 

31 Args: 

32 saved_model_dir: saved_model path to convert. 

33 tag_set: Set of tag(s) of the MetaGraphDef to load. 

34 

35 Returns: 

36 The meta_graph_def used for tflite conversion. 

37 

38 Raises: 

39 ValueError: No valid MetaGraphDef for given tag_set. 

40 """ 

41 with session.Session(graph=ops.Graph()) as sess: 

42 return loader.load(sess, tag_set, saved_model_dir) 

43 

44 

45def get_signature_def(meta_graph, signature_key): 

46 """Get the signature def from meta_graph with given signature_key. 

47 

48 Args: 

49 meta_graph: meta_graph_def. 

50 signature_key: signature_def in the meta_graph_def. 

51 

52 Returns: 

53 The signature_def used for tflite conversion. 

54 

55 Raises: 

56 ValueError: Given signature_key is not valid for this meta_graph. 

57 """ 

58 signature_def_map = meta_graph.signature_def 

59 signature_def_keys = set(signature_def_map.keys()) 

60 logging.info( 

61 "The given SavedModel MetaGraphDef contains SignatureDefs with the " 

62 "following keys: %s", signature_def_keys) 

63 if signature_key not in signature_def_keys: 

64 raise ValueError("No '{}' in the SavedModel\'s SignatureDefs. Possible " 

65 "values are '{}'.".format(signature_key, 

66 ",".join(signature_def_keys))) 

67 return signature_def_map[signature_key] 

68 

69 

70def get_inputs_outputs(signature_def): 

71 """Get inputs and outputs from SignatureDef. 

72 

73 Args: 

74 signature_def: SignatureDef in the meta_graph_def for conversion. 

75 

76 Returns: 

77 The inputs and outputs in the graph for conversion. 

78 """ 

79 inputs_tensor_info = signature_def.inputs 

80 outputs_tensor_info = signature_def.outputs 

81 

82 def gather_names(tensor_info): 

83 return [tensor_info[key].name for key in tensor_info] 

84 

85 inputs = gather_names(inputs_tensor_info) 

86 outputs = gather_names(outputs_tensor_info) 

87 return inputs, outputs 

88 

89 

90def _get_tensors(graph, signature_def_tensor_names=None, 

91 user_tensor_names=None): 

92 """Gets the tensors associated with the tensor names. 

93 

94 Either signature_def_tensor_names or user_tensor_names should be provided. If 

95 the user provides tensors, the tensors associated with the user provided 

96 tensor names are provided. Otherwise, the tensors associated with the names in 

97 the SignatureDef are provided. 

98 

99 Args: 

100 graph: GraphDef representing graph. 

101 signature_def_tensor_names: Tensor names stored in either the inputs or 

102 outputs of a SignatureDef. (default None) 

103 user_tensor_names: Tensor names provided by the user. (default None) 

104 

105 Returns: 

106 List of tensors. 

107 

108 Raises: 

109 ValueError: 

110 signature_def_tensors and user_tensor_names are undefined or empty. 

111 user_tensor_names are not valid. 

112 """ 

113 tensors = [] 

114 if user_tensor_names: 

115 # Sort the tensor names. 

116 user_tensor_names = sorted(user_tensor_names) 

117 

118 tensors = util.get_tensors_from_tensor_names(graph, user_tensor_names) 

119 elif signature_def_tensor_names: 

120 tensors = [ 

121 graph.get_tensor_by_name(name) 

122 for name in sorted(signature_def_tensor_names) 

123 ] 

124 else: 

125 # Throw ValueError if signature_def_tensors and user_tensor_names are both 

126 # either undefined or empty. 

127 raise ValueError( 

128 "Specify either signature_def_tensor_names or user_tensor_names") 

129 

130 return tensors 

131 

132 

133@convert_phase(Component.PREPARE_TF_MODEL, SubComponent.FREEZE_SAVED_MODEL) 

134def freeze_saved_model(saved_model_dir, input_arrays, input_shapes, 

135 output_arrays, tag_set, signature_key): 

136 """Converts a SavedModel to a frozen graph. 

137 

138 Args: 

139 saved_model_dir: SavedModel directory to convert. 

140 input_arrays: List of input tensors to freeze graph with. Uses input arrays 

141 from SignatureDef when none are provided. 

142 input_shapes: Dict of strings representing input tensor names to list of 

143 integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}). 

144 Automatically determined when input shapes is None (e.g., {"foo" : None}). 

145 output_arrays: List of output tensors to freeze graph with. Uses output 

146 arrays from SignatureDef when none are provided. 

147 tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to 

148 analyze. All tags in the tag set must be present. 

149 signature_key: Key identifying SignatureDef containing inputs and outputs. 

150 

151 Returns: 

152 frozen_graph_def: Frozen GraphDef. 

153 in_tensors: List of input tensors for the graph. 

154 out_tensors: List of output tensors for the graph. 

155 graph: `Graph` object. 

156 

157 Raises: 

158 ValueError: 

159 SavedModel doesn't contain a MetaGraphDef identified by tag_set. 

160 signature_key is not in the MetaGraphDef. 

161 assets/ directory is in the MetaGraphDef. 

162 input_shapes does not match the length of input_arrays. 

163 input_arrays or output_arrays are not valid. 

164 """ 

165 # Read SignatureDef. 

166 meta_graph = get_meta_graph_def(saved_model_dir, tag_set) 

167 signature_def = get_signature_def(meta_graph, signature_key) 

168 inputs, outputs = get_inputs_outputs(signature_def) 

169 

170 # Check SavedModel for assets directory. 

171 collection_def = meta_graph.collection_def 

172 if constants.ASSETS_KEY in collection_def: 

173 raise ValueError("SavedModels with assets/ directory are not supported.") 

174 

175 graph = ops.Graph() 

176 with session.Session(graph=graph) as sess: 

177 loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir) 

178 

179 # Gets input and output tensors. 

180 # TODO(zhixianyan): Use TFLite supported Op list to filter outputs. 

181 in_tensors = _get_tensors(graph, inputs, input_arrays) 

182 out_tensors = _get_tensors(graph, outputs, output_arrays) 

183 util.set_tensor_shapes(in_tensors, input_shapes) 

184 

185 frozen_graph_def = util.freeze_graph(sess, in_tensors, out_tensors) 

186 return frozen_graph_def, in_tensors, out_tensors, sess.graph