Implement SGD Optimizer with Warm-up in PyTorch – PyTorch Tutorial

By | October 25, 2022
0 Comment

When we are using SGD optimizer to train a pytorch model, we may use warm-up strategy to improve the training efficiency. In this tutorial, we will introduce you how to implement this strategy in pytorch.

Preliminary

We can use pytorch-gradual-warmup-lr, you can download it here.

https://github.com/LvJC/pytorch-gradual-warmup-lr

Implement SGD optimizer with warm-up

Here is an example code:

  1. import torch
  2. from warmup_scheduler import GradualWarmupScheduler
  4. from matplotlib import pyplot as plt
  5. lr_list = []
  6. model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
  7. LR = 0.001
  9. optimizer = torch.optim.SGD(model, lr=LR, momentum=0.9, weight_decay=1e-3)
  10. scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size= 1, gamma=0.97)
  11. scheduler_warmup = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=6, after_scheduler=scheduler)
  14. for epoch in range(1,200):
  15. data_size = 400
  16. for i in range(data_size):
  17. optimizer.zero_grad()
  18. #loss
  19. optimizer.step()
  20. scheduler_warmup.step(epoch)
  21. #
  22. lr_list.append(scheduler_warmup.get_lr())
  24. print(lr_list)
  25. plt.plot(range(len(lr_list)),lr_list,color = 'r')
  26. plt.show()
import torch
from warmup_scheduler import GradualWarmupScheduler

from matplotlib import pyplot as plt
lr_list = []
model = [torch.nn.Parameter(torch.randn(2, 2, requires_grad=True))]
LR = 0.001

optimizer = torch.optim.SGD(model, lr=LR, momentum=0.9, weight_decay=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size= 1, gamma=0.97)
scheduler_warmup = GradualWarmupScheduler(optimizer, multiplier=1, total_epoch=6, after_scheduler=scheduler)


for epoch in range(1,200):
    data_size = 400
    for i in range(data_size):
        optimizer.zero_grad()
        #loss
        optimizer.step()
    scheduler_warmup.step(epoch)
    #
    lr_list.append(scheduler_warmup.get_lr())

print(lr_list)
plt.plot(range(len(lr_list)),lr_list,color = 'r')
plt.show()

Run this code, we will see the learning rate as follows:

Implement SGD Optimizer wih Warm-up in PyTorch - PyTorch Tutorial

However, we start to train with epoch = 1, the first learning rate is 0.00016666666666666666

If epoch = 0, the first learning rate will be 0

For example:

  1. for epoch in range(0,200):
  2. data_size = 400
  3. for i in range(data_size):
  4. optimizer.zero_grad()
  5. #loss
  6. optimizer.step()
  7. scheduler_warmup.step(epoch)
  8. #
  9. lr_list.append(scheduler_warmup.get_lr())
  11. print(lr_list)
  12. plt.plot(range(len(lr_list)),lr_list,color = 'r')
  13. plt.show()
for epoch in range(0,200):
    data_size = 400
    for i in range(data_size):
        optimizer.zero_grad()
        #loss
        optimizer.step()
    scheduler_warmup.step(epoch)
    #
    lr_list.append(scheduler_warmup.get_lr())

print(lr_list)
plt.plot(range(len(lr_list)),lr_list,color = 'r')
plt.show()

We will see:

Implement SGD Optimizer wih Warm-up epoch = 0 in PyTorch - PyTorch Tutorial

In order to get a correct warm-up learning rate, we shoud start to train with epoch = 1.

Leave a Reply