Adds label-smoothing on a pre-computed output from a Transformers model. Args: epsilon (:obj:`float`, `optional`, defaults to 0.1): The label smoothing factor. ignore_index (:obj:`int`, `optional`, defaults to -100): The index in the labels to ignore
| 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 |