我在尝试实现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)>>> ttensor([[-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]])为什么这些不同的实现会产生不同的结果？
回答：
这个方法有效
import torchdef 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_sumt = torch.randn(3, 2) s1 = my_softmax(t)s2 = t.softmax(1)print(torch.allclose(s1, s2))True