Simple, easy-to-integrate PyTorch implementation of GradNorm: https://arxiv.org/abs/1711.02257.
Despite its popularity, GradNorm has rather few mature implementations in open source. There are several other PyTorch implementations of GradNorm:
- https://github.com/lucidrains/gradnorm-pytorch
- https://github.com/LucasBoTang/GradNorm
- https://github.com/brianlan/pytorch-grad-norm
- https://github.com/hosseinshn/GradNorm
However, each of these has one or more shortcomings that can cause issues when using them in larger PyTorch training frameworks:
- The library requires delegating execution of
.backward()
to it, preventing frameworks like PyTorch Lightning or Keras from being able to call the.backward()
, as in 1 and 2. - The library is not standalone, as in 3.
- The library lacks verifiably correct behavior compared to the original paper, as in 1, 3-4.
This repo aims to address this, by being completely separate from the model's training loop and having verified performance.
Usage is easy, just 3 lines:
gradnorm = GradNorm(model.last_shared_layer, alpha=0.12, number_of_tasks=10,
lr=1e-3, lr_warmup=1000, device='cuda:0')
# in training loop
L_grad = gradnorm.gradnorm(task_losses)
gradnorm.apply_grads(L_grad)
Insert it into your training loop like so:
for epoch in range(epochs):
for x, y_true in dataloader:
x, y_true = x.to(device), y_true.to(device)
y_pred = model(x)
# Your training loop's loss calculation for each task
task_losses = F.mse_loss(y_pred, y_true, reduction='none').mean(dim=-1)
L_i = task_losses.mean(dim=0)
L = torch.sum(gradnorm.w_i * L_i) # <--- Use gradnorm weights w_i here
# Your training loop's backward call
optimizer.zero_grad()
L.backward(retain_graph=True) # retain_graph=True is currently required
# Compute the GradNorm loss
L_grad = gradnorm.gradnorm(L_i) # <--- Line 1
gradnorm.apply_grads(L_grad) # <--- Line 2
optimizer.step()
We re-create the toy model experiment from the paper with the below results.
![image](https://private-user-images.githubusercontent.com/17986781/294347259-44d25e11-84a5-423d-8e49-946988e0fd36.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjU3NDczNTIsIm5iZiI6MTcyNTc0NzA1MiwicGF0aCI6Ii8xNzk4Njc4MS8yOTQzNDcyNTktNDRkMjVlMTEtODRhNS00MjNkLThlNDktOTQ2OTg4ZTBmZDM2LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA5MDclMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwOTA3VDIyMTA1MlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTFmODI5ZWVkOGZkNGNjODYwNDBhODA0NjA4N2FjNzQ0NDRlYWI3NjkyZmUzN2E3NDcyODMwZjI0MzM3ZjEzMGUmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.HAZ710qP1EPyE0xqdLY-1wve2JIhUC2uKMhPoaGYx5E)
Where the paper results are:
![image](https://private-user-images.githubusercontent.com/17986781/294348116-a84225f4-3d9e-4fb4-a6b9-cb639fcb8383.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjU3NDczNTIsIm5iZiI6MTcyNTc0NzA1MiwicGF0aCI6Ii8xNzk4Njc4MS8yOTQzNDgxMTYtYTg0MjI1ZjQtM2Q5ZS00ZmI0LWE2YjktY2I2MzlmY2I4MzgzLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA5MDclMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwOTA3VDIyMTA1MlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTg4YjZjZDZiYTQ2ZGYyODM2MzkyZmI2NzFiZDJiNDBiMDY2NGJkZjEwMmFiOTdkNTM2ODVjN2MyNzM4MzQ3ZWUmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.JWrO6DRPHYHSh-J5biLY5kVjKDjGRTQ2DiGR2MqN5JU)
Experimenting with different learning rates yields slightly different curves. Adding a learning rate warmup to the toy model training yields even smoother weight adjustments.
![image](https://private-user-images.githubusercontent.com/17986781/294349155-f0be12e4-0f44-4f34-b531-85ec48380236.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjU3NDczNTIsIm5iZiI6MTcyNTc0NzA1MiwicGF0aCI6Ii8xNzk4Njc4MS8yOTQzNDkxNTUtZjBiZTEyZTQtMGY0NC00ZjM0LWI1MzEtODVlYzQ4MzgwMjM2LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA5MDclMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwOTA3VDIyMTA1MlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWE2OTEwMWY2YzA1ZTBmNDAwNTQxNWEyY2VjMTBhM2E2OTE5ZmQxNTMyZmI1OWUxNDc2YzU2ZmFmMGFlMWQwMzMmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.0Np5-IUMS12g9E4LKpb_lETIKKL04VKgMr6EgWpnTxg)