Realize cosine learning rate based on PyTorch

1. Need to use the library

Set the learning rate and model

import math
import matplotlib.pyplot as plt
import torch.optim as optim
from torchvision.models import resnet18

lr_rate = 0.1
model = resnet18(num_classes=10)

2.LambdaLR实现cosine learning rate

Set up lambda, optimizer and scheduler 

lambda1 = lambda epoch: (epoch / 4000) if epoch < 4000 else 0.5 * (math.cos((epoch - 4000)/(100 * 1000 - 4000) * math.pi) + 1)
optimizer = optim.SGD(model.parameters(), lr=lr_rate, momentum=0.9, nesterov=True)
scheduler = optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)

3.learning rate设置

index = 0
x = []
y = []
for epoch in range(100):
    for batch in range(1000):
        x.append(index)
        y.append(optimizer.param_groups[0]['lr'])
        index += 1
        scheduler.step()

4. Visualize the learning rate

plt.figure(figsize=(10, 8), dpi=200)
plt.xlabel('batch stop')
plt.ylabel('learning rate')
plt.plot(x, y, color='r', linewidth=2.0, label='modify data')
plt.legend(loc='upper right')
plt.savefig('result.png')
plt.show()

5. Result of learning rate change

TensorBoard learning rate可视化

 

Guess you like

Origin blog.csdn.net/linghu8812/article/details/105122161