axial-attention:实施轴向注意力-有效地处理多维数据-源码

轴向注意 在Pytorch中实施。 一种简单而强大的技术，可以有效处理多维数据。 它为我和许多其他研究人员创造了奇迹。 只需在数据中添加一些位置编码，然后将其传递到此方便的类中，即可指定要嵌入的尺寸以及要旋转的轴向尺寸。 所有的排列，整形，都将为您解决。 实际上，这篇论文由于过于简单而被拒绝了。 然而，自那以后，它已成功用于许多应用中，包括， 。 只是去展示。 安装 $ pip install axial_attention 用法 图像 import torch from axial_attention import AxialAttention img = torch . randn ( 1 , 3 , 256 , 256 ) attn = AxialAttention ( dim = 3 , # embedding dimension