Classification losses error about theseus

[kaylode]

ok, I've analyzed this error. It is because MixupCutmix occasionally transforms the targets tensor into soft one-hot encoding , whereas the current FocalLoss and SmoothCE requires argmax labels.

For FocalLoss it can be easily fixed by this:

class FocalLoss(nn.Module):
    def forward(self, outputs: Dict[str, Any], batch: Dict[str, Any], device: torch.device):
        outputs = outputs['outputs']
        targets = move_to(batch['targets'], device)
        num_classes = outputs.shape[-1]

        # Need to be one hot encoding
        if outputs.shape != targets.shape:
            targets = nn.functional.one_hot(targets, num_classes=num_classes)
            targets = targets.float().squeeze()

        loss = sigmoid_focal_loss(outputs, targets, self.alpha, self.gamma, self.reduction)
        loss_dict = {"L": loss.item()}
        return loss, loss_dict

For SmoothCE, timm has SoftTargetCrossEntropy class which is suitable for soft one-hot encoding:

from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy

class SmoothCELoss(nn.Module):
    def __init__(self, smoothing: float=0.1, **kwargs):
        super(SmoothCELoss, self).__init__(**kwargs)
        self.smooth_criterion = LabelSmoothingCrossEntropy()
        self.soft_criterion = SoftTargetCrossEntropy()

    def forward(self, outputs: Dict[str, Any], batch: Dict[str, Any], device: torch.device):
        pred = outputs['outputs']
        target = move_to(batch["targets"], device)

        if pred.shape == target.shape:
            loss = self.soft_criterion(pred, target)
        else:
            loss = self.smooth_criterion(pred, target.view(-1).contiguous())
        loss_dict = {"CE": loss.item()}
        return loss, loss_dict

I will make a PR and you can test it out