自定义卷积注意力算子的CUDA实现

背景

注意力机制的一种卷积替代方式

我们在上面这篇文章讲到了设计适当的卷积来替代注意力机制

最终设计的卷积如下

Out_{i,c} = \sum^k_{j=1} X_{ i+j-\left \lceil \frac{k+1}{2} \right \rceil, c} \cdot W_{\left \lceil \frac{cH}{d} \right \rceil, j} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (1)

其中X \in R^{B,n,d}, W \in R^{H, k}, 且有W_{h,:} = Softmax(W_{h, :});c代表d这个维度;

要实现这样的卷积,有一种折中的方式是,令

\\X => X' \in R^{BH,n,\frac{d}{H}} \\W => W' \in R^{BH, n, n} \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ \space\ (2)

X'/W'的具体设计请参考上述文章,此处不再赘述;

为了更高效的实现这样的卷积,考虑对该自定义卷积算子进行CUDA实现;


方案

CUDA 实现自定义卷积一般需要几个固定的步骤,在 CUDA实现focal_loss 这篇文章中详细说明了;这里引用一张图,如下

图1:CUDA实现算子的一般步骤介绍

这里CUDA实现自定义卷积算子的步骤和上面差不多,最大的区别是 step4/5,尤其是step5的kernel设计;

以下按照上述流程讲述实现方法;

---------

待续。。。

你可能感兴趣的:(技术问题,#,CUDA,#,深度学习,深度学习,CUDA,神经网络)