如何对深度学习中softmax前向传播的for循环进行numpy向量化优化

你原有的循环实现是逐样本（输入第0维为样本维度）计算带数值稳定性优化的Softmax，完全向量化的实现代码如下：

def forward(self, input_t):
    # 逐行计算样本最大值，保留维度适配广播运算
    curr = input_t - np.max(input_t, axis=1, keepdims=True)
    exp_curr = np.exp(curr)
    # 逐行计算指数和后做广播除法完成归一化
    output_t = exp_curr / np.sum(exp_curr, axis=1, keepdims=True)
    self.store = output_t
    return output_t

实现说明 #

• 核心逻辑和你原有循环完全等价：保留了减去每行最大值的操作，避免指数运算出现数值溢出
• 所有聚合计算（取最大值、求和）都指定axis=1对应原来循环里的逐行处理，keepdims=True保证运算时的维度匹配，不需要额外做维度reshape即可触发numpy广播机制完成逐元素运算
• 无显式循环的实现处理大批量输入时性能远高于循环版本

内容的提问来源于stack exchange，提问作者user15788706