softmax python实现

softmax将k维度的向量压缩至另一个k维度向量，新向量每个元素都在0,1之间，加起来等于1。
$\sigma {(\mathbf{z})}_j=\frac{{\mathbf{e}}^{z_j}}{{\sum }_{k=1}^{\mathbf{K}}{\mathbf{e}}^{z_k}},j\in 1,...,K$
数值稳定问题：
${\mathbf{e}}^c$当c很大的时候，${\mathbf{e}}^c$会overflow。
${\mathbf{e}}^c$当c很小的时候，${\mathbf{e}}^c$会underflow。
已知softmax(x)=softmax(x+c)。所以可以这样，x是一个向量，令x_max为x元素的最大值，x_i为x的任意元素，把所有x_i都减去x_max，得到一个新的向量z。由softmax(x)=softmax(x+c)知softmax(z)=softmax(x)。算softmax(z)可以防止overflow和underflow。
用pytorch的softmax时不支持整型，也不支持float16。
自己写的softmax不支持整型，支持所有浮点型。
应养成习惯，在初始化ndarray时指定数据类型，第一是有些功能只支持某些数据类型，第二是避免内存浪费，比如如果float32就可以满足要求没必要float64，因为如果不加以限制，在反向传播时可能会占用过多内存，在实际模型部署时也不利于模型加速。

'''

1. 使用numpy实现Softmax（假设2个样本，给定[[1, 2, 3], [2，1，3]])；
2. 使用torch.nn.functional.softmax() 验证结果是否一致。
'''
import numpy as np
import torch
import torch.nn.functional as F

# 支持1维度或2维度或一批的数据
def stable_softmax(x):
    z = x - np.max(x, axis=-1, keepdims=True)
    numerator = np.exp(z)
    denominator = np.sum(numerator, axis=-1, keepdims=True)
    softmax = numerator / denominator
    return softmax


def softmax_my(predict):
    '''
    使用numpy实现Softmax
    input:
        numpy.ndarray
        [[1, 2, 3],
        [2，1，3]]
    output:
        softmax value: numpy.ndarray
    '''

    row_size, col_size = predict.shape

    # 使用max()，其中参数 axis = 1 表示二维数组中沿着横轴取最大值
    max_i = np.max(predict, axis=-1)


    # 每一行减去本行最大的数字，reshape
    reshaped_max_i = max_i.reshape(-1, 1)
    z = predict - reshaped_max_i
    # max_i.shape:(2,)
    # predict.shape:(2, 3)
    # reshaped_max_i.shape:(2, 1)
    # z.shape:(2, 3)



    # 计算每个元素e的指数次幂
    numerator = np.exp(z)

    # 对每一行进行求和操作
    denominator = np.sum(numerator, axis=-1, keepdims=True)

    # 每一行 predict_exp / predict_exp_row_sum
    return numerator/denominator




if __name__ == '__main__':
    '''
    假设两个样本
    [[1, 2, 3],
    [2，1，3]]
    '''
    # 如果想尝试随机数
    # np.random.seed(0)
    # predict = np.random.randn(2, 3) # 返回标准正态分布的一组数据
    # print(predict)
    predict = np.array([[1, 2, 3], [2, 1, 3]], dtype=np.float16)
    softmax_value = softmax_my(predict)
    print('softmax结果：', softmax_value)

    # 验证softmax是否每行和为1
    print(softmax_value.sum(axis=1))

    # torch.nn.functional.softmax(input, dim)
    # 参数：dim:指明维度，dim=0表示按列计算；dim=1表示按行计算
    predict_tensor = torch.tensor([[1, 2, 3], [2, 1, 3]], dtype=torch.float32) # F.softmax不支持整型数，也不支持float16, 会报错RuntimeError: "softmax_lastdim_kernel_impl" not implemented for 'Half'
    softmax_torch = F.softmax(predict_tensor, dim=1)