loss.py (1439B)
1 from typing import Optional 2 3 import torch 4 5 6 class Loss(torch.nn.Module): 7 def __init__(self, config, label_weights): 8 super().__init__() 9 self.config = config 10 self.label_weights = label_weights 11 # reduction none because we weight it with class weights and then we do the average ourself. 12 if self.config.multilabel: 13 self.loss = torch.nn.BCEWithLogitsLoss(reduction="none") 14 else: 15 self.loss = torch.nn.CrossEntropyLoss(reduction="none", label_smoothing=self.config.get("label_smoothing", 0)) 16 17 def forward(self, logits: torch.Tensor, labels: Optional[torch.Tensor] = None): 18 if labels is None: 19 loss = None 20 elif self.config.multilabel: 21 loss = self.loss(logits, labels) 22 loss = loss*self.label_weights[labels%4] # %4: artisanal way to avoid the -100 when indexing 23 loss = loss.mean() 24 else: 25 loss = self.loss(logits.view(labels.numel(), 4), labels.view(-1)) # numel = nb of elements in labels in total (product of dimensions) 26 loss = loss*self.label_weights[labels.view(-1)%4] # %4: artisanal way to avoid the -100 when indexing 27 loss = loss.mean() 28 29 if self.config.multilabel: 30 prediction = torch.nn.functional.sigmoid(logits, dim=2) 31 else: 32 prediction = torch.nn.functional.softmax(logits, dim=2) 33 34 return loss, prediction