在PyTorch中自定义交叉熵损失函数

Question

我已经对PyTorch交叉熵损失函数进行了自定义实现（因为我需要在以后引入更多的灵活性）。我打算用这个模型进行训练，但它需要相当长的时间来训练，而且可用资源不能仅用于测试函数是否正确实现。我已经实现了向量化实现，因为它运行起来更快。

以下是我编写的代码:

def custom_cross(my_pred,true,batch_size=BATCH_SIZE):
    loss= -torch.mean(torch.sum(true.view(batch_size, -1) * torch.log(my_pred.view(batch_size, -1)), dim=1))
    return loss

我希望你能提出一种更优化的实现方案，或者指出我当前的实现方案存在的错误。该模型将使用Nvidia Tesla K-80进行训练。

- Inder

---

由于log(0)未定义，我认为您需要在预测值和目标值中添加一个epsilon(1e-10)。 - Craig.Li

如果my_pred是softmax的结果，您可以使用LogSoftmax来避免一些指数。 - Jindřich

@Craig.Li，这是非常有价值的建议，谢谢你。 - Inder

@Jindřich 是的，我的 my_pred 确实是 softmax 的结果，但很抱歉我不确定这将如何影响或可能的实现方式，请问可否再多解释一些，谢谢。 - Inder

如果你使用softmax softmax = e^{X} / [sum e^{X}]，你总是会从除法中得到舍入误差，并且你会对X进行指数运算，这不是最快的操作。如果你使用logsoftmax = X - log [sum e^{X}]，你可以避免舍入误差，并且摆脱对X进行指数和对数运算。由于聪明的实现，sum e^{X}的对数与仅求和一样快，而且没有舍入误差。 - Jindřich

@Jindřich，非常感谢您的精彩分享，也感谢您详细的解释。 - Inder

Answer 1

如果您只需要交叉熵，您可以利用PyTorch定义的优势。

import torch.nn.functional as F
loss_func = F.cross_entropy

建议提供更优化的实现

PyTorch有很多F.损失函数，但您可以使用纯Python轻松编写自己的损失函数。PyTorch会自动为您的函数创建快速的GPU或矢量化CPU代码。

因此，您可以查看PyTorch的原始实现，但我认为是这个：

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

这里是交叉熵损失的原始实现，现在您可以进行修改：

此处。

nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)

对于你需要的东西，只要你拥有它。

- prosti