diff options
author | Alexander Hansson <Alexander.Hansson@arm.com> | 2023-06-30 15:41:13 +0000 |
---|---|---|
committer | Alexander Hansson <alexander.hansson@arm.com> | 2023-07-11 11:02:10 +0100 |
commit | da8741a14c3774d3161f59019d3003a2ee944400 (patch) | |
tree | cf4aa63fd073d0a3b0f99b4ea6716b4a8f1be10f /ethosu/vela/tflite_graph_optimiser.py | |
parent | 1d5e859973ff18f3e4285f0ca04251ca246a182c (diff) | |
download | ethos-u-vela-da8741a14c3774d3161f59019d3003a2ee944400.tar.gz |
MLBEDSW-7653: Extend Mean support for depth axis
If any of H,W axes have shape 1, the IFM can be reshaped to support
reduction over the depth axis.
Signed-off-by: Alexander Hansson <Alexander.Hansson@arm.com>
Change-Id: I432ff1c399b7cee4ca5f0a8f4461e9c0a936d804
Diffstat (limited to 'ethosu/vela/tflite_graph_optimiser.py')
-rw-r--r-- | ethosu/vela/tflite_graph_optimiser.py | 35 |
1 files changed, 22 insertions, 13 deletions
diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py index a12eeb37..31d3ae1a 100644 --- a/ethosu/vela/tflite_graph_optimiser.py +++ b/ethosu/vela/tflite_graph_optimiser.py @@ -2004,16 +2004,9 @@ def convert_mean_to_depthwise_conv(op, arch, nng): intermediate_shape.insert(i, 1) # Reshape to 4D - if dims == 2: - # Reshape WxC -> 1xHxWx1 to support both axes - reduce_axis = [False] + reduce_axis + [False] - ifm_shape = [1] + ifm_shape + [1] - intermediate_shape = [1] + intermediate_shape + [1] - elif dims == 3: - # Reshape to 4D HxWxC -> 1xHxWxC - reduce_axis = [False] + reduce_axis - ifm_shape = [1] + ifm_shape - intermediate_shape = [1] + intermediate_shape + reduce_axis = full_shape(4, reduce_axis, False) + ifm_shape = full_shape(4, ifm_shape, 1) + intermediate_shape = full_shape(4, intermediate_shape, 1) # If all dimensions to reduce have shape 1, the operation is essentially a memcpy. # We can then remove the whole op by propagating ofm to previous ops @@ -2022,9 +2015,25 @@ def convert_mean_to_depthwise_conv(op, arch, nng): op = bypass_memory_only_ops(op, arch, nng) return op - # Compute kernel sizes for our convolutions. - # batch and depth axes are only supported if their shapes are 1. - # hence reduction in batch or depth axis is implicit. + # Support mean over depth-axis by left-shifting the C channel + # From semantics checks we can assume that one of H,W,C has shape 1 + if reduce_axis[3] and ifm_shape[3] > 1: + assert 1 in ifm_shape[1:], "Mean reduction over depth channel, but none of H,W,C has shape 1" + # If W=1 reshape NxHx1xC -> NxHxCx1, else reshape Nx1xWxC -> NxWxCx1 + idx_to_del = 2 if ifm_shape[2] == 1 else 1 + + # Delete axis with size 1 + del reduce_axis[idx_to_del] + del ifm_shape[idx_to_del] + del intermediate_shape[idx_to_del] + + # Add another element to set channel-axis to one + reduce_axis.append(False) + ifm_shape.append(1) + intermediate_shape.append(1) + + # Compute kernel sizes for our convolutions + # Batch axis is implicit as it is only supported if batch size is 1. h = ifm_shape[1] if reduce_axis[1] else 1 w = ifm_shape[2] if reduce_axis[2] else 1 |