blob: a2f744d36fc61417ddf7a526ed61f9c91fc7b246 [file] [log] [blame]
Tim Hall79d07d22020-04-27 18:20:16 +01001# Copyright (C) 2020 Arm Limited or its affiliates. All rights reserved.
2#
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
Tim Hallc8310b12020-06-17 14:53:11 +010023from .errors import TensorError
Diego Russoe8a10452020-04-21 17:39:10 +010024from .nn_graph import Graph
25from .nn_graph import Subgraph
Diego Russoea6111a2020-04-14 18:41:58 +010026from .operation import Operation
Diego Russoe8a10452020-04-21 17:39:10 +010027from .tensor import QuantizationParameters
28from .tensor import Tensor
29from .tflite.BuiltinOperator import BuiltinOperator
30from .tflite.Model import Model
31from .tflite_mapping import builtin_operator_map
32from .tflite_mapping import DataType
33from .tflite_mapping import datatype_map
34from .tflite_mapping import datatype_map_numpy
Tim Hall79d07d22020-04-27 18:20:16 +010035
36
37def decode_str(s):
38 if s is None:
39 return ""
40 return s.decode("utf-8")
41
42
Louis Verhaard3c07c972020-05-07 08:12:58 +020043def clone_and_reshape_tensor(src_tens, reorder):
Tim Hall79d07d22020-04-27 18:20:16 +010044
Louis Verhaard3c07c972020-05-07 08:12:58 +020045 tens = src_tens.clone("_reshape")
46 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
56 op = Operation("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 != []:
79 TensorError(tens, "This subgraph input tensor has unexpected driving operators.")
80
Tim Hall79d07d22020-04-27 18:20:16 +010081 op = Operation("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:
86 op = Operation("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())
108 dtype = datatype_map[tens_data.Type()]
Tim Hall79d07d22020-04-27 18:20:16 +0100109 tens = Tensor(shape, dtype, name)
Tim Hall79d07d22020-04-27 18:20:16 +0100110 quant = tens_data.Quantization()
111
Tim Hall79d07d22020-04-27 18:20:16 +0100112 tens.quantization = QuantizationParameters()
Tim Halle4e58e12020-05-08 09:50:21 +0100113 if quant is not None:
Diego Russod0eee262020-04-23 18:14:37 +0100114 tens.quantization.min = self.len1_array_to_scalar(quant.MinAsNumpy())
115 tens.quantization.max = self.len1_array_to_scalar(quant.MaxAsNumpy())
116 tens.quantization.scale_f32 = self.len1_array_to_scalar(quant.ScaleAsNumpy())
117 tens.quantization.zero_point = self.len1_array_to_scalar(quant.ZeroPointAsNumpy())
Tim Hall79d07d22020-04-27 18:20:16 +0100118
119 if dtype == DataType.uint8:
120 tens.quantization.quant_min = 0
121 tens.quantization.quant_max = (1 << dtype.bits) - 1
122 elif dtype in set((DataType.int8, DataType.int16, DataType.int32, DataType.int64)):
123 tens.quantization.quant_min = -(1 << (dtype.bits - 1))
124 tens.quantization.quant_max = (1 << (dtype.bits - 1)) - 1
Tim Hall79d07d22020-04-27 18:20:16 +0100125
126 if tens.quantization.scale_f32 is None and tens.quantization.zero_point is None:
127 tens.quantization = None
128
129 tens.values = None
130 buf = self.graph.buffers[tens_data.Buffer()]
131 if buf is not None:
132 tens.values = np.array(buf.view(datatype_map_numpy[tens_data.Type()]).reshape(shape))
133 if tens.quantization is not None:
134 tens.quant_values = tens.values
135 tens.values = tens.quantization.dequantize(tens.quant_values)
136 return tens
137
Tim Hallc8310b12020-06-17 14:53:11 +0100138 def parse_operator(self, op_index, op_data):
Tim Hall79d07d22020-04-27 18:20:16 +0100139 op_type, opt_serializer = self.graph.operator_codes[op_data.OpcodeIndex()]
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +0200140 inputs = [self.tensors[idx] if idx != -1 else None for idx in op_data.InputsAsNumpy()]
141 outputs = [self.tensors[idx] if idx != -1 else None for idx in op_data.OutputsAsNumpy()]
Tim Hall79d07d22020-04-27 18:20:16 +0100142 name = "unknown_op_name"
143 if len(outputs):
144 name = outputs[0].name
145 op = Operation(op_type, name)
Tim Hallc8310b12020-06-17 14:53:11 +0100146 op.op_index = op_index
Tim Hall79d07d22020-04-27 18:20:16 +0100147 op.inputs = inputs
148 op.outputs = outputs
149 for out in op.outputs:
150 out.ops = [op]
151
152 activation_function_to_split_out = None
153
154 if op_type.startswith("DepthwiseConv2d") or op_type.startswith("Conv2D"):
Louis Verhaard3c07c972020-05-07 08:12:58 +0200155 inputs[1] = clone_and_reshape_tensor(inputs[1], (1, 2, 3, 0))
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +0200156 if len(inputs) < 3 or (len(inputs) < 4 and "Backprop" in op_type):
157 # No Bias tensor
158 inputs.append(None)
159 if inputs[-1]:
160 inputs[-1] = clone_and_reshape_tensor(inputs[-1], (0,))
Tim Hall79d07d22020-04-27 18:20:16 +0100161
162 if op_type.startswith("FullyConnected"):
Louis Verhaard3c07c972020-05-07 08:12:58 +0200163 inputs[1] = clone_and_reshape_tensor(inputs[1], (1, 0))
Jacob Bohlin67e0d8f2020-08-20 10:53:02 +0200164 if len(inputs) < 3:
165 # No Bias tensor
166 inputs.append(None)
167 if inputs[-1]:
168 inputs[-1] = clone_and_reshape_tensor(inputs[-1], (0,))
Tim Hall79d07d22020-04-27 18:20:16 +0100169
170 if opt_serializer is not None:
Tim Hallc8310b12020-06-17 14:53:11 +0100171 op.attrs = opt_serializer.deserialize(op_data)
Tim Hall79d07d22020-04-27 18:20:16 +0100172
Michael McGeagh7b245fd2020-07-31 12:50:57 +0100173 if op_type == "Reshape" and "new_shape" not in op.attrs:
174 # Reshape should have an attrib "new_shape" but if it is missing, add it based on the output shape
175 op.attrs["new_shape"] = outputs[0].shape
176
Tim Hall79d07d22020-04-27 18:20:16 +0100177 if "stride_w" in op.attrs:
178 op.attrs["strides"] = (1, op.attrs["stride_h"], op.attrs["stride_w"], 1)
179 if "filter_width" in op.attrs:
180 op.attrs["ksize"] = (1, op.attrs["filter_height"], op.attrs["filter_width"], 1)
181 if "dilation_w_factor" in op.attrs:
182 op.attrs["dilation"] = (1, op.attrs["dilation_h_factor"], op.attrs["dilation_w_factor"], 1)
183 if "depth_multiplier" in op.attrs:
184 op.attrs["channel_multiplier"] = op.attrs["depth_multiplier"]
185
186 if "fused_activation_function" in op.attrs:
187 if op_type in set(("ConcatTFLite",)):
188 act = op.attrs["fused_activation_function"]
189 del op.attrs["fused_activation_function"]
190 if act is not None:
191 activation_function_to_split_out = act
192
193 if activation_function_to_split_out is not None:
194 act_op = Operation(activation_function_to_split_out, name + activation_function_to_split_out)
195 out_tens = op.outputs[0]
196 intermediate_tens = out_tens.clone("_act_intermediate")
Michael McGeaghc5b549b2020-08-07 11:54:28 +0100197 act_op.set_output_tensor(out_tens)
Tim Hall79d07d22020-04-27 18:20:16 +0100198 intermediate_tens.ops = [op]
199 op.outputs[0] = intermediate_tens
200 act_op.inputs = [intermediate_tens]
201
Diego Russod0eee262020-04-23 18:14:37 +0100202 @staticmethod
203 def len1_array_to_scalar(arr):
204 # The following flatbuffer quantisation fields all return a scalar value of 0 if they are not definied in
205 # the input buffer. This is represented in Vela by using None.
206 # Otherwise, the fields returned are a single or multi-element array. In which case, single element arrays
207 # are converted to scalars
208 if isinstance(arr, int) and arr == 0:
209 return None
210 if len(arr) == 1:
211 return arr[0]
212 return arr
213
Tim Hall79d07d22020-04-27 18:20:16 +0100214
215class TFLiteGraph:
216 def __init__(
Diego Russoea6111a2020-04-14 18:41:58 +0100217 self, filename, batch_size=1, feed_dict={}, output_node_names=[], initialisation_nodes=[],
Tim Hall79d07d22020-04-27 18:20:16 +0100218 ):
219
220 self.op_times = {}
221 if batch_size is None:
222 batch_size = 1
223 self.batch_size = batch_size
224 self.name = os.path.splitext(os.path.basename(filename))[0]
225 self.initialisation_nodes = initialisation_nodes
226
227 with open(filename, "rb") as f:
228 buf = bytearray(f.read())
229
230 model = Model.GetRootAsModel(buf, 0)
231
232 self.buffers = []
233 for idx in range(model.BuffersLength()):
234 self.buffers.append(self.parse_buffer(model.Buffers(idx)))
235
236 self.operator_codes = []
237 for idx in range(model.OperatorCodesLength()):
238 self.operator_codes.append(self.parse_operator_code(model.OperatorCodes(idx)))
239
240 self.subgraphs = []
241 for idx in range(model.SubgraphsLength()):
242 self.subgraphs.append(TFLiteSubgraph(self, model.Subgraphs(idx)))
243
244 self.nng = Graph(self.name, self.batch_size)
245 for tflite_sg in self.subgraphs:
246 sg = Subgraph(tflite_sg.name)
247 sg.original_inputs = tflite_sg.inputs # Preserve the original input order
248 sg.output_tensors = tflite_sg.outputs
249 self.nng.subgraphs.append(sg)
250
Michael McGeagh22f74e12020-08-07 16:21:03 +0100251 # Preserve the original metadata
252 for idx in range(model.MetadataLength()):
253 meta = model.Metadata(idx)
254 name = meta.Name()
255 if name is not None:
256 buf_data = self.buffers[meta.Buffer()]
257 self.nng.metadata.append((name, buf_data))
258
Tim Hall79d07d22020-04-27 18:20:16 +0100259 def parse_buffer(self, buf_data):
260 if buf_data.DataLength() == 0:
261 return None
262 data = buf_data.DataAsNumpy()
263 return data
264
265 def parse_operator_code(self, code):
266 c = code.BuiltinCode()
Tim Hallc30f4952020-06-15 20:47:35 +0100267 if c not in builtin_operator_map:
Louis Verhaard678645b2020-06-15 15:22:47 +0200268 msg = "The input file contains operator code {} which is currently not supported".format(c)
269 raise InputFileError(self.name, msg)
Tim Hall79d07d22020-04-27 18:20:16 +0100270 op_type, ser = builtin_operator_map[c]
271 if c == BuiltinOperator.CUSTOM:
272 op_type += decode_str(code.CustomCode())
273 return op_type, ser
274
275
276def read_tflite(
Diego Russoea6111a2020-04-14 18:41:58 +0100277 filename, batch_size=1, feed_dict={}, output_node_names=[], initialisation_nodes=[],
Tim Hall79d07d22020-04-27 18:20:16 +0100278):
Diego Russoea6111a2020-04-14 18:41:58 +0100279 tflite_graph = TFLiteGraph(filename, batch_size, feed_dict, output_node_names, initialisation_nodes)
Tim Hall79d07d22020-04-27 18:20:16 +0100280 nng = tflite_graph.nng
281 nng.refresh_after_modification()
282 return nng