如何编写自定义 CrossEntropyLoss

Question

我正在 Pytorch 中学习逻辑回归，为了更好地理解我正在定义一个自定义 CrossEntropyLoss，如下所示：

def softmax(x):
    exp_x = torch.exp(x)
    sum_x = torch.sum(exp_x, dim=1, keepdim=True)

    return exp_x/sum_x

def log_softmax(x):
    return torch.exp(x) - torch.sum(torch.exp(x), dim=1, keepdim=True)

def CrossEntropyLoss(outputs, targets):
    num_examples = targets.shape[0]
    batch_size = outputs.shape[0]
    outputs = log_softmax(outputs)
    outputs = outputs[range(batch_size), targets]

    return - torch.sum(outputs)/num_examples

我还进行了自己的逻辑回归（用于预测 FashionMNIST），如下所示：

input_dim = 784 # 28x28 FashionMNIST data
output_dim = 10

w_init = np.random.normal(scale=0.05, size=(input_dim,output_dim))
w_init = torch.tensor(w_init, requires_grad=True).float()
b = torch.zeros(output_dim)

def my_model(x):
    bs = x.shape[0]
    return x.reshape(bs, input_dim) @ w_init + b

为了验证我的自定义交叉熵损失，我将其与 Pytorch 的 nn.CrossEntropyLoss 进行了比较，将其应用于 FashionMNIST 数据，如下所示：

criterion = nn.CrossEntropyLoss()

for X, y in trn_fashion_dl:
    outputs = my_model(X)
    my_outputs = softmax(outputs)

    my_ce = CrossEntropyLoss(my_outputs, y)
    pytorch_ce = criterion(outputs, y)

    print (f'my custom cross entropy: {my_ce.item()}\npytorch cross entroopy: {pytorch_ce.item()}')
    break 

我的问题是关于结果 my_ce（我的交叉熵）与 pytorch_ce（pytorch 交叉熵）的不同之处：

my custom cross entropy: 9.956839561462402
pytorch cross entroopy: 2.378990888595581

提前感谢您的帮助！

Answer 1

您的代码中有两个错误。

1. log_softmax(x) 应该是这样的，

def log_softmax(x):
    return torch.log(softmax(x))

1. 计算自己的CE损失时，应输入outputs而不是my_outputs。因为您将在自己的 CE 损失函数中计算 softmax。应该是这样的，

outputs = my_model(X)
my_ce = CrossEntropyLoss(outputs, y)
pytorch_ce = criterion(outputs, y)

然后你会得到相同的结果。

my custom cross entropy: 3.584486961364746
pytorch cross entroopy: 3.584486961364746

Answer 2

您的log_softmax fn 似乎是错误的。它应该是：

def log_softmax(x):
    return torch.log(softmax(x))

但由于您的 softmax 在数值上不稳定，它可能会有些不稳定。您可以通过以下方式对其进行改进：

def log_softmax(x):
    return x - torch.logsumexp(x,dim=1)

请注意，我使用了身份 log (exp{x}/sum exp(x)) = x - log (sum exp(x))

另见https://pytorch.org/docs/stable/torch.html?highlight=logsumexp#torch.logsumexp