在模型中如果设置了requires_grad=True,则表示该层要进行梯度计算,标记为False则不计算梯度,在迁移学习中一般会设置成False,这样会大量减少算力。
而optim中的parameters是定义要对那些层进行参数优化
一般在迁移学习的代码过程中我们会先把加载的模型所有层定义成requires_grad=False,再
将模型编辑成我们需要的样子,例如将全连接层的输出定义成我们要的输出。
然后根据没有冻结的层创建优化器。
# 加载模型
model_ft = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
# 冻结模型
for param in model_ft.parameters():
param.requires_grad = False
# 编辑模型
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, 10)
# 取出未冻结的层
param_update= []
for param in model_ft.parameters():
if param.requires_grad:
param_update.append(param)
# 定义优化器
optimizer_ft = optim.Adam(param_update, lr=0.001)
这样做的好处是可以节省计算梯度和优化的算力。
但是这样做有一个问题:
我在代码中会保存优化器以便于之后的继续测试
state = {
'state_dict': mymodel.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(state, "save_test.pth")
读取的时候
model_ft = torchvision.models.resnet18(weights=torchvision.models.ResNet18_Weights.DEFAULT)
# 此时不冻结任何层 对全部层进行训练
optimizer = optim.Adam(model_ft.parameters(), lr=1e-2)
checkpoint = torch.load("save_test.pth")
optimizer.load_state_dict(checkpoint['optimizer'])
如果保存的时候对优化器设置的是只针对某几个层进行优化,而读取的时候希望对所有的层进行优化程序就会报错:
File "D:\anaconda3\envs\pytorch_gpu\lib\site-packages\torch\optim\optimizer.py", line 171, in load_state_dict
raise ValueError("loaded state dict contains a parameter group "
ValueError: loaded state dict contains a parameter group that doesn't match the size of optimizer's group
所以,在读取优化器记录的时候要注意,如果保存的数据是针对某些层保存的,而加载的时候又希望对所有层进行训练,这时候优化器数据就没必要加载了,要么写个try,要么就把这行注释掉。