Uncertainty Loss不确定损失 处理多任务学习，pytorch代码实现

可以使用以下代码实现 Uncertainty Loss：

import torch
import torch.nn as nn

class UncertaintyLoss(nn.Module):
    def __init__(self, num_tasks, reduction='mean'):
        super(UncertaintyLoss, self).__init__()
        self.num_tasks = num_tasks
        self.loss_fn = nn.CrossEntropyLoss(reduction=reduction)

    def forward(self, input, target, lambda_=1.0):
        bs, num_classes = input.shape[-2], input.shape[-1]
        task_loss = torch.stack([self.loss_fn(input[:, i], target[:, i]) for i in range(self.num_tasks)])
        weights = torch.exp(-lambda_ * task_loss)
        total_loss = torch.sum(weights * task_loss) / torch.sum(weights)

        return total_loss

其中，num_tasks 表示任务数量，lambda_ 表示不确定性的权重。在使用时，将输入的 loss 替换为 UncertaintyLoss 即可。