blob: cf40b5b5800903d6e58669a82fd028eaab00d77b [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 write to a TensorFlow Lite format file. Supports adding in file identifiers.
Tim Hall79d07d22020-04-27 18:20:16 +010018import flatbuffers
Diego Russoe8a10452020-04-21 17:39:10 +010019import flatbuffers.number_types as N
20import numpy as np
21from flatbuffers import encode
Diego Russoea6111a2020-04-14 18:41:58 +010022from flatbuffers.builder import UOffsetTFlags
23
Diego Russoe8a10452020-04-21 17:39:10 +010024from .nn_graph import PassPlacement
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020025from .tensor import MemType
Diego Russoe8a10452020-04-21 17:39:10 +010026from .tensor import TensorPurpose
Tim Hall79d07d22020-04-27 18:20:16 +010027from .tflite import Buffer
28from .tflite import Metadata
Diego Russoe8a10452020-04-21 17:39:10 +010029from .tflite import Model
30from .tflite import Operator
31from .tflite import OperatorCode
32from .tflite import QuantizationParameters
33from .tflite import SubGraph
34from .tflite import Tensor
35from .tflite_mapping import builtin_operator_inv_map
36from .tflite_mapping import BuiltinOperator
37from .tflite_mapping import custom_prefix
38from .tflite_mapping import datatype_inv_map
39
40# ugh, the python flatbuffer interface is missing a method to add in file identifier. patching it in here:
Tim Hall79d07d22020-04-27 18:20:16 +010041
42tflite_version = 3
43tflite_file_identifier = "TFL" + str(tflite_version)
44
45
Tim Hall79d07d22020-04-27 18:20:16 +010046def FinishWithFileIdentifier(self, rootTable, fid):
47 if fid is None or len(fid) != 4:
48 raise Exception("fid must be 4 chars")
49
50 flags = N.Uint8Flags
51 prepSize = 4
52 self.Prep(self.minalign, prepSize + len(fid))
53 for i in range(3, -1, -1):
54 self.head = self.head - flags.bytewidth
55 encode.Write(flags.packer_type, self.Bytes, self.Head(), ord(fid[i]))
56
57 return self.Finish(rootTable)
58
59
60flatbuffers.Builder.FinishWithFileIdentifier = FinishWithFileIdentifier
61
62
63def make_vector(v):
64 try:
65 len(v)
66 return v
67 except TypeError:
68 return [v]
69
70
71class TFLiteSerialiser:
72 def __init__(self, nng):
73 self.builder = flatbuffers.Builder(0)
74 self.nng = nng
75
76 self.scratch_buf_id = 0 # Always assign scratch to buffer 0
Patrik Gustavssoneca2e952020-05-27 09:15:11 +020077 self.scratch_fast_buf_id = 1 # Always assign scratch_fast to buffer 1
Tim Hall79d07d22020-04-27 18:20:16 +010078 self.buffer_offsets_map = {}
79 self.buffers_to_write = [] # have an empty array there
80
81 self.input_tensors = []
82 self.ops_to_ignore = set(("Const", "Placeholder", "SubgraphInput"))
83
84 self.tensors_to_reshape = {}
85
86 self.subgraphs_to_write = [sg for sg in self.nng.subgraphs if sg.placement == PassPlacement.Cpu]
87
88 all_ops = []
89 for sg in self.subgraphs_to_write:
90 for ps in sg.passes:
91 for op in ps.ops:
92 if op.type not in self.ops_to_ignore:
93 all_ops.append(op)
94 if op.type.startswith("Conv2D") or op.type.startswith("DepthwiseConv2d"):
95 self.tensors_to_reshape[op.inputs[1]] = (3, 0, 1, 2)
96 if op.type.startswith("FullyConnected"):
97 self.tensors_to_reshape[op.inputs[1]] = (1, 0)
98
99 self.operator_codes = list(sorted(set(op.type for op in all_ops)))
100 self.operator_code_map = {}
101
102 def write_byte_vector(self, v, alignment=1):
103 builder = self.builder
104 builder.StartVector(1, len(v), alignment)
105 for e in v[::-1]:
106 builder.PrependByte(e)
107 return builder.EndVector(len(v))
108
109 def write_int_vector(self, v):
110 builder = self.builder
111 builder.StartVector(4, len(v), 4)
112 for e in v[::-1]:
113 builder.PrependInt32(e)
114 return builder.EndVector(len(v))
115
116 def write_long_vector(self, v):
117 builder = self.builder
118 builder.StartVector(8, len(v), 8)
119 for e in v[::-1]:
120 builder.PrependInt64(e)
121 return builder.EndVector(len(v))
122
123 def write_float_vector(self, v):
124 builder = self.builder
125 builder.StartVector(4, len(v), 4)
126 for e in v[::-1]:
127 builder.PrependFloat32(e)
128 return builder.EndVector(len(v))
129
130 def write_offset_vector(self, v):
131 builder = self.builder
132 builder.StartVector(4, len(v), 4)
133 for e in v[::-1]:
134 builder.PrependUOffsetTRelative(e)
135 return builder.EndVector(len(v))
136
Tim Hallc8310b12020-06-17 14:53:11 +0100137 def assign_buffers_to_tensors(self, tensors, scratch_tensor):
138 if scratch_tensor is not None:
139 scratch_tensor_mem_area = scratch_tensor.mem_area
Tim Hall25f605c2020-05-18 18:04:26 +0100140 else:
141 scratch_tensor_mem_area = None # all tensors are initialised to MemArea.Unknown
142
Tim Hall79d07d22020-04-27 18:20:16 +0100143 buffer_map = {}
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200144
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200145 buf_idx = 2
Tim Hall79d07d22020-04-27 18:20:16 +0100146
147 for tens in tensors:
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200148 # Set buffer ids depending on allocation
149 if tens.is_allocated_in_tensor_arena(scratch_tensor_mem_area):
Tim Hall79d07d22020-04-27 18:20:16 +0100150 buffer_map[tens] = self.scratch_buf_id
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200151 elif tens.mem_type == MemType.Scratch_fast:
152 # For Scratch_fast when not co-allocated with scratch in the TensorArena:
153 buffer_map[tens] = self.scratch_fast_buf_id
Tim Hall79d07d22020-04-27 18:20:16 +0100154 else:
155 buffer_map[tens] = buf_idx
156 buf_idx += 1
157
Tim Hallc8310b12020-06-17 14:53:11 +0100158 # Initialize buffers_to_write to a length equal to number of buffers so
Tim Hall79d07d22020-04-27 18:20:16 +0100159 # they can be appended at the correct index during tensor serialization
160 self.buffers_to_write = [None] * (buf_idx)
161
162 return buffer_map
163
164 def serialise_operator_code(self, idx, code):
165 builder = self.builder
166 custom_code_offset = None
167 if code.startswith(custom_prefix):
168 tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]
169 custom_code_offset = builder.CreateString(code[len(custom_prefix) :])
170 else:
171 try:
172 tf_code, opt_serializer = builtin_operator_inv_map[code]
173 except KeyError:
174 print(
Diego Russoea6111a2020-04-14 18:41:58 +0100175 "Warning: Writing operation %s, which does not have a direct TensorFlow Lite mapping,"
176 "as a custom operation" % (code,)
Tim Hall79d07d22020-04-27 18:20:16 +0100177 )
178 tf_code, opt_serializer = builtin_operator_inv_map[custom_prefix]
179
180 if tf_code == BuiltinOperator.CUSTOM:
181 assert code == "NpuOp" # Currently only support serialising NPU operators as a custom op
182 custom_code_offset = builder.CreateString("ethos-u")
183
Tim Hallc8310b12020-06-17 14:53:11 +0100184 self.operator_code_map[code] = (idx, tf_code, opt_serializer)
Tim Hall79d07d22020-04-27 18:20:16 +0100185
186 OperatorCode.OperatorCodeStart(builder)
187 OperatorCode.OperatorCodeAddBuiltinCode(builder, tf_code)
188 if custom_code_offset is not None:
189 OperatorCode.OperatorCodeAddCustomCode(builder, custom_code_offset)
190
191 return OperatorCode.OperatorCodeEnd(builder)
192
193 def serialise_quantization_parameters(self, quant):
194 builder = self.builder
195
196 min = None
197 max = None
198 scale = None
199 zero_point = None
200 if quant is not None:
201 if quant.min is not None:
202 min = self.write_float_vector(make_vector(quant.min))
203 if quant.max is not None:
204 max = self.write_float_vector(make_vector(quant.max))
205 if quant.scale_f32 is not None:
206 scale = self.write_float_vector(make_vector(quant.scale_f32))
207 if quant.zero_point is not None:
208 zero_point = self.write_long_vector(make_vector(quant.zero_point))
209
210 QuantizationParameters.QuantizationParametersStart(builder)
211 if min is not None:
212 QuantizationParameters.QuantizationParametersAddMin(builder, min)
213 if max is not None:
214 QuantizationParameters.QuantizationParametersAddMax(builder, max)
215 if scale is not None:
216 QuantizationParameters.QuantizationParametersAddScale(builder, scale)
217 if zero_point is not None:
218 QuantizationParameters.QuantizationParametersAddZeroPoint(builder, zero_point)
219 return QuantizationParameters.QuantizationParametersEnd(builder)
220
221 def serialise_tensor(self, tens):
222 builder = self.builder
223 tens_shape = tens.shape
224 values = tens.quant_values
225 if values is None:
226 values = tens.values
227
228 if values is None:
229 values = np.empty(shape=(0), dtype=np.uint8)
230
231 if tens in self.tensors_to_reshape:
232 reorder = self.tensors_to_reshape[tens]
233 tens_shape = [tens_shape[idx] for idx in reorder]
234 values = values.transpose(reorder)
235
236 if tens.purpose == TensorPurpose.Scratch:
237 tens_shape = [0]
Tim Hall79d07d22020-04-27 18:20:16 +0100238
239 buf_id = self.buffer_map[tens]
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200240 self.buffers_to_write[buf_id] = values.flatten().view(np.uint8)
Tim Hall79d07d22020-04-27 18:20:16 +0100241
242 shape = self.write_int_vector(tens_shape)
243
244 name = builder.CreateString(tens.name)
245 quant = self.serialise_quantization_parameters(tens.quantization)
246
247 Tensor.TensorStart(builder)
248 Tensor.TensorAddShape(builder, shape)
249 Tensor.TensorAddType(builder, datatype_inv_map[tens.dtype])
250 # All tensors must have a valid backing buffer, even if it is empty.
251 # Empty buffers should be kept unique for TensorFlow Lite Micro
252 Tensor.TensorAddBuffer(builder, buf_id)
253 Tensor.TensorAddName(builder, name)
254 Tensor.TensorAddQuantization(builder, quant)
255
256 res = Tensor.TensorEnd(builder)
257 return res
258
259 def serialise_operator(self, op):
260 builder = self.builder
261
262 inputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.inputs])
263 outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.outputs])
264
265 op_idx, tflop, opt_serializer = self.operator_code_map[op.type]
266
267 builtin_opt_offset = None
268 custom_opt_offset = None
269 if opt_serializer is not None:
270 attrs = dict(op.attrs)
271 if "strides" in attrs:
272 attrs["stride_h"] = attrs["strides"][1]
273 attrs["stride_w"] = attrs["strides"][2]
274 if "ksize" in attrs:
275 attrs["filter_height"] = attrs["ksize"][1]
276 attrs["filter_width"] = attrs["ksize"][2]
277 if "dilation" in attrs:
278 attrs["dilation_h_factor"] = attrs["dilation"][1]
279 attrs["dilation_w_factor"] = attrs["dilation"][2]
280 if "channel_multiplier" in attrs:
281 attrs["depth_multiplier"] = attrs["channel_multiplier"]
282
283 builtin_opt_offset, custom_opt_offset = opt_serializer.serialize(builder, attrs)
284
285 mutating_variable_inputs_offset = self.write_byte_vector([])
286 Operator.OperatorStart(builder)
287 Operator.OperatorAddOpcodeIndex(builder, op_idx)
288 Operator.OperatorAddInputs(builder, inputs_offset)
289 Operator.OperatorAddOutputs(builder, outputs_offset)
290
291 if builtin_opt_offset is not None:
292 Operator.OperatorAddBuiltinOptionsType(builder, opt_serializer.builtin_opt_type)
293 Operator.OperatorAddBuiltinOptions(builder, builtin_opt_offset)
294 if custom_opt_offset is not None:
295 Operator.OperatorAddCustomOptions(builder, custom_opt_offset)
296 Operator.OperatorAddCustomOptionsFormat(builder, opt_serializer.custom_opt_format)
297
298 Operator.OperatorAddMutatingVariableInputs(builder, mutating_variable_inputs_offset)
299 return Operator.OperatorEnd(builder)
300
301 def serialise_subgraph(self, sg):
302 builder = self.builder
303 tensor_set = set()
304
305 all_ops = []
306 for ps in sg.passes:
307 for op in ps.ops:
308 if op.type not in self.ops_to_ignore:
309 all_ops.append(op)
310
311 for op in all_ops:
312 for tens in op.inputs + op.outputs:
313 tensor_set.add(tens)
314
315 all_tensors = [tens for nm, idx, tens in sorted((tens.name, idx, tens) for idx, tens in enumerate(tensor_set))]
316
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200317 scratch_tensors = [tens for tens in all_tensors if tens.name.endswith("scratch")]
318
319 for tens in all_tensors:
320 if tens.name.endswith("scratch_fast"):
321 scratch_fast_tensor = tens
Tim Hallc8310b12020-06-17 14:53:11 +0100322
323 if len(scratch_tensors) == 0:
324 scratch_tensor = None
325 else:
326 assert len(scratch_tensors) == 1, "Multiple scratch tensors"
327 scratch_tensor = scratch_tensors[0]
328
Tim Hall79d07d22020-04-27 18:20:16 +0100329 self.tensor_map = {tens: idx for idx, tens in enumerate(all_tensors)}
Tim Hallc8310b12020-06-17 14:53:11 +0100330 self.buffer_map = self.assign_buffers_to_tensors(all_tensors, scratch_tensor)
Tim Hall79d07d22020-04-27 18:20:16 +0100331
332 tensors_offset = self.write_offset_vector([self.serialise_tensor(tens) for tens in all_tensors])
333
Tim Hall79d07d22020-04-27 18:20:16 +0100334 # Make sure the input_tensors haven't been modified
335 assert all(inp in sg.original_inputs for inp in sg.input_tensors)
Tim Hallc8310b12020-06-17 14:53:11 +0100336 inputs = [self.tensor_map[tens] for tens in sg.original_inputs]
337
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200338 # Add the Scratch Tensors as input to the NPU subgraph to get them allocated by TensorFlow Lite Micro
Tim Hallc8310b12020-06-17 14:53:11 +0100339 scratch_tensor_idx = self.tensor_map.get(scratch_tensor, None)
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200340 scratch_fast_tensor_idx = self.tensor_map.get(scratch_fast_tensor, None)
341
Tim Hallc8310b12020-06-17 14:53:11 +0100342 if scratch_tensor_idx is not None and scratch_tensor_idx not in inputs:
343 inputs.append(scratch_tensor_idx)
344
Patrik Gustavsson3ab94522020-06-29 17:36:55 +0200345 if scratch_fast_tensor_idx is not None and scratch_fast_tensor_idx not in inputs:
346 inputs.append(scratch_fast_tensor_idx)
347
Tim Hallc8310b12020-06-17 14:53:11 +0100348 inputs_offset = self.write_int_vector(inputs)
Tim Hall79d07d22020-04-27 18:20:16 +0100349 outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in sg.output_tensors])
350
351 operators_offset = self.write_offset_vector([self.serialise_operator(op) for op in all_ops])
352
353 SubGraph.SubGraphStart(builder)
354 SubGraph.SubGraphAddTensors(builder, tensors_offset)
355 SubGraph.SubGraphAddInputs(builder, inputs_offset)
356 SubGraph.SubGraphAddOutputs(builder, outputs_offset)
357
358 SubGraph.SubGraphAddOperators(builder, operators_offset)
359
360 return SubGraph.SubGraphEnd(builder)
361
362 def write_aligned_bytes(self, buf):
363 builder = self.builder
364 builder.nested = True
365 data = bytes(buf)
366 length_bytes = UOffsetTFlags.py_type(len(data))
367 builder.Prep(16, length_bytes) # Reserve aligned storage
368 builder.head = UOffsetTFlags.py_type(builder.Head() - length_bytes) # Update FlatBuffer internal pointer
369 builder.Bytes[builder.Head() : builder.Head() + length_bytes] = data # Assign bytes to aligned area
370 return builder.EndVector(length_bytes)
371
372 def serialise_buffer(self, buf):
373 builder = self.builder
374 data = None
375 if buf is not None:
376 data = self.write_aligned_bytes(buf)
377 Buffer.BufferStart(builder)
378 if data is not None:
379 Buffer.BufferAddData(builder, data)
380 return Buffer.BufferEnd(builder)
381
382 def serialise_metadata(self, metadata):
383 builder = self.builder
384 name = builder.CreateString(metadata[0])
385
386 Metadata.MetadataStart(builder)
387 Metadata.MetadataAddName(builder, name)
388 Metadata.MetadataAddBuffer(builder, metadata[1])
389
390 return Metadata.MetadataEnd(builder)
391
392 def serialise_model(self):
393 builder = self.builder
394 operator_code_offset = self.write_offset_vector(
395 [self.serialise_operator_code(idx, code) for idx, code in enumerate(self.operator_codes)]
396 )
397
398 description = builder.CreateString("Vela Optimised")
399
400 subgraph_offset = self.write_offset_vector([self.serialise_subgraph(sg) for sg in self.subgraphs_to_write])
401
402 # Fill the metadata buffer
403 version = np.int32(0)
404 subgraph_idx = np.int32(len(self.subgraphs_to_write)) # Only 1 supported currently
405 nbr_tensors = np.int32(len(self.tensor_map))
406
407 # An offset of -1 indicates that the tensor will be allocated online by Tensorflow Lite Micro
408 offsets = [np.int32(-1)] * nbr_tensors
409
410 # Ensure that the order of the offsets match the order of the tensors
411 for tens, idx in self.tensor_map.items():
Patrik Gustavssoneca2e952020-05-27 09:15:11 +0200412 # Set offsets for tensor allocated in Tensor Arena or in the scratch_fast area
Charles Xu04ce34c2020-06-23 12:42:28 +0200413 if tens.mem_type in set((MemType.Scratch, MemType.Scratch_fast)) and tens.address is not None:
Tim Hall79d07d22020-04-27 18:20:16 +0100414 offsets[idx] = np.int32(tens.address)
415
416 metadata_buffer = np.array([version, subgraph_idx, nbr_tensors] + offsets)
417 self.buffers_to_write.append(metadata_buffer)
418
419 buffers_offset = self.write_offset_vector([self.serialise_buffer(buf) for buf in self.buffers_to_write])
420
421 metadata_list = [("OfflineMemoryAllocation", len(self.buffers_to_write) - 1)]
422 metadata_offset = self.write_offset_vector([self.serialise_metadata(metadata) for metadata in metadata_list])
423
424 Model.ModelStart(builder)
425 Model.ModelAddVersion(builder, tflite_version)
426 Model.ModelAddOperatorCodes(builder, operator_code_offset)
427 Model.ModelAddSubgraphs(builder, subgraph_offset)
428 Model.ModelAddDescription(builder, description)
429 Model.ModelAddBuffers(builder, buffers_offset)
430 Model.ModelAddMetadata(builder, metadata_offset)
431 return Model.ModelEnd(builder)
432
433 def serialise(self):
434
435 model = self.serialise_model()
436
437 self.builder.FinishWithFileIdentifier(model, tflite_file_identifier)
438
439 return self.builder.Output()
440
441 def write(self, filename):
442 with open(self.filename, "wb") as f:
443 f.write(self.serialised_buf)
444
445
446def write_tflite(nng, filename):
447 writer = TFLiteSerialiser(nng)
448 buf = writer.serialise()
449
450 with open(filename, "wb") as f:
451 f.write(buf)