diff options
author | Michalis Spyrou <michalis.spyrou@arm.com> | 2019-01-02 15:54:03 +0000 |
---|---|---|
committer | Michalis Spyrou <michalis.spyrou@arm.com> | 2019-01-08 13:11:25 +0000 |
commit | 8d1b7186ae18944678be04694d51ce7c0c9c9aa9 (patch) | |
tree | 0c9f372c7e488b59b999e1fcd0df7d5226ccd593 /src/runtime/NEON | |
parent | 5a97b28d483eefa5810f5cf57356086090c8c894 (diff) | |
download | ComputeLibrary-8d1b7186ae18944678be04694d51ce7c0c9c9aa9.tar.gz |
COMPMID-1865 NEReduceMean fails on shape validation
Also handle negative axis
Change-Id: I28e48702d926c2f4aea7b1b674b51bebb01ce5f8
Reviewed-on: https://review.mlplatform.org/464
Reviewed-by: Matthew Bentham <matthew.bentham@arm.com>
Reviewed-by: Isabella Gottardi <isabella.gottardi@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Diffstat (limited to 'src/runtime/NEON')
-rw-r--r-- | src/runtime/NEON/functions/NEReduceMean.cpp | 55 |
1 files changed, 45 insertions, 10 deletions
diff --git a/src/runtime/NEON/functions/NEReduceMean.cpp b/src/runtime/NEON/functions/NEReduceMean.cpp index c5229daae3..dc610d55de 100644 --- a/src/runtime/NEON/functions/NEReduceMean.cpp +++ b/src/runtime/NEON/functions/NEReduceMean.cpp @@ -39,16 +39,37 @@ Status NEReduceMean::validate(const ITensorInfo *input, const Coordinates &reduc ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input); ARM_COMPUTE_RETURN_ERROR_ON(reduction_axis.num_dimensions() > input->num_dimensions()); - for(unsigned int i = 0; i < reduction_axis.num_dimensions(); ++i) + TensorShape out_shape = input->tensor_shape(); + const unsigned int reduction_ops = reduction_axis.num_dimensions(); + const int input_dims = input->num_dimensions(); + Coordinates axis_local = reduction_axis; + + // Convert negative axis + for(unsigned int i = 0; i < reduction_ops; ++i) + { + axis_local[i] = wrap_around(axis_local[i], input_dims); + } + + std::sort(axis_local.begin(), axis_local.begin() + reduction_ops); + for(unsigned int i = 0; i < reduction_ops; ++i) { - if(output->total_size() > 0) + ARM_COMPUTE_RETURN_ERROR_ON(axis_local[i] > 3); + ARM_COMPUTE_RETURN_ERROR_ON(static_cast<unsigned int>(axis_local[i]) > input->num_dimensions() - 1); + if(output->total_size() > 0 && keep_dims) { - ARM_COMPUTE_RETURN_ERROR_ON(output->dimension(reduction_axis[i]) != 1); - ARM_COMPUTE_RETURN_ERROR_ON(static_cast<unsigned int>(reduction_axis[i]) > input->num_dimensions() - 1); + ARM_COMPUTE_RETURN_ERROR_ON(output->dimension(axis_local[i]) != 1); + } + if(keep_dims) + { + out_shape.set(axis_local[i], 1); + } + else + { + out_shape.remove_dimension(axis_local[i] - i); } - - ARM_COMPUTE_RETURN_ON_ERROR(NEReductionOperationKernel::validate(input, output, reduction_axis[i], ReductionOperation::MEAN_SUM)); } + const TensorInfo out_info = input->clone()->set_tensor_shape(out_shape); + ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_SHAPES(output, &out_info); return Status{}; } @@ -62,22 +83,32 @@ void NEReduceMean::configure(ITensor *input, const Coordinates &reduction_axis, _reduced_outs = arm_compute::support::cpp14::make_unique<Tensor[]>(_reduction_ops - (keep_dims ? 1 : 0)); _keep_dims = keep_dims; + Coordinates axis_local = reduction_axis; + const int input_dims = input->info()->num_dimensions(); + const unsigned int reduction_ops = reduction_axis.num_dimensions(); + + // Convert negative axis + for(unsigned int i = 0; i < reduction_ops; ++i) + { + axis_local[i] = wrap_around(axis_local[i], input_dims); + } + // Perform reduction for every axis for(unsigned int i = 0; i < _reduction_ops; ++i) { TensorShape out_shape = i == 0 ? input->info()->tensor_shape() : (_reduced_outs.get() + i - 1)->info()->tensor_shape(); - out_shape.set(reduction_axis[i], 1); + out_shape.set(axis_local[i], 1); auto in = (i == 0) ? input : (_reduced_outs.get() + i - 1); if(i == _reduction_ops - 1 && keep_dims) { - _reduction_kernels[i].configure(in, output, reduction_axis[i], ReductionOperation::MEAN_SUM); + _reduction_kernels[i].configure(in, output, axis_local[i], ReductionOperation::MEAN_SUM); } else { _reduced_outs[i].allocator()->init(TensorInfo(out_shape, input->info()->num_channels(), input->info()->data_type())); _memory_group.manage(_reduced_outs.get() + i); - _reduction_kernels[i].configure(in, _reduced_outs.get() + i, reduction_axis[i], ReductionOperation::MEAN_SUM); + _reduction_kernels[i].configure(in, _reduced_outs.get() + i, axis_local[i], ReductionOperation::MEAN_SUM); } } @@ -91,9 +122,13 @@ void NEReduceMean::configure(ITensor *input, const Coordinates &reduction_axis, if(!keep_dims) { TensorShape out_shape = input->info()->tensor_shape(); + + // We have to sort the reduction axis vectors in order for remove_dimension + // to work properly + std::sort(axis_local.begin(), axis_local.begin() + _reduction_ops); for(unsigned int i = 0; i < _reduction_ops; ++i) { - out_shape.remove_dimension(reduction_axis[i]); + out_shape.remove_dimension(axis_local[i] - i); } auto_init_if_empty(*output->info(), input->info()->clone()->set_tensor_shape(out_shape)); _reshape.configure(_reduced_outs.get() + _reduction_ops - 1, output); |