适配pytorch softmax函数

Question

我目前正在研究 softmax 函数，我想将最初实现的函数用于一些小型测试。

我去过文档，但没有太多关于该功能的有用信息。这是pytorch python implementation：

def __init__(self, dim=None):
    super(Softmax, self).__init__()
    self.dim = dim

def __setstate__(self, state):
    self.__dict__.update(state)
    if not hasattr(self, 'dim'):
        self.dim = None

def forward(self, input):
    return F.softmax(input, self.dim, _stacklevel=5)

在哪里可以找到 F.softmax 实现？

我想尝试的一个例子是这里描述的 soft-margin softmax：Soft-Margin Softmax for Deep Classification

从哪里开始最好？ 提前致谢！

Answer 1

PyTorch 和 Numpy 中的 Softmax 实现

一个Softmax函数定义如下：

上述公式的直接实现如下：

def softmax(x):
    return np.exp(x) / np.exp(x).sum(axis=0)

由于np.exp(x)，上述实现可能会遇到算术溢出。

为了避免溢出，我们可以将softmax方程中的分子和分母除以常数C。那么softmax函数就变成了：

上述方法在 PyTorch 中实现，我们将log(C) 视为-max(x)。下面是 PyTorch 的实现：

def softmax_torch(x): # Assuming x has atleast 2 dimensions
    maxes = torch.max(x, 1, keepdim=True)[0]
    x_exp = torch.exp(x-maxes)
    x_exp_sum = torch.sum(x_exp, 1, keepdim=True)
    probs = x_exp/x_exp_sum
    return probs 

对应的 Numpy 等价物如下：

def softmax_np(x):
    maxes = np.max(x, axis=1, keepdims=True)[0]
    x_exp = np.exp(x-maxes)
    x_exp_sum = np.sum(x_exp, 1, keepdims=True)
    probs = x_exp/x_exp_sum
    return probs 

我们可以将结果与 PyTorch 实现进行比较 - torch.nn.functional.softmax 使用下面的 sn-p：

import torch
import numpy as np
if __name__ == "__main__":
    x = torch.randn(1, 3, 5, 10)
    std_pytorch_softmax = torch.nn.functional.softmax(x)
    pytorch_impl = softmax_torch(x)
    numpy_impl = softmax_np(x.detach().cpu().numpy())
    print("Shapes: x --> {}, std --> {}, pytorch impl --> {}, numpy impl --> {}".format(x.shape, std_pytorch_softmax.shape, pytorch_impl.shape, numpy_impl.shape))
    print("Std and torch implementation are same?", torch.allclose(std_pytorch_softmax, pytorch_impl))
    print("Std and numpy implementation are same?", torch.allclose(std_pytorch_softmax, torch.from_numpy(numpy_impl)))

参考：

• Softmax discussion
• Discussion on Softmax implementation at PyTorch forumn
• An SO thread on implementation of Softmax in Python