From a43f93a23234fc32b0684874b616956e1af6e895 Mon Sep 17 00:00:00 2001 From: =?utf8?q?=D0=90=D0=BD=D0=B4=D1=80=D0=B5=D0=B9=20=D0=A2=D0=B8=D1=89?= =?utf8?q?=D0=B5=D0=BD=D0=BA=D0=BE/AI=20Tools=20Lab=20/SRR/Staff=20Enginee?= =?utf8?q?r/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?= Date: Thu, 6 Dec 2018 18:04:37 +0300 Subject: [PATCH] Improved Pooling and GlobalAveragePool Implementation for ONNX (#2540) Some bugs in ONNX AveragePool and MaxPool were fixed. GlobalAveragePool for ONNX was supported as simple AveragePool with a set of default parameters. Signed-off-by: Andrew V. Tischenko --- .../nnc/passes/onnx_frontend/ONNXImporterImpl.cpp | 2 +- contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp | 60 +++++++++++++++------- 2 files changed, 42 insertions(+), 20 deletions(-) diff --git a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp index a2f0cad..a453471 100644 --- a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp +++ b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp @@ -50,6 +50,7 @@ static void collectUnsupportedOps(std::unique_ptr& model) { switch (ir_op_type->opCode) { case ONNXOpCode::opAdd: case ONNXOpCode::opAveragePool: + case ONNXOpCode::opGlobalAveragePool: case ONNXOpCode::opBatchNormalization: case ONNXOpCode::opConcat: case ONNXOpCode::opConv: @@ -228,7 +229,6 @@ mir::Graph *ONNXImporterImpl::createIR() { outputs = _opCreator.convertElementwise(input_nodes, mir::ops::ElementwiseOp::OpType::max); break; case ONNXOpCode::opGlobalAveragePool: - break; case ONNXOpCode::opAveragePool: case ONNXOpCode::opMaxPool: outputs = _opCreator.convertPool(input_nodes, onnx_op_type->opCode, onnx_node); diff --git a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp index d2cfb94..6d6faf1 100644 --- a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp +++ b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp @@ -131,18 +131,25 @@ std::vector ONNXOpCreator::convertConcat(InputOps& inputs, } std::vector ONNXOpCreator::convertPool(InputOps& inputs, ONNXOpCode op_code, - const onnx::NodeProto& onnx_node) { - auto* kshape = findAttribute(onnx_node, "kernel_shape"); - assert(kshape && kshape->ints_size()); - auto* strides = findAttribute(onnx_node, "strides"); - assert(strides && strides->ints_size()); - Shape onnx_kernel_shape = ShapeHelper::createShape(kshape->ints(), kshape->ints_size()); - Shape onnx_strides_shape = ShapeHelper::createShape(strides->ints(), strides->ints_size()); - + const onnx::NodeProto& onnx_node) { ops::PoolOp::BorderType border_type; ops::PoolOp::PoolingType pool_type; + std::vector result; + std::vector padding_before{0, 0}; + std::vector padding_after{0, 0}; + switch (op_code) { + case ONNXOpCode::opGlobalAveragePool: + // GlobalAveragePool is equivalent to AveragePool with kernel size equal + // to the spatial dimension of input tensor + return createOp(inputs[0]->getOutput(0), + ops::PoolOp::PoolingType::AVG, + inputs[0]->getOutputShape(0), // kernel_shape + Shape({1, 1}), // strides_shape + padding_before, padding_after,// no padding + ops::PoolOp::BorderType::ZEROFILLED, + ops::PoolOp::RoundMode::floor); case ONNXOpCode::opAveragePool: border_type = ops::PoolOp::BorderType::ZEROFILLED; pool_type = ops::PoolOp::PoolingType::AVG; @@ -154,17 +161,32 @@ std::vector ONNXOpCreator::convertPool(InputOps& inputs, ONNXOpCode default: assert(false); } - // FIXME: it's a hack to be compatible with current implementation of PoolOp - we expect - // 3 dims for 2D picture - Shape window_shape{onnx_kernel_shape.dim(0), onnx_kernel_shape.dim(1)}; - // FIXME: it's another hack identical to the above one - Shape strides_shape{onnx_strides_shape.dim(0), onnx_strides_shape.dim(1)}; - std::vector padding_before{0, 0}; - std::vector padding_after{0, 0}; - // TODO: ONNX has more parameters for pooling. We should use them. - auto pooling = createOp(inputs[0]->getOutput(0), pool_type, window_shape, - strides_shape, padding_before, padding_after, border_type); - return pooling; + // Proceed with Average or Max Pool + auto* kshape = findAttribute(onnx_node, "kernel_shape"); + assert(kshape && kshape->ints_size()); + auto* strides = findAttribute(onnx_node, "strides"); + assert(strides && strides->ints_size()); + auto* pads = findAttribute(onnx_node, "pads"); + + Shape kernel_shape = ShapeHelper::createShape(kshape->ints(), kshape->ints_size()); + Shape strides_shape = ShapeHelper::createShape(strides->ints(), strides->ints_size()); + + if (pads) { + assert(pads->ints_size() >= 2); + padding_before[0] = pads->ints(0); + padding_before[1] = pads->ints(1); + // TODO: ONNX padding could be for the beginning and ending along each axis that's why we + // should select the interesting ones. + if (pads->ints_size() == 4) { + padding_after[0] = pads->ints(2); + padding_after[1] = pads->ints(3); + } + + } + result = createOp(inputs[0]->getOutput(0), pool_type, kernel_shape, strides_shape, + padding_before, padding_after, border_type, + ops::PoolOp::RoundMode::floor); + return result; } std::vector ONNXOpCreator::convertSoftmax(InputOps& inputs, -- 2.7.4