如何为NumPy数组（含3D）实现Softmax？解决自定义函数报错问题

问题背景 #

自定义了如下Softmax函数：

def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x) 
    return e_x / e_x.sum(axis = 1)

处理形状为50*3的logits数组时，调用该函数出现维度不匹配报错；使用scipy.special.logit能得到概率值，但需要自定义函数同时适配2D和3D的logits数组。

报错原因 #

原函数中e_x.sum(axis=1)对50*3的数组求和后，得到的是形状为(50,)的一维数组，而e_x是(50,3)的二维数组。numpy广播机制要求两个数组的维度能兼容，这里一维数组和二维数组直接做除法会触发维度不匹配。

正确解决方案 #

通用Softmax函数（支持任意维度） #

Softmax的核心是对每个样本的类别维度单独归一化，因此需要保证求和操作后维度与原数组匹配，同时添加数值溢出防护：

import numpy as np

def softmax(x, axis=-1):
    """对输入数组在指定轴上计算Softmax，默认最后一维为类别轴
    参数：
        x: 输入logits数组，支持2D/3D等任意维度
        axis: 进行Softmax计算的轴，默认最后一维（类别维度）
    返回：
        归一化后的概率数组，形状与输入一致
    """
    # 减去每个样本的最大得分，避免指数运算数值溢出
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    # 在指定轴求和并保持维度，确保除法时广播正常
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

适配不同维度的调用方式 #

• 处理50*3的2D logits：
  类别维度是第1轴（索引从0开始），直接调用：

  y_pred_logits_activated = softmax(y_pred_logit, axis=1)
  

• 处理3D logits数组：
  若3D数组的最后一维是类别维度（比如形状为(样本数, 序列长度, 类别数)），直接使用默认参数即可：

  y_pred_logits_activated = softmax(3d_logits_array)
  

注意：错误修改的问题 #

你之前修改的函数：

def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x)
    return np.exp(x)/sum(np.exp(x))

虽然能运行，但本质是对整个数组的所有元素求和归一化，不符合Softmax“每个样本单独归一化类别概率”的定义，会导致结果完全错误，请勿使用。

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