Pytorch学习笔记——优化器的使用
如何建立优化器
torch中优化器在torch.optim
中实现。从库中导入对应的优化器即可,此处需要注意的是,如果模型是在cuda上运行的,则优化器需要在模型移动到GPU上后再建立。
建立优化器的方式如下:
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = torch.optim.Adam([var1, var2], lr=0.0001)
一些特殊用法
当需要对模型的不同地方用不同的学习率训练时,给优化器传入的不再是模型的参数,而是传入一个可迭代对象(列表)。该列表中的元素为字典,字典必须包含一个一个params
的键,对应的值为模型的一部分参数,其他的键值则应该匹配该优化器可以接受的参数,例如:
optim.SGD([
{'params': model.base.parameters()},
{'params': model.classifier.parameters(), 'lr': 1e-3}
], lr=1e-2, momentum=0.9)
用法
所有的优化器都实现了.step()
方法,调用该方法将会更新参数。该方法有两种用法,在此介绍最为常用的一种用法:optimizer.step()
这是大多数优化器均支持的一个简化版本,当计算梯度后,可调用该方法。例如:
for input, target in dataset:
optimizer.zero_grad() # 梯度清零
output = model(input)
loss = loss_fn(output, target)
loss.backward() # 反向传播
optimizer.step()