pytorch 中损失函数的奇怪行为

Weird behaviour of loss function in pytorch

我正在计算一个自定义成本函数,它只是用交叉熵的指数除以参数 \eta。在第一次迭代期间(大约 20 次),训练损失正在减少,但在那之后,我突然得到一个 nan,我不明白为什么会这样。

我使用的代码如下:

e_loss = []
eta = 2 #just an example of value of eta I'm using 
criterion = nn.CrossEntropyLoss()
for e in range(epoch):
    train_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        client_model.train()
        optimizer.zero_grad()
        output = client_model(data)
        loss = torch.exp(criterion(output, target)/eta) # this is the line where I input my custom loss function
        loss.backward()
        optimizer.step()
        train_loss += loss.item()*data.size(0)
    train_loss = train_loss/len(train_loader) # average losses
    e_loss.append(train_loss)

直接使用exp在输入无界时非常不稳定。如果网络非常自信地预测错误 class(b/c -log(x) 变为 inf,因为 x 变为 0),交叉熵损失可以 return 非常大的值。像这样的模型的单个不准确预测会导致数值精度,从而导致梯度变为 nan,这将立即导致模型的权重和输出变为 nan。

例如

>>> import torch
>>> import torch.nn.functional as F
>>> torch.exp(F.cross_entropy(torch.tensor([[-50.0, 50.0]]), torch.tensor([0])))
tensor(inf)