mtl-uncertainty-loss 0.1.0

Uncertainty-weighted multi-task loss for PyTorch (Kendall et al. 2018)

mtl_uncertainty_loss

Uncertainty-weighted multi-task loss for PyTorch, based on Kendall et al. (2018): "Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics" (arXiv:1705.07115)

Features

• Arbitrary number of tasks
• Learnable uncertainty parameters (log σ per task)
• PyTorch nn.Module interface — log_sigmas are registered parameters and updated during training
• Easy to monitor uncertainties via get_sigmas()

Installation

pip install mtl_uncertainty_loss

Usage

from mtl_uncertainty_loss import UncertaintyWeightedMultiTaskLoss

loss_fn = UncertaintyWeightedMultiTaskLoss(num_tasks=3)

# Pass already-reduced (scalar) per-task losses
total_loss = loss_fn([loss1, loss2, loss3])
total_loss.backward()

print(loss_fn.get_sigmas())  # tensor of current σ values

Testing

pip install -e ".[dev]"
pytest tests

License

MIT