From bb1b09ea438067d71d6dc3bbf55a3f6b41aaa75f Mon Sep 17 00:00:00 2001 From: Michael McGeagh Date: Wed, 19 Aug 2020 11:24:17 +0100 Subject: MLBEDSW-2783 Vela crashed on empty tflite file There may be cases where after optimisations, there are no operators contained within the subgraph. Upon serialising and writing out the vela optimised tflite file, it would crash for such a corner case. This fixes it allowing it to not crash but instead write out the empty tflite file. Signed-off-by: Michael McGeagh Change-Id: Ia879d1ffdbab21706b15e99aa107fb2d8d4dd3de --- ethosu/vela/tflite_writer.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) (limited to 'ethosu/vela') diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py index 92b5c6b0..8f523c7d 100644 --- a/ethosu/vela/tflite_writer.py +++ b/ethosu/vela/tflite_writer.py @@ -259,8 +259,10 @@ class TFLiteSerialiser: def serialise_operator(self, op): builder = self.builder - inputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.inputs]) - outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.outputs]) + inputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in op.inputs if tens in self.tensor_map]) + outputs_offset = self.write_int_vector( + [self.tensor_map[tens] for tens in op.outputs if tens in self.tensor_map] + ) op_idx, tflop, opt_serializer = self.operator_code_map[op.type] @@ -334,7 +336,7 @@ class TFLiteSerialiser: # Make sure the input_tensors haven't been modified assert all(inp in sg.original_inputs for inp in sg.input_tensors) - inputs = [self.tensor_map[tens] for tens in sg.original_inputs] + inputs = [self.tensor_map[tens] for tens in sg.original_inputs if tens in self.tensor_map] # Add the Scratch Tensors as input to the NPU subgraph to get them allocated by TensorFlow Lite Micro scratch_tensor_idx = self.tensor_map.get(scratch_tensor, None) @@ -347,7 +349,9 @@ class TFLiteSerialiser: inputs.append(scratch_fast_tensor_idx) inputs_offset = self.write_int_vector(inputs) - outputs_offset = self.write_int_vector([self.tensor_map[tens] for tens in sg.output_tensors]) + outputs_offset = self.write_int_vector( + [self.tensor_map[tens] for tens in sg.output_tensors if tens in self.tensor_map] + ) operators_offset = self.write_offset_vector([self.serialise_operator(op) for op in all_ops]) -- cgit v1.2.1