Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: support adaptive avg pool 2d and 3d dynamo converters #2632

Merged
merged 5 commits into from
Apr 12, 2024
Merged

feat: support adaptive avg pool 2d and 3d dynamo converters #2632

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
db11ce6
WIP: need to fix the bug that output values are incorrect sometimes. …
zewenli98 Jan 21, 2024
e2e0aa5
feat: support adaptive avg pool2d and pool3d dynamo converters
zewenli98 Feb 1, 2024
cbc09ae
refactor to separate 2d and 3d functions for maintainability, fix bug…
zewenli98 Apr 8, 2024
61b74f3
combine two pools
zewenli98 Apr 9, 2024
f1b1ad0
update tests and enforce input type
zewenli98 Apr 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 39 additions & 1 deletion py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2185,7 +2185,7 @@ def aten_ops_avg_pool(


@dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool1d.default)
def aten_ops_adaptive_avg_pool(
def aten_ops_adaptive_avg_pool1d(
ctx: ConversionContext,
target: Target,
args: Tuple[Argument, ...],
Expand All @@ -2202,6 +2202,44 @@ def aten_ops_adaptive_avg_pool(
)


@dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool2d.default)
@dynamo_tensorrt_converter(torch.ops.aten._adaptive_avg_pool2d.default)
def aten_ops_adaptive_avg_pool2d(
ctx: ConversionContext,
target: Target,
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return impl.pool.adaptive_avg_pool2d(
ctx,
target,
source_ir=SourceIR.ATEN,
name=name,
input=args[0],
output_size=args[1],
)


@dynamo_tensorrt_converter(torch.ops.aten.adaptive_avg_pool3d.default)
@dynamo_tensorrt_converter(torch.ops.aten._adaptive_avg_pool3d.default)
def aten_ops_adaptive_avg_pool3d(
ctx: ConversionContext,
target: Target,
args: Tuple[Argument, ...],
kwargs: Dict[str, Argument],
name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
return impl.pool.adaptive_avg_pool3d(
ctx,
target,
source_ir=SourceIR.ATEN,
name=name,
input=args[0],
zewenli98 marked this conversation as resolved.
Show resolved Hide resolved
output_size=args[1],
)


def max_pool_param_validator(pool_node: Node) -> bool:
dilation = args_bounds_check(pool_node.args, 4, 1)
ceil_mode = args_bounds_check(pool_node.args, 5, False)
Expand Down
Loading
Loading