多维下numpy实现torch下的log_softmax

一维情况下，来源：link

import numpy as np
import torch
import torch.nn.functional as F

def my_softmax(x):
    exp_x = np.exp(x)
    return exp_x/np.sum(exp_x)

def my_log_softmax1(x):
    return np.log(my_softmax(x))

def my_log_softmax2(x):
    return x-np.log(np.sum(np.exp(x)))

x = np.array([1.5, 2.2, 3.1, 0.9, 1.2, 1.7])
xt = torch.from_numpy(x)

print('F.softmax:', F.log_softmax(xt))
print('my_softmax', my_log_softmax1(x))
print('my_softmax', my_log_softmax2(x))

二维情况下

import numpy as np
import torch
import torch.nn.functional as F


def my_log_softmax2(x, axis):
    return x - np.log(np.sum(np.exp(x), axis=axis))[..., np.newaxis]


x = np.array([[1.5, 2.2, 3.1, 0.9, 1.2, 1.7], [1, 2, 3, 4, 5, 6]])
xt = torch.from_numpy(x)

torch_res = F.log_softmax(xt)
np_res = my_log_softmax2(x)
np.testing.assert_equal(torch_res.numpy(), np_res)