blob: daea1bf86467749813e37ab08fe93ce9de0c1960 [file] [log] [blame]
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +01001# Copyright (C) 2020-2021 Arm Limited or its affiliates. All rights reserved.
Tim Hall79d07d22020-04-27 18:20:16 +01002#
3# SPDX-License-Identifier: Apache-2.0
4#
5# Licensed under the Apache License, Version 2.0 (the License); you may
6# not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9# www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an AS IS BASIS, WITHOUT
13# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
Tim Hall79d07d22020-04-27 18:20:16 +010016# Description:
17# Functions used to read from a TensorFlow Lite format file.
Diego Russoea6111a2020-04-14 18:41:58 +010018import os.path
Tim Hall79d07d22020-04-27 18:20:16 +010019
20import numpy as np
Tim Hall79d07d22020-04-27 18:20:16 +010021
Louis Verhaard678645b2020-06-15 15:22:47 +020022from .errors import InputFileError
Diego Russoe8a10452020-04-21 17:39:10 +010023from .nn_graph import Graph
24from .nn_graph import Subgraph
Louis Verhaarde8a5a782020-11-02 18:04:27 +010025from .operation import create_activation_function
Louis Verhaardaee5d752020-09-30 09:01:52 +020026from .operation import Op
Diego Russoea6111a2020-04-14 18:41:58 +010027from .operation import Operation
Diego Russoe8a10452020-04-21 17:39:10 +010028from .tensor import QuantizationParameters
29from .tensor import Tensor
30from .tflite.BuiltinOperator import BuiltinOperator
31from .tflite.Model import Model
32from .tflite_mapping import builtin_operator_map
33from .tflite_mapping import DataType
34from .tflite_mapping import datatype_map
35from .tflite_mapping import datatype_map_numpy
Tim Hall79d07d22020-04-27 18:20:16 +010036
37
38def decode_str(s):
39 if s is None:
40 return ""
41 return s.decode("utf-8")
42
43
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +010044def clone_and_reshape_tensor(src_tens, reorder, set_unique):
45 tens = src_tens.clone("_reshape", set_unique)
Louis Verhaard3c07c972020-05-07 08:12:58 +020046 tens.shape = [src_tens.shape[idx] for idx in reorder]
47 tens.bandwidth_shape = tens.shape
48 tens.storage_shape = tens.shape
Tim Hall79d07d22020-04-27 18:20:16 +010049
Louis Verhaard3c07c972020-05-07 08:12:58 +020050 if tens.values is not None:
51 tens.values = tens.values.transpose(reorder)
Tim Hall79d07d22020-04-27 18:20:16 +010052
Louis Verhaard3c07c972020-05-07 08:12:58 +020053 if tens.quant_values is not None:
54 tens.quant_values = tens.quant_values.transpose(reorder)
55
Louis Verhaardaee5d752020-09-30 09:01:52 +020056 op = Operation(Op.Const, tens.name)
Michael McGeaghc5b549b2020-08-07 11:54:28 +010057 op.set_output_tensor(tens)
Louis Verhaard3c07c972020-05-07 08:12:58 +020058 return tens
Tim Hall79d07d22020-04-27 18:20:16 +010059
60
61class TFLiteSubgraph:
62 def __init__(self, graph, subgraph):
63 self.graph = graph
64 self.name = decode_str(subgraph.Name())
65
66 self.tensors = []
67 for idx in range(subgraph.TensorsLength()):
68 self.tensors.append(self.parse_tensor(subgraph.Tensors(idx)))
69
70 for idx in range(subgraph.OperatorsLength()):
Tim Hallc8310b12020-06-17 14:53:11 +010071 self.parse_operator(idx, subgraph.Operators(idx))
Tim Hall79d07d22020-04-27 18:20:16 +010072
Tim Hallc8310b12020-06-17 14:53:11 +010073 self.outputs = self.get_tensors_from_indices_remove_duplicates(subgraph.OutputsAsNumpy(), "output")
74 self.inputs = self.get_tensors_from_indices_remove_duplicates(subgraph.InputsAsNumpy(), "input")
Tim Hall79d07d22020-04-27 18:20:16 +010075
76 # Fix up tensors without operations. Generate either Placeholder or Constant ops
77 for tens in self.inputs:
Tim Hallc8310b12020-06-17 14:53:11 +010078 if tens.ops != []:
Michael McGeagh528a56d2020-12-16 11:33:21 +000079 tens.error("This subgraph input tensor has unexpected driving operators.")
Tim Hallc8310b12020-06-17 14:53:11 +010080
Louis Verhaardaee5d752020-09-30 09:01:52 +020081 op = Operation(Op.Placeholder, tens.name)
Michael McGeaghc5b549b2020-08-07 11:54:28 +010082 op.set_output_tensor(tens)
Tim Hall79d07d22020-04-27 18:20:16 +010083
84 for tens in self.tensors:
85 if not tens.ops:
Louis Verhaardaee5d752020-09-30 09:01:52 +020086 op = Operation(Op.Const, tens.name)
Michael McGeaghc5b549b2020-08-07 11:54:28 +010087 op.set_output_tensor(tens)
Tim Hall79d07d22020-04-27 18:20:16 +010088
Tim Hallc8310b12020-06-17 14:53:11 +010089 def get_tensors_from_indices_remove_duplicates(self, indices, warning_str):
90 tensors = []
91 for idx in indices:
92 tensor = self.tensors[idx]
93 if tensor not in tensors:
94 tensors.append(tensor)
95 else:
96 print(
97 "Warning: Subgraph {0} tensor ({1}) with idx = {2} already seen. Removing the duplicate.".format(
98 warning_str, tensor, idx
99 )
100 )
101
102 return tensors
103
Tim Hall79d07d22020-04-27 18:20:16 +0100104 def parse_tensor(self, tens_data):
105 np_shape = tens_data.ShapeAsNumpy()
106 shape = list(np_shape) if type(np_shape) is np.ndarray else []
107 name = decode_str(tens_data.Name())
Dwight Lidmane05de452020-11-05 15:56:08 +0100108 tens_dtype = tens_data.Type()
109 dtype = datatype_map[tens_dtype]
Tim Hall79d07d22020-04-27 18:20:16 +0100110 tens = Tensor(shape, dtype, name)
Tim Hall79d07d22020-04-27 18:20:16 +0100111 quant = tens_data.Quantization()
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100112 tens.is_variable = tens_data.IsVariable()
Tim Hall79d07d22020-04-27 18:20:16 +0100113
Tim Hall79d07d22020-04-27 18:20:16 +0100114 tens.quantization = QuantizationParameters()
Tim Halle4e58e12020-05-08 09:50:21 +0100115 if quant is not None:
Diego Russod0eee262020-04-23 18:14:37 +0100116 tens.quantization.min = self.len1_array_to_scalar(quant.MinAsNumpy())
117 tens.quantization.max = self.len1_array_to_scalar(quant.MaxAsNumpy())
118 tens.quantization.scale_f32 = self.len1_array_to_scalar(quant.ScaleAsNumpy())
119 tens.quantization.zero_point = self.len1_array_to_scalar(quant.ZeroPointAsNumpy())
Tim Hall79d07d22020-04-27 18:20:16 +0100120
121 if dtype == DataType.uint8:
122 tens.quantization.quant_min = 0
123 tens.quantization.quant_max = (1 << dtype.bits) - 1
Michael McGeaghf3e3ad72020-12-02 12:39:03 +0000124 elif dtype in (DataType.int8, DataType.int16, DataType.int32, DataType.int64):
Tim Hall79d07d22020-04-27 18:20:16 +0100125 tens.quantization.quant_min = -(1 << (dtype.bits - 1))
126 tens.quantization.quant_max = (1 << (dtype.bits - 1)) - 1
Tim Hall79d07d22020-04-27 18:20:16 +0100127
128 if tens.quantization.scale_f32 is None and tens.quantization.zero_point is None:
129 tens.quantization = None
130
131 tens.values = None
132 buf = self.graph.buffers[tens_data.Buffer()]
Louis Verhaardf4e12be2020-12-18 14:23:06 +0100133 if buf is not None:
134 np_dtype = datatype_map_numpy[tens_dtype]
135 if dtype == DataType.string:
136 tens.values = np.array(buf.view(np_dtype))
137 else:
138 tens.values = np.array(buf.view(np_dtype).reshape(shape))
139 if tens.quantization is not None:
140 tens.quant_values = tens.values
141 tens.values = tens.quantization.dequantize(tens.quant_values)
Tim Hall79d07d22020-04-27 18:20:16 +0100142 return tens
143
Tim Hallc8310b12020-06-17 14:53:11 +0100144 def parse_operator(self, op_index, op_data):
Louis Verhaardaee5d752020-09-30 09:01:52 +0200145 op_type, opt_serializer, custom_code = self.graph.operator_codes[op_data.OpcodeIndex()]
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +0200146 inputs = [self.tensors[idx] if idx != -1 else None for idx in op_data.InputsAsNumpy()]
147 outputs = [self.tensors[idx] if idx != -1 else None for idx in op_data.OutputsAsNumpy()]
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100148 intermediates = []
149 if op_data.IntermediatesLength():
150 intermediates = [self.tensors[idx] if idx != -1 else None for idx in op_data.IntermediatesAsNumpy()]
151
Tim Hall79d07d22020-04-27 18:20:16 +0100152 name = "unknown_op_name"
153 if len(outputs):
154 name = outputs[0].name
155 op = Operation(op_type, name)
Tim Hallc8310b12020-06-17 14:53:11 +0100156 op.op_index = op_index
Tim Hall79d07d22020-04-27 18:20:16 +0100157 op.inputs = inputs
158 op.outputs = outputs
Fredrik Svedberg8d0f4892021-02-16 21:59:50 +0100159 op.intermediates = intermediates
Tim Hall79d07d22020-04-27 18:20:16 +0100160 for out in op.outputs:
161 out.ops = [op]
162
Louis Verhaardaee5d752020-09-30 09:01:52 +0200163 if op.type.is_depthwise_conv2d_op() or op.type.is_conv2d_op() or op.type == Op.FullyConnected:
Andreas Nevalainend8c032d2020-09-11 10:25:09 +0200164 if inputs[1].values is not None:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200165 if op.type == Op.FullyConnected:
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100166 inputs[1] = clone_and_reshape_tensor(inputs[1], (1, 0), False)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200167 else:
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100168 inputs[1] = clone_and_reshape_tensor(inputs[1], (1, 2, 3, 0), False)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200169 if op.type.needs_bias() and len(inputs) <= op_type.info.indices.biases[0]:
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +0200170 # No Bias tensor
171 inputs.append(None)
Patrik Gustavssone2dbed22020-10-06 10:14:36 +0200172 if inputs[-1] and inputs[-1].values is not None:
Patrik Gustavsson34359582020-11-03 10:24:08 +0100173 # Since bias tensor is used for both bias and scale,
Patrik Gustavsson6ae0e422020-11-04 12:43:50 +0100174 # a clone with a unique equivalence_id is needed
175 inputs[-1] = clone_and_reshape_tensor(inputs[-1], (0,), True)
Tim Hall79d07d22020-04-27 18:20:16 +0100176
177 if opt_serializer is not None:
Tim Hallc8310b12020-06-17 14:53:11 +0100178 op.attrs = opt_serializer.deserialize(op_data)
Tim Hall79d07d22020-04-27 18:20:16 +0100179
Louis Verhaardaee5d752020-09-30 09:01:52 +0200180 if op_type == Op.Reshape and "new_shape" not in op.attrs:
Michael McGeagh7b245fd2020-07-31 12:50:57 +0100181 # Reshape should have an attrib "new_shape" but if it is missing, add it based on the output shape
182 op.attrs["new_shape"] = outputs[0].shape
183
Louis Verhaardaee5d752020-09-30 09:01:52 +0200184 if op_type == Op.Cast:
Andreas Nevalainend8c032d2020-09-11 10:25:09 +0200185 # Cast op should have "in/out_data_type" attribs add if missing
186 if "in_data_type" not in op.attrs:
187 op.attrs["in_data_type"] = inputs[0].dtype
188 if "out_data_type" not in op.attrs:
189 op.attrs["out_data_type"] = outputs[0].dtype
190
Tim Hall79d07d22020-04-27 18:20:16 +0100191 if "stride_w" in op.attrs:
192 op.attrs["strides"] = (1, op.attrs["stride_h"], op.attrs["stride_w"], 1)
193 if "filter_width" in op.attrs:
194 op.attrs["ksize"] = (1, op.attrs["filter_height"], op.attrs["filter_width"], 1)
195 if "dilation_w_factor" in op.attrs:
196 op.attrs["dilation"] = (1, op.attrs["dilation_h_factor"], op.attrs["dilation_w_factor"], 1)
197 if "depth_multiplier" in op.attrs:
198 op.attrs["channel_multiplier"] = op.attrs["depth_multiplier"]
199
Fredrik Svedbergbdf09f92020-11-18 11:30:21 +0100200 if op_type == Op.DepthwiseConv2DBias and op.attrs["depth_multiplier"] == 0:
201 # The depth multiplier is implicit and is calculated as weight channels / ifm channels
202 # Note however that the weights have been reshaped above.
203 # The original value is cached above in channel_multiplier
204 op.attrs["depth_multiplier"] = op.weights.shape[2] // op.ifm.shape[-1]
205
Louis Verhaarde8a5a782020-11-02 18:04:27 +0100206 faf = op.attrs.pop("fused_activation_function", None)
207 if faf is not None:
208 op.activation = create_activation_function(faf)
Louis Verhaardaee5d752020-09-30 09:01:52 +0200209 if custom_code is not None:
210 op.attrs["custom_code"] = custom_code
Tim Hall79d07d22020-04-27 18:20:16 +0100211
Diego Russod0eee262020-04-23 18:14:37 +0100212 @staticmethod
213 def len1_array_to_scalar(arr):
214 # The following flatbuffer quantisation fields all return a scalar value of 0 if they are not definied in
215 # the input buffer. This is represented in Vela by using None.
216 # Otherwise, the fields returned are a single or multi-element array. In which case, single element arrays
217 # are converted to scalars
218 if isinstance(arr, int) and arr == 0:
219 return None
220 if len(arr) == 1:
221 return arr[0]
222 return arr
223
Tim Hall79d07d22020-04-27 18:20:16 +0100224
225class TFLiteGraph:
Michael McGeagh6f725262020-12-03 15:21:36 +0000226 def __init__(self, filename, batch_size, feed_dict, output_node_names, initialisation_nodes):
Tim Hall79d07d22020-04-27 18:20:16 +0100227
228 self.op_times = {}
229 if batch_size is None:
230 batch_size = 1
231 self.batch_size = batch_size
232 self.name = os.path.splitext(os.path.basename(filename))[0]
233 self.initialisation_nodes = initialisation_nodes
234
235 with open(filename, "rb") as f:
236 buf = bytearray(f.read())
237
238 model = Model.GetRootAsModel(buf, 0)
239
240 self.buffers = []
241 for idx in range(model.BuffersLength()):
242 self.buffers.append(self.parse_buffer(model.Buffers(idx)))
243
244 self.operator_codes = []
245 for idx in range(model.OperatorCodesLength()):
246 self.operator_codes.append(self.parse_operator_code(model.OperatorCodes(idx)))
247
248 self.subgraphs = []
249 for idx in range(model.SubgraphsLength()):
250 self.subgraphs.append(TFLiteSubgraph(self, model.Subgraphs(idx)))
251
252 self.nng = Graph(self.name, self.batch_size)
253 for tflite_sg in self.subgraphs:
254 sg = Subgraph(tflite_sg.name)
255 sg.original_inputs = tflite_sg.inputs # Preserve the original input order
256 sg.output_tensors = tflite_sg.outputs
257 self.nng.subgraphs.append(sg)
258
Michael McGeagh22f74e12020-08-07 16:21:03 +0100259 # Preserve the original metadata
260 for idx in range(model.MetadataLength()):
261 meta = model.Metadata(idx)
262 name = meta.Name()
263 if name is not None:
264 buf_data = self.buffers[meta.Buffer()]
265 self.nng.metadata.append((name, buf_data))
266
Tim Hall79d07d22020-04-27 18:20:16 +0100267 def parse_buffer(self, buf_data):
268 if buf_data.DataLength() == 0:
269 return None
270 data = buf_data.DataAsNumpy()
271 return data
272
273 def parse_operator_code(self, code):
274 c = code.BuiltinCode()
Tim Hall42abec12021-02-04 21:31:57 +0000275 if c == 0:
276 c = code.DeprecatedBuiltinCode()
Tim Hallc30f4952020-06-15 20:47:35 +0100277 if c not in builtin_operator_map:
Michael McGeagh7a6f8432020-12-02 15:29:22 +0000278 raise InputFileError(
279 self.name, f"The input file contains operator code '{c}' which is currently not supported"
280 )
Tim Hall79d07d22020-04-27 18:20:16 +0100281 op_type, ser = builtin_operator_map[c]
Louis Verhaardaee5d752020-09-30 09:01:52 +0200282 custom_code = None
Tim Hall79d07d22020-04-27 18:20:16 +0100283 if c == BuiltinOperator.CUSTOM:
Louis Verhaardaee5d752020-09-30 09:01:52 +0200284 custom_code = decode_str(code.CustomCode())
285 return op_type, ser, custom_code
Tim Hall79d07d22020-04-27 18:20:16 +0100286
287
Michael McGeagh6f725262020-12-03 15:21:36 +0000288def read_tflite(filename, batch_size, feed_dict, output_node_names, initialisation_nodes):
Diego Russoea6111a2020-04-14 18:41:58 +0100289 tflite_graph = TFLiteGraph(filename, batch_size, feed_dict, output_node_names, initialisation_nodes)
Tim Hall79d07d22020-04-27 18:20:16 +0100290 nng = tflite_graph.nng
291 nng.refresh_after_modification()
292 return nng