diff options
| author | Michael McGeagh <michael.mcgeagh@arm.com> | 2020-08-19 11:24:17 +0100 |
|---|---|---|
| committer | tim.hall <tim.hall@arm.com> | 2020-08-20 13:48:22 +0000 |
| commit | bb1b09ea438067d71d6dc3bbf55a3f6b41aaa75f (patch) | |
| tree | dfa7e64bf6cb72f366a1feaf6a5a38fb590508f2 /ethosu/vela/tflite_writer.py | |
| parent | a3aad93eec911f6bce69d7fb645d72ed82e59e16 (diff) | |
| download | ethos-u-vela-bb1b09ea438067d71d6dc3bbf55a3f6b41aaa75f.tar.gz | |
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 <michael.mcgeagh@arm.com>
Change-Id: Ia879d1ffdbab21706b15e99aa107fb2d8d4dd3de
Diffstat (limited to 'ethosu/vela/tflite_writer.py')
| -rw-r--r-- | ethosu/vela/tflite_writer.py | 12 |
1 files changed, 8 insertions, 4 deletions
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]) |