我已经完成了 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 进行训练。