Skip to content

[tmva][sofie] AveragePool with ceil_mode divides the partial/overhang window by the full kernel size #22523

@harz05

Description

@harz05

Check duplicate issues.

  • Checked for duplicates

Description

When an ONNX AveragePool node has ceil_mode=1, the last sliding window can extend past the end of the input (the "overhang" window that ceil_mode is meant to keep). SOFIE's generated code averages such a window by dividing the sum of the real cells by the full kernel size, instead of by the number of cells actually present. The result is wrong for that window.

The divisor in ROperator_Pool.hxx is computed at run time only when count_include_pad == 0 && doPadding. Otherwise it is the compile-time constant kernel area (kh, kw*kh, ...). A window can also be partial because of ceil_mode, which this condition does not cover, so:

  • count_include_pad = 0 (the default) with no padding + ceil_mode -> constant divisor, wrong.
  • count_include_pad = 1 + ceil_mode -> constant divisor includes the overhang cells, wrong.

Example (1D): input [1,2,3,4,5,6], kernel 3, stride 2, no padding, ceil_mode = 1. Output has 3 windows; the third covers input indices [4,5] (index 6 is the overhang).

  • Expected (PyTorch avg_pool1d, count_include_pad=False; same as the ONNX definition): [2.0, 4.0, 5.5] ((5+6)/2 = 5.5)
  • SOFIE generated code: [2.0, 4.0, 3.6667] ((5+6)/3 = 3.6667)

The first two windows are full so they match; only the overhung window is wrong.

Reproducer

Three small steps to reproduce-

  1. Create the model (Python, needs the onnx package):
##make_avgpool_ceil.py
import onnx
from onnx import helper, TensorProto
node = helper.make_node('AveragePool', ['x'], ['y'],
                        kernel_shape=[3], strides=[2], ceil_mode=1)
g = helper.make_graph([node], 'avgpool_ceil',
    [helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 1, 6])],
    [helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 1, 3])])
m = helper.make_model(g, opset_imports=[helper.make_opsetid('', 14)])
m.ir_version = 7
onnx.save(m, 'avgpool_ceil.onnx')

Run: python3 make_avgpool_ceil.py

  1. Generate the SOFIE code (gen.C):
// gen.C
void gen() {
   using namespace TMVA::Experimental::SOFIE;
   RModelParser_ONNX parser;
   RModel model = parser.Parse("avgpool_ceil.onnx");
   model.Generate();
   model.OutputGenerated("avgpool_ceil.hxx");
}

Run: root -l -b -q gen.C

  1. Run inference and print the output (run.C):
// run.C
#include "avgpool_ceil.hxx"
void run() {
   std::vector<float> x = {1, 2, 3, 4, 5, 6};
   TMVA_SOFIE_avgpool_ceil::Session s;
   auto y = s.infer(x.data());
   for (float v : y) std::cout << v << " ";
   std::cout << std::endl;
}

Run: root -l -b -q run.C
Prints: 2 4 3.66667 (expected: 2 4 5.5)

(If the generated namespace differs, check the top of avgpool_ceil.hxx and adjust
TMVA_SOFIE_avgpool_ceil accordingly.)

ROOT version

6.41.01

Installation method

Built from source

Operating system

Ubuntu 22.04.2 LTS

Additional context

No response

Metadata

Metadata

Type

No type
No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions