多重损失的梯度裁剪

Question

我需要使用多重损失来训练神经网络。最基本的方法是对损失求和，然后做一个梯度步骤

optimizer.zero_grad()
total_loss = loss_1 + loss_2
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()

然而，有时一个损失可能会接管，我希望双方都做出同样的贡献。我考虑过像下面这样的单次传球后的剪裁损失

optimizer.zero_grad()
loss_1.backward(retain_graph=True)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
loss_2.backward(retain_graph=True)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()

但这仍然是一个问题，因为如果第一个损失很大，梯度将已经接近 max_norm。

我想过这样的事情

optimizer.zero_grad()
loss_1.backward(retain_graph=True)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
grad_1 = ... # save gradient

optimizer.zero_grad()
loss_2.backward(retain_graph=True)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
grad_2 = ... # save gradient

optimizer.zero_grad()
grad = grad_1 + grad_2
# manually apply grad (how?)

optimizer.step()

有更好的方法吗？我如何以一种很好的方式手动将渐变应用于模型？