Spaces:
Build error
Build error
File size: 1,756 Bytes
e487255 |
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 |
import torch
from torch import nn
class CrossEntropyForMulticlassLoss(torch.nn.CrossEntropyLoss):
# This loss applies cross entropy after reducing the number of prediction
# dimensions to the number of classes in the target
# TODO: loss.item() doesn't work so the displayed losses are Nans
def __init__(self, num_classes, weight=None, size_average=None, ignore_index: int = -100,
reduce=None, reduction: str = 'mean', label_smoothing: float = 0.0) -> None:
super().__init__(size_average=size_average, reduce=reduce, reduction=reduction, ignore_index=ignore_index)
self.num_classes = num_classes
def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
loss = torch.zeros_like(input[:, :, 0])
for b in range(target.shape[1]):
l = super().forward(input[:, b, 0:len(torch.unique(target[:, b]))], target[:, b])
loss[:, b] += l
return loss.flatten()
def JointBCELossWithLogits(output, target):
# output shape: (S, B, NS) with NS = Number of sequences
# target shape: (S, B, SL)
# Loss = -log(mean_NS(prod_SL(p(target_SL, output_NS))))
# Here at the moment NS = SL
output = output.unsqueeze(-1).repeat(1, 1, 1, target.shape[-1]) # (S, B, NS, SL)
output = output.permute(2, 0, 1, 3) # (NS, S, B, SL)
print(target.shape, output.shape)
loss = (target * torch.sigmoid(output)) + ((1-target) * (1-torch.sigmoid(output)))
loss = loss.prod(-1)
loss = loss.mean(0)
loss = -torch.log(loss)
loss = loss.mean()
return loss
class ScaledSoftmaxCE(nn.Module):
def forward(self, x, label):
logits = x[..., :-10]
temp_scales = x[..., -10:]
logprobs = logits.softmax(-1)
|