Skip to content

[torchlib] Support divisor_override in avg_pool2d and avg_pool3d - #3070

Open
Raashish Aggarwal (raashish1601) wants to merge 1 commit into
microsoft:mainfrom
raashish1601:feature/avg-pool-divisor-override
Open

Raashish Aggarwal (raashish1601) wants to merge 1 commit into
microsoft:mainfrom
raashish1601:feature/avg-pool-divisor-override

Conversation

@raashish1601

Copy link
Copy Markdown

aten_avg_pool2d and aten_avg_pool3d accepted divisor_override but ignored it (there was a TODO for it), so the exported model silently divided by the window size instead:

x = torch.arange(16.0).reshape(1, 1, 4, 4)
F.avg_pool2d(x, 2, 2, divisor_override=1)  # [10, 18, 42, 50]
# exported model returns [2.5, 4.5, 10.5, 12.5]

AveragePool can't take a custom divisor, so with divisor_override the window sums are computed with LpPool(p=1) on the positive and negative parts (sum = LpPool(relu(x)) - LpPool(relu(-x))) and divided by divisor_override. Padding contributes zeros and LpPool supports ceil_mode, so this matches PyTorch, which divides the sum of the in-bounds values by divisor_override regardless of count_include_pad. Without divisor_override the graph is unchanged.

Testing:

  • Removed the divisor_override xfails for nn.functional.avg_pool2d / avg_pool3d, and excluded divisor_override samples from the avg_pool3d ceil_mode xfail, since they no longer go through AveragePool. pytest tests/function_libs/torch_lib/ops_test.py -k avg_pool passes with onnxruntime 1.23.0 (the version in noxfile.py): 5 passed, 33 subtests passed, 7 xfailed.
  • I also compared LpPool-based sums with F.avg_pool2d(..., divisor_override=3) over 96 combinations of input size, kernel, stride, padding, ceil_mode and count_include_pad on onnxruntime 1.23.0 and 1.30.0. All match. Exporting nn.AvgPool2d(..., divisor_override=...) and 3d/unbatched cases with torch.onnx.export(dynamo=True) also matches eager.
  • ruff check and format are clean on the changed files.

divisor_override was ignored, so the exported model divided by the window size instead. Compute the window sums with LpPool (p=1) over the positive and negative parts and divide by divisor_override, which also handles padding and ceil_mode.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

1 participant