python之如何向量化多维矩阵的 Softmax 概率

我正在尝试通过 assignment 1斯坦福 cs244n 类(class)。问题 1b 强烈建议对 Softmax 函数进行优化。我设法得到了N维向量的Softmax。我还得到了 MxN 维矩阵的 Softmax，但在列中使用了 for 循环。我有以下代码:

def softmax(x): 
    orig_shape = x.shape 
 
    # Matrix 
    if len(x.shape) > 1: 
        softmax = np.zeros(orig_shape) 
        for i,col in enumerate(x): 
            softmax[i] = np.exp(col - np.max(col))/np.sum(np.exp(col - np.max(col))) 
    # Vector 
    else: 
        softmax = np.exp(x - np.max(x))/np.sum(np.exp(x - np.max(x))) 
    return softmax 

我可以实现更优化的 Matrix 实现吗？

请您参考如下方法:

使用 NumPy broadcasting关于相关ufuncs这涵盖了通用维数的 ndarrays -

exp_max = np.exp(x - np.max(x,axis=-1,keepdims=True)) 
out = exp_max/np.sum(exp_max,axis=-1,keepdims=True)