复制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