Softmax 的 CUDA 实现答案

【问题标题】：CUDA implementation of SoftmaxSoftmax 的 CUDA 实现
【发布时间】：2019-12-04 04:53:41
【问题描述】：

我希望使用 CUDA 提高我的 softmax 层的速度。由于缺乏 python 和 CUDA 的示例，我希望在这里得到一些建议。我已经建立了一个幼稚的实现并寻求相同的建议。

@cuda.jit
def softmax(X, w, b):

    m = X.shape[0]

    probs = np.zeros((m, 120))
    startX=cuda.grid(2)
    gridX=cuda.gridDim.x * cuda.blockDim.x;
    for i in range(startX, m):
        X_slice = X[i,:,:,:]
        z = np.dot(X_slice,w).reshape(1, w.shape[-1])
        z_exp = np.exp(z) 
        z_probs = z_exp/np.sum(z_exp) 
        probs[i,:] =z_probs

    A_prev = (X, w, b)
    return probs, A_prev

【问题讨论】：

我的建议是尝试运行您的代码，当它不可避免地不起作用时，返回一个您需要回答的具体问题。在此之前，我会引导您访问 Numba 文档的 this part

标签： python cuda numba

【解决方案1】：

我已经建立了一个幼稚的实现并正在寻找相同的建议

好的，让我们根据文档中的this page 来看看你的代码：

@cuda.jit
def softmax(X, w, b):

    m = X.shape[0]                                     # OK

    probs = np.zeros((m, 120))                         # Illegal, array creation
    startX=cuda.grid(2)                                # OK
    gridX=cuda.gridDim.x * cuda.blockDim.x;            # OK but unused
    for i in range(startX, m):                         # Illegal, startX is not a scalar
        X_slice = X[i,:,:,:]                           # Illegal, array creation
        z = np.dot(X_slice,w).reshape(1, w.shape[-1])  # Illegal, array method
        z_exp = np.exp(z)                              # Illegal, array method
        z_probs = z_exp/np.sum(z_exp)                  # Illegal, array method
        probs[i,:] =z_probs                            # OK (except probs can't be created in kernel)

    A_prev = (X, w, b)                                 # OK tuples are supported
    return probs, A_prev                               # illegal, kernels can't return anything

所以基本上，您的整个代码都是非法的，并且永远无法工作。如果您需要一些建议，请参阅我链接到的文档：

为了获得最佳性能，用户应该编写代码，使得每个线程一次处理一个元素

强调我的。 CUDA Python 编程不仅仅是在您已经拥有的一些代码前面添加@cuda.jit 并期望它能够快速运行。您需要为完整算法想象元素明智的代码，然后编写代码来执行它。对于上面的代码，这可能需要多个内核。

【讨论】：