我已经实现了一个自定义的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进行训练。

回答：

如果你只需要交叉熵，你可以利用PyTorch已经定义好的函数。

import torch.nn.functional as Floss_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)

变成你需要的形式，这样你就有了。