如何实现比 PyTorch 还要快的置换/转置操作？

作者：盐析白兔 | 2024-06-02 08:49:12

踩

介绍

在 NLP 领域占据主导地位的 Transformer 和 CV 领域的后起之秀 Vision Transformer 的模型中都可以看到 Transpose/Permute 操作。特别是在多头注意力中，需要这个操作来改变数据的维度排列。

显然，作为一个使用频率较高的操作，Transpose/Permute 操作的 CUDA 实现会影响实际网络的训练速度。本文将介绍 OneFlow 中 Permute Kernel 的优化技术，并在实验中与 PyTorch 的 Permute 和原生 Copy 操作进行比较。结果表明，深度优化的 Permute 操作比 PyTorch 更快、带宽效率更高，并且带宽利用率接近原生 Copy 操作。

Permute排列实现

Permute的作用是变换张量数据维度的顺序，例如：

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/盐析白兔/article/detail/662041