diff options
Diffstat (limited to 'src/runtime/NEON/functions/NEPoolingLayer.cpp')
| -rw-r--r-- | src/runtime/NEON/functions/NEPoolingLayer.cpp | 47 |
1 files changed, 36 insertions, 11 deletions
diff --git a/src/runtime/NEON/functions/NEPoolingLayer.cpp b/src/runtime/NEON/functions/NEPoolingLayer.cpp index bc0b6f86d3..cbfd68485f 100644 --- a/src/runtime/NEON/functions/NEPoolingLayer.cpp +++ b/src/runtime/NEON/functions/NEPoolingLayer.cpp @@ -31,7 +31,7 @@ using namespace arm_compute; NEPoolingLayer::NEPoolingLayer() - : _pooling_layer_kernel(), _border_handler(), _is_global_pooling_layer(false) + : _pooling_layer_kernel(), _border_handler(), _is_global_pooling_layer(false), _data_layout(DataLayout::NCHW) { } @@ -40,17 +40,31 @@ void NEPoolingLayer::configure(ITensor *input, ITensor *output, const PoolingLay // Check if we have Global Pooling Layer _is_global_pooling_layer = (input->info()->dimension(0) == pool_info.pool_size().width) && (input->info()->dimension(1) == pool_info.pool_size().height); + // Get data layout + _data_layout = input->info()->data_layout(); + // Configure pooling kernel _pooling_layer_kernel.configure(input, output, pool_info); - // Configure border depending on operation required (quantize border in case of asymmetric data_type) - BorderMode border_mode = (pool_info.pool_type() == PoolingType::MAX) ? BorderMode::REPLICATE : BorderMode::CONSTANT; - PixelValue zero_value(0.f); - if(is_data_type_quantized_asymmetric(input->info()->data_type()) && !pool_info.exclude_padding()) + switch(_data_layout) { - zero_value = PixelValue(static_cast<uint32_t>(input->info()->quantization_info().offset)); + case DataLayout::NCHW: + { + // Configure border depending on operation required (quantize border in case of asymmetric data_type) + BorderMode border_mode = (pool_info.pool_type() == PoolingType::MAX) ? BorderMode::REPLICATE : BorderMode::CONSTANT; + PixelValue zero_value(0.f); + if(is_data_type_quantized_asymmetric(input->info()->data_type()) && !pool_info.exclude_padding()) + { + zero_value = PixelValue(static_cast<uint32_t>(input->info()->quantization_info().offset)); + } + _border_handler.configure(input, _pooling_layer_kernel.border_size(), border_mode, zero_value); + break; + } + case DataLayout::NHWC: + break; + default: + ARM_COMPUTE_ERROR("Data layout not supported"); } - _border_handler.configure(input, _pooling_layer_kernel.border_size(), border_mode, zero_value); } Status NEPoolingLayer::validate(const ITensorInfo *input, const ITensorInfo *output, const PoolingLayerInfo &pool_info) @@ -60,9 +74,20 @@ Status NEPoolingLayer::validate(const ITensorInfo *input, const ITensorInfo *out void NEPoolingLayer::run() { - // Fill border - NEScheduler::get().schedule(&_border_handler, Window::DimY); + switch(_data_layout) + { + case DataLayout::NCHW: + // Fill border + NEScheduler::get().schedule(&_border_handler, Window::DimY); - // Run pooling layer - NEScheduler::get().schedule(&_pooling_layer_kernel, _is_global_pooling_layer ? Window::DimZ : Window::DimY); + // Run pooling layer + NEScheduler::get().schedule(&_pooling_layer_kernel, _is_global_pooling_layer ? Window::DimZ : Window::DimY); + break; + case DataLayout::NHWC: + // Run pooling layer + NEScheduler::get().schedule(&_pooling_layer_kernel, Window::DimX); + break; + default: + ARM_COMPUTE_ERROR("Data layout not supported"); + } }
\ No newline at end of file |