import mathimport matplotlib.pyplot as pltimport torch.optim as optimfrom torchvision.models import resnet18lr_rate = 0.0001model = resnet18(num_classes=10)# T_max = 1000epoch_total = 25epoch_iter = 609warm_up = 800lambda1 = lambda epoch: (epoch / warm_up) if epoch < warm_up else 0.5 * (math.cos((epoch - warm_up)/(epoch_total*epoch_iter - warm_up) * 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)index = 0x = []y = []for epoch in range(epoch_total): for batch in range(609): x.append(index) y.append(optimizer.param_groups[0]['lr']) index += 1 scheduler.step()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()
结果: