AutomaticWeightedLoss
A PyTorch implementation of Liebel L, Körner M. Auxiliary tasks in multi-task learning[J]. arXiv preprint arXiv:1805.06334, 2018.
The above paper improves the paper "Multi-task learning using uncertainty to weigh losses for scene geometry and semantics" to avoid the loss of becoming negative during training.
Requirements
- Python
- PyTorch
How to Train with Your Model
- Clone the repository
git clone [email protected]:Mikoto10032/AutomaticWeightedLoss.git- Create an AutomaticWeightedLoss module
from AutomaticWeightedLoss import AutomaticWeightedLoss
awl = AutomaticWeightedLoss(2) # we have 2 losses
loss1 = 1
loss2 = 2
loss_sum = awl(loss1, loss2)- Create an optimizer to learn weight coefficients
from torch import optim
model = Model()
optimizer = optim.Adam([
{'params': model.parameters()},
{'params': awl.parameters(), 'weight_decay': 0}
])- A complete example
from torch import optim
from AutomaticWeightedLoss import AutomaticWeightedLoss
model = Model()
awl = AutomaticWeightedLoss(2) # we have 2 losses
loss_1 = ...
loss_2 = ...
# learnable parameters
optimizer = optim.Adam([
{'params': model.parameters()},
{'params': awl.parameters(), 'weight_decay': 0}
])
for i in range(epoch):
for data, label1, label2 in data_loader:
# forward
pred1, pred2 = Model(data)
# calculate losses
loss1 = loss_1(pred1, label1)
loss2 = loss_2(pred2, label2)
# weigh losses
loss_sum = awl(loss1, loss2)
# backward
optimizer.zero_grad()
loss_sum.backward()
optimizer.step()Something to Say
Actually, it is not always effective, but I hope it can help you.