Legacy Torch implementation for Dilated Neighborhood Attention (DiNAT)

[shi-labs]

Legacy Torch implementation for Dilated Neighborhood Attention (DiNAT)

[alihassanijr]

Hello and thank you for your interest.

We will look into it for future versions, but we strongly encourage everyone to set up NATTEN instead of using the torch-only implementation.
The reason for that is simply the pure torch version has a very large memory footprint and is quite slow, therefore it won't even scale up to larger tasks (see the chart below).

DiNA's torch implementation would be considerably slower and would use a lot more memory, because one of the key components that reduces speed is the extra post-padding step in that implementation, which attempts to pad an already massive tensor containing extracted sliding windows. Doing that for DiNA would be probably worse because dilation leads to more padding, and more padding, higher latency.

If you do run into any issues with installing NATTEN, please let us know.

[alihassanijr]

I'm closing this issue now because we're moving our extension to its own separate repository, and it is going to be pip-installable from now on.
We're still working on a CPU implementation, that is definitely on our list of features.
You are welcome to open an issue there if you'd like.
We actually encourage that, because the new NATTEN has dropped the Legacy pure-torch implementations, as we will soon have CPU implementations that will be much faster.

As for the pip installs, if you're using CUDA you can refer to our website and get a pip install command for your specific torch version.

pip install natten will eventually serve as our CPU-only build, like most other extension.

Thanks again for your interest.

[HerrYu123]

Hello,
Congratulations on your great work!
I would also like to know if there is a torch-based implementation for (Dilated) Neighborhood Attention. This would help me understand your work more concretely.

[alihassanijr]

Hi,
Thank you for your interest.
Unfortunately there is no pure torch implementation, and we are unsure whether it would work with the same replicated padding trick we used for NA.

We used to include a torch implementation of NA, but that was removed when we moved NATTEN and replaced with CPU kernels that run much more efficiently.

[HerrYu123]

Thanks a lot!