Average pooling clamped divisor should be done on all conditions where the kernel can go out of bounds #4144
Conversation
…e the kernel can go out of bounds.
…rnel/stride/padding elements have to be processed in reversed order relative to the spatial dimensions.
|
@ivangarcia44 This is not a good practice to close and open a new PR just because conflicts have arisen in the branch. Ideally, you should have addressed the conflicts and updated this PR just like other contributors do. At the very least, it creates confusion and makes it hard to follow the discussion on 2 different PR for the same changes. Also, I would expect that this PR (or the new replacement of this) is not merged until a thorough review of the patch is done. The last PR had already made some changes which should not have been done, so I expect all to be more careful this time. |
|
I will do the merge in this PR and close the new one. |
|
Hi all, updated the changes after merging the conflict with @vivekkhandelwal1 's changes. Please review when you get a chance. Thank you @vivekkhandelwal1 |
| // | ||
| // indexStartOffset = ceil((kernelSize - 1)/2) - padding | ||
| // | ||
| // clampedKernelSize = | ||
| // min(outIntIndex * stride + indexStartOffset + floor((kernelSize - 1)/2) | ||
| // + 1, | ||
| // InputSpatialDimSize + padding) - | ||
| // max(outIntIndex * stride + indexStartOffset - ceil((kernelSize - 1)/2), | ||
| // -padding) | ||
| // |
There was a problem hiding this comment.
Hi @ivangarcia44, can you please add any link to the code/implementation in PyTorch from where you adopted this?
There was a problem hiding this comment.
Hi Vivek, the original author of the torch-linalg avg_pool2d lowering (@AmosLewis ) based the 2D implementation of the divisor computation on this PyTorch code:
For the new E2E test suite I added, the test AvgPool2dCeilPaddingStridedIncludePadding discovered that there was a numerical mismatch due to the divisor computation. Because of this I wrote a new average pooling divisor formula based on the PyTorch avg_pool2d doc and numerical values I observed. I tried to make the comments, variable names and code clearer than the previous version which helped me to correct the numerical issues I found and make all E2E tests pass.
Taking a second look at the PyTorch code, and the @AmosLewis code, I noticed I missed a detail in the 1D/3D generalization of the algorithm: 1) count_include_pad parameter used in the divisor computation 2) Reversing of kernel/stride/pad parameter elements order. Although the 2nd issue could have been present in the original code. The second issue was found by 3 new E2E tests from this PR with non-uniform values for these parameters.
I updated the PR to use the PyTorch based algorithm with these two fixes and all E2E tests are passing now.
|
Hi all, this is a reminder to review when you get a chance. Thank you @vivekkhandelwal1 |
|
Could you please take a look at this PR when you have a chance? I’ve addressed all the previous feedback, and GitHub requires your approval before it can be merged. This PR includes three numerical correctness bug fixes for the average pooling operator, which are important for our stakeholders working on safety-critical applications. It’s been open for a month, so I wanted to check in and see if there’s anything else needed from my side to move this forward. Thank you very much! @silvasean |
Hi @ivangarcia44, I'm sorry for the delay. I got caught up in some other work. Can I review it in a couple of days? |
Its ok, next week is perfectly fine. Thank you for following up! |
There was a problem hiding this comment.
@ivangarcia44, thanks for the changes and for adding the clarifications. The PR looks good. If you want, then you may incorporate the comment suggestion before merging the PR.
I have 2 doubts:
- I am unable to understand the need for reversing the kernel/stride/padding element order.
- In response to this comment of mine, you said:
Because of this I wrote a new average pooling divisor formula based on the PyTorch avg_pool2d doc and numerical values I observed.
Honestly, as per my opinion doing something based on observation won't be correct. There must be some rationale/reference for the same. Also. now I don't see that comment in the code.
Committing Vivek's comment update suggestion. Co-authored-by: Vivek Khandelwal <[email protected]>
Hi Vivek, I will integrate your requested comment change. Here are my answers for the two questions above:
I will submit a commit for point number 1 above. Please let me know if I missed anything. Thanks, |
In this pull request I added various E2E tests to cover edge cases on the average pooling torch to linalg lowering algorithm. These new tests uncovered various numerical issues that were addressed in this same PR.
One of the issues is the IREE test failure found in #4079 which triggered this work.
Background:
In the most common case, the divisor for the average pooling is just the product of kernel dimensions. But with padding, and ceil mode options, some elements need to be discounted from the divisor computation. This change fixes two components of this:
Fix the condition that determines if the divisor is just the product of kernel dimensions or the clamped divisor comptutation.
Add missing isCountIncludePad logic divisor computation algorithm and reversal of kernel/stride/padding parameters element order. Both were missing from the first generalization change.
@vivekkhandelwal1
@silvasean
@zjgarvey
@ramiro050
@JianzheXiao
@AmosLewis
@rsuderman
@nirvedhmeshram
@sahas3
@Hanumanth04
@dixinzhou
@rafaelubalmw