label_smoother_sum.py (1748B)
1 #!/usr/bin/env python 2 # -*- coding:utf-8 -*- 3 import torch 4 from dataclasses import dataclass 5 6 7 @dataclass 8 class SumLabelSmoother: 9 """ 10 Adds label-smoothing on a pre-computed output from a Transformers model. 11 12 Args: 13 epsilon (:obj:`float`, `optional`, defaults to 0.1): 14 The label smoothing factor. 15 ignore_index (:obj:`int`, `optional`, defaults to -100): 16 The index in the labels to ignore when computing the loss. 17 """ 18 19 epsilon: float = 0.1 20 ignore_index: int = -100 21 22 def __call__(self, model_output, labels): 23 logits = model_output["logits"] if isinstance(model_output, dict) else model_output[0] 24 log_probs = -torch.nn.functional.log_softmax(logits, dim=-1) 25 if labels.dim() == log_probs.dim() - 1: 26 labels = labels.unsqueeze(-1) 27 28 padding_mask = labels.eq(self.ignore_index) 29 # In case the ignore_index is -100, the gather will fail, so we replace labels by 0. The padding_mask 30 # will ignore them in any case. 31 labels.clamp_min_(0) 32 nll_loss = log_probs.gather(dim=-1, index=labels) 33 smoothed_loss = log_probs.sum(dim=-1, keepdim=True) 34 35 nll_loss.masked_fill_(padding_mask, 0.0) 36 smoothed_loss.masked_fill_(padding_mask, 0.0) 37 38 # Take the mean over the label dimensions, then divide by the number of active elements (i.e. not-padded): 39 # num_active_elements = padding_mask.numel() - padding_mask.long().sum() 40 nll_loss = nll_loss.sum() # / num_active_elements 41 smoothed_loss = smoothed_loss.sum() # / (num_active_elements * log_probs.shape[-1]) 42 eps_i = self.epsilon / log_probs.size(-1) 43 return (1 - self.epsilon) * nll_loss + eps_i * smoothed_loss