feat: support 1D, 2D, and 3D avg and max pooling dynamo converters - #2317
Conversation
github-actions
Bot
added
component: api [Python]
gs-olive
left a comment
gs-olive
left a comment
There was a problem hiding this comment.
Left a few code comments, otherwise looks good to me
| if stride is None: | ||
| stride = kernel_size | ||
| else: |
There was a problem hiding this comment.
By the function documentation and the schema, it should not be allowed for stride to be None, so this can be removed
There was a problem hiding this comment.
Thanks for pointing it out! According to the schema, the stride could be [] by default, so I made changes accordingly.
| @@ -0,0 +1,106 @@ | |||
| from typing import Optional, Sequence, Union | |||
|
|
|||
| # @manual=//deeplearning/trt/python:py_tensorrt | |||
| from torch.fx.node import Target | ||
| from torch_tensorrt.dynamo.conversion.converter_utils import extend_attr_to_tuple | ||
| from torch_tensorrt.fx.converters.converter_utils import ( | ||
| SourceIR, |
There was a problem hiding this comment.
Switch to using the torch_tensorrt.dynamo version of SourceIR
| if stride is None: | ||
| stride = kernel_size | ||
| else: |
There was a problem hiding this comment.
See above comment (can remove)
|
|
||
|
|
||
| # Note: MaxPool1d uses max_pool2d as it converts to 2D first. | ||
| @dynamo_tensorrt_converter(torch.ops.aten.max_pool2d.default, capability_validator=max_pool_param_validator) # type: ignore[misc] |
There was a problem hiding this comment.
Could max_pool1d support be added here as well? Schema
There was a problem hiding this comment.
I can add torch.ops.aten.max_pool1d.default but it won't be used. Even for torch.nn.AvgPool1d, it still calls torch.ops.aten.avg_pool2d.default, as you can see in the test file: https://github.com/pytorch/TensorRT/pull/2317/files#diff-9fce39bc42c66d2866c41665779cab7da0a4d3fe54576925e2b66c17a1cf1ebfR20-R43
But anyways, the 1d schema looks same as others, so I added here.
There was a problem hiding this comment.
Thanks for that - I plan to add a lowering pass which will lead us to that converter, so it will still be helpful.
…/TensorRT into pool_dynamo_converters
gs-olive
left a comment
gs-olive
left a comment
There was a problem hiding this comment.
Looks good to me, pending CI!
4 participants
Description
Support 1D, 2D, and 3D avg and max pooling dynamo converters.
Fixes #2200
Type of change
Checklist: