Using max pooling kernels without computing indices

1

Hi PyTorch forum,

I’m trying to address a bottleneck in some code of mine where a big F.max_pool1d is being applied (very long input length and long kernel length of 121). I noticed that this was a bottleneck from running in VizTracer, where I also learned that internally it’s calling into an ATen max_pool2d_with_indices. I’m running on a GPU (CUDA).

I don’t need gradients or anything, so I was wondering if it would be possible to avoid the indices. I noticed some merged PRs in PyTorch relating to implementing no-indices optimizations (like MaxPool1d without indices optimization by heitorschueroff · Pull Request #43745 · pytorch/pytorch · GitHub), and I was curious if it’s possible for end users to access those code paths from Python. I’ve tried inference_mode and set_grad_enabled(False), but I still see the with_indices function in the call stack.

Just curious if that’s possible – thanks!

2

AFAICT, there is currently no native implementation without indices (both adaptive: pytorch/aten/src/ATen/native/cuda/AdaptiveMaxPooling2d.cu at main · pytorch/pytorch · GitHub and regular : pytorch/aten/src/ATen/native/cuda/DilatedMaxPool2d.cu at main · pytorch/pytorch · GitHub ).

So you cannot call into this implementation without the indices tensor, regardless at what API level.