Spaces:
Build error
Build error
File size: 1,696 Bytes
2890e34 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
import torch
from torch import nn
from torch.nn import KLDivLoss
from torch.nn import LogSoftmax
class LabelSmoothingLoss(nn.Module):
def __init__(self, label_smoothing=0.0, unreliable_label=None, ignore_index=-100):
"""
If label_smoothing == 0.0, it is equivalent to xentropy
"""
assert 0.0 <= label_smoothing <= 1.0
super(LabelSmoothingLoss, self).__init__()
self.ignore_index = ignore_index
self.label_smoothing = label_smoothing
self.loss_fn = KLDivLoss(reduction='batchmean')
self.unreliable_label = unreliable_label
self.max_gap = 100.
self.log_softmax = LogSoftmax(1)
def forward(self, output, target):
"""
output: logits
target: labels
"""
vocab_size = output.shape[1]
mask = (target != self.ignore_index)
output, target = output[mask], target[mask]
output = self.log_softmax(output)
def get_smooth_prob(ls):
smoothing_value = ls / (vocab_size - 1)
prob = output.new_full((target.size(0), vocab_size), smoothing_value)
prob.scatter_(1, target.unsqueeze(1), 1 - ls)
return prob
if self.unreliable_label is not None:
smoothed_prob = get_smooth_prob(self.label_smoothing)
hard_prob = get_smooth_prob(0.0)
unreliable_mask = (target == self.unreliable_label).to(torch.float)
model_prob = ((smoothed_prob.T * unreliable_mask) + (hard_prob.T * (1 - unreliable_mask))).T
else:
model_prob = get_smooth_prob(self.label_smoothing)
loss = self.loss_fn(output, model_prob)
return loss
|