学习率调整之stepLR

在PyTorch的torch.optim包提供了非常多的可实现参数自动优化的类（SGD、Adam..）及学习率调整的类lr_scheduler

class torch.optim.lr_scheduler.stepLR(optimizer, step_size,gamma=0.1,last_epoch=-1)

参数：

optimizer —要更新学习率的优化器

step_size — 每经step_size个epoch,更新一次学习率lr

gamma — 更新lr的乘法因子

last_epoch — 默认为-1表示从头开始训练

代码示例：

import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from cnn_finetune import make_model
'''模型定义'''
model = make_model(model_name='vgg16', pretrained=True, num_classes=2,input_size= 224,224))

lr = 0.1
optimizer = optim.Adam(model.parameters(), lr=lr)  # 用Adam优化算法更新权重
criterion = nn.CrossEntropyLoss()  # 损失函数
scheduler = lr_scheduler.StepLR(optimizer, 8, gamma=0.1, last_epoch=-1)
for epoch in range(25):
    optimizer.zero_grad()
    optimizer.step()
    print("第%d轮的学习率：%f"%(epoch, optimizer.param_groups[0]['lr']))
    scheduler.step()

输出结果：

来源：maly_Sunshine