torch.optim是一个用于实现优化算法的包,支持常见的优化算法。本文介绍一下这个包下面的一些常见内容,如optimizer和lr_scheduler的构造和使用。
optimizer用于根据已得到的梯度来更新网络参数。
optimizer(优化器)的构造:
构造优化器需要传入一个包含要学习参数的迭代器,此外还有一些学习策略参数如学习率等。
以常用的两个优化器SGD和Adam为例:
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)
optimizer(优化器)的属性或方法:
optimizer也有state_dict(),其中包含两部分:
#model define...
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])
输出:
Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]
也可以这样来查看
optimizer=optim.Adam(model.parameters(),0.01)
print("optimizer.param_groups的长度:{}".format(len(optimizer.param_groups)))
for param_group in optimizer.param_groups:
print(param_group.keys())
print([type(value) for value in param_group.values()])
print('查看学习率: ',param_group['lr'])
optimizer.param_groups的长度:1
dict_keys(['params', 'lr', 'betas', 'eps', 'weight_decay', 'amsgrad'])
[<class 'list'>, <class 'float'>, <class 'tuple'>, <class 'float'>, <class 'int'>, <class 'bool'>]
查看学习率: 0.01
优化器也可以对不同的层制定不同的优化器参数如学习率等。
eg:
optim.SGD([
{'params': model.base.parameters()},
{'params': model.classifier.parameters(), 'lr': 1e-3}
], lr=1e-2, momentum=0.9)
上面例子中默认的学习率和动量为0.01和0.9,而classifier的学习率单独设置为了0.001.
注意:如果要将模型转到GPU中,即调用.cuda(),那么构造optimizer要放在model.cuda()之后,因为模型转到GPU后其参数会不同。
举一个较综合、详细的例子:
model=the_model()
model.cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
model.train()
epochs=10
for epoch in range(epochs):
total_loss=0.0
for input, target in dataset:
with torch.no_grad():
if torch.cuda.is_available():
Variable(torch.from_numpy(input).type(torch.FloatTensor)).cuda()#input target transfer to gpu,
#so does target
else:
#keep on cpu
optimizer.zero_grad()
output = model(input)
loss = loss_fn(output, target)
loss.backward()
optimizer.step()
total_loss+=loss
#save model to inference
torch.save(model.state_dict(),"path/to/save.pth")
#也可保存模型、优化器、epoch等以恢复训练
上述示例未加入lr_scheduler,可以将上面单个周期的训练封装,一个周期内包括训练和验证两部分,这两部分大体内容相同,然后每一个周期之后调用scheduler.step()。
调整学习率:
torch.optim.lr_scheduler提供了几种根据epoch数调整学习率的方法。
torch.optim.lr_scheduler.ReduceLROnPlateau允许基于某些验证度量来动态降低学习率。
lr_scheduler也有state_dict(),以dict的形式返回scheduler的状态。
其使用时的一般形式为:
>>> scheduler = ...
>>> for epoch in range(100):
>>> train(...)
>>> validate(...)
>>> scheduler.step()
注意优化器的step()要在scheduler的step()之前。
这里主要介绍两种,
第一类:
torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch=-1, verbose=False)
optimizer (Optimizer) – 封装的optimizer.
lr_lambda (function or list) – 一个关于epoch的计算函数, 如果optimizer.param_groups中有多个需要优化的参数组,那么这里就有一个由多个计算函数组成的list。
last_epoch (int) – 最后一个epoch的值. 默认: -1.
verbose (bool) – 如果为True打印输出每次更新的值,默认为Fasle。
示例:
# Assuming optimizer has two groups.
lambda1 = lambda epoch: epoch // 30
lambda2 = lambda epoch: 0.95 ** epoch
scheduler = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2])
for epoch in range(100):
train(...)
validate(...)
scheduler.step()
第二类:
torch.optim.lr_scheduler.StepLR(optimizer, step_size, gamma=0.1, last_epoch=-1, verbose=False)
step_size (int) – 多少个epoch后开始学习率decay
gamma (float) – 学习率衰退的因子,默认 0.1。每次衰退时乘以这个。
示例:
>>> # Assuming optimizer uses lr = 0.05 for all groups
>>> # lr = 0.05 if epoch < 30
>>> # lr = 0.005 if 30 <= epoch < 60
>>> # lr = 0.0005 if 60 <= epoch < 90
>>> # ...
>>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
>>> for epoch in range(100):
>>> train(...)
>>> validate(...)
>>> scheduler.step()
文章主要参考pytorch官方文档。