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 <a.tischenko@partner.samsung.com>

diff --git a/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp b/contrib/nnc/passes/onnx_frontend/ONNXImporterImpl.cpp
index a2f0cad..a453471 100644 (file)
--- 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<onnx::ModelProto>& 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 (file)
--- a/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp
+++ b/contrib/nnc/passes/onnx_frontend/ONNXOpCreator.cpp
@@ -131,18 +131,25 @@ std::vector<Operation*> ONNXOpCreator::convertConcat(InputOps& inputs,
 }
 
 std::vector<Operation*> 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<Operation*> result;
+  std::vector<int32_t> padding_before{0, 0};
+  std::vector<int32_t> 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<ops::PoolOp>(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<Operation*> 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<int32_t> padding_before{0, 0};
-  std::vector<int32_t> padding_after{0, 0};
-  // TODO: ONNX has more parameters for pooling. We should use them.
-  auto pooling = createOp<ops::PoolOp>(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<ops::PoolOp>(inputs[0]->getOutput(0), pool_type, kernel_shape, strides_shape,
+                                 padding_before, padding_after, border_type,
+                                 ops::PoolOp::RoundMode::floor);
+  return result;
 }
 
 std::vector<Operation*> ONNXOpCreator::convertSoftmax(InputOps& inputs,