复制pytorch的softmax

Question

我试图在pytorch中实现softmax函数，但我的实现的输出无法与pytorch的实现的输出相匹配。

我之所以尝试这样做，是因为我想继续实现一个屏蔽的softmax，它不会在分母的和中包括某些索引，并为这些屏蔽的索引设置输出。

我想计算一个矩阵，其中输出中的每一行和为1。我目前的实现是：

def my_softmax(x):
    exp = x.exp()
    return exp / exp.sum(1, keepdim=True)

然而，我的实现和pytorch的输出是不同的：

>>> t = torch.randn(3, 2)
>>> t
tensor([[-1.1881, -0.1085],
        [ 0.5825,  1.0719],
        [-0.5309, -1.3774]])
>>> my_softmax(t)
tensor([[0.2536, 0.7464],
        [0.3800, 0.6200],
        [0.6998, 0.3002]])
>>> t.softmax(1)
tensor([[0.2536, 0.7464],
        [0.3800, 0.6200],
        [0.6998, 0.3002]])
>>> my_softmax(t) == t.softmax(1)
tensor([[False,  True],
        [False, False],
        [ True,  True]])

为什么这些不同的实现会产生不同的结果？

Answer 1

这是可行的

import torch

def my_softmax(x):

    means = torch.mean(x, 1, keepdim=True)[0]
    x_exp = torch.exp(x-means)
    x_exp_sum = torch.sum(x_exp, 1, keepdim=True)

    return x_exp/x_exp_sum

t = torch.randn(3, 2) 

s1 = my_softmax(t)
s2 = t.softmax(1)

print(torch.allclose(s1, s2))
True

附注:摘自讨论https://discuss.pytorch.org/t/how-to-implement-the-exactly-same-softmax-as-f-softmax-by-pytorch/44263