如何在pytorch中自己实现softmax运算

本文记录如何实现softmax运算。

线性回归模型适用于输出为连续值的情景

softmax运算使输出更适合离散值的预测和训练

softmax用于分类问题：

预备知识：首先我们看看这个关于Tensor的运算：sum()函数

>>X = torch.tensor([1, 2, 3])

tensor([1, 2, 3])

>>X.sum(dim=1, keepdim=True)

tensor([ 6])

Sum函数将每一行的所有数相加，keepdim=True表示将总数保留在相应的dim轴，关于dim轴的问题看我这个博客。点这里

以及tensor.exp()函数：

>>torch.exp(torch.tensor([0, 2.0]))
tensor([1.0000, 7.3891])

它返回的是tensor([e^0, e^2]), 注意这个的参数tensor必须是浮点数。

 

这样我们就可以使用 总数/每个数=占总数的百分比

1+2+3=6 

如 X就变成了tensor([1/6, 2/6, 3/6])

这样一行的数相加为1，至此实现softmax运算。

我们可以将这个运算封装成一个方法

def softmax(X):

    X_exp = X.exp()

    partition = X_exp.sum(dim=1, keepdim=True)

    return X_exp / partition  # 这里使用了numpy的广播机制

现在来使用一下，创建一个tensor

>>X = torch.rand((1, 5))

tensor([[0.0160, 0.7090, 0.3592, 0.2295, 0.9136]])

>>X_prob = softmax(X)

>>X_prob

tensor([[0.1235, 0.2468, 0.1740, 0.1528, 0.3029]])

     不用验算，相加肯定为1

 

参考：https://tangshusen.me/Dive-into-DL-PyTorch/#/chapter03_DL-basics/3.4_softmax-regression