MLBEDSW-7949: [MLCE] Remove duplicate cpu tensors
- If a npu op is followed by a convolution op than runs on the cpu,
the optimized file ends up containing a duplicated tensor called _cpu.
Functionality wise not a problem but the graph will look strange in a
graph viewer.
- This error was introduced when removing duplicate weights
tensors but the above use case was not considered in that patch.
- The fix is to make sure that only the weight and bias tensor are
modified.
Change-Id: I576f13650f1f9d3d50a421ab7100fc8b5ab62657
Signed-off-by: Johan Alfven <johan.alfven@arm.com>
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index 5aa38a7..6d85567 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -103,10 +103,10 @@
self.align_nng_inputs_to_tflite(op)
all_ops.append(op)
if op.type.is_conv2d_op() or op.type.is_depthwise_conv2d_op() or op.type == Op.FullyConnected:
- # Op is run on CPU, make sure original tensor are written back
+ # Op is run on CPU, make sure the original weight and bias tensors are written back
# instead of the cloned/reshaped (see tflite_reader)
for idx, inp in enumerate(op.inputs):
- if inp is not None and inp.src_tensor is not None:
+ if inp != op.ifm and inp is not None and inp.src_tensor is not None:
op.inputs[idx] = inp.src_tensor
# list of tuple(Op, string, op.version); the custom code is only used for 3rd party custom operators