保存预训练的模型参数:
torch.save(net.state_dict(), ‘net_params.pkl’)
加载模型:
pretrained_dict = torch.load(‘net_params.pkl’)
初始化:
net = Net() # 创建net
net_state_dict = net.state_dict() # 获取已创建net的state_dict
接着将pretrained_dict里不属于net_state_dict的键剔除掉:
pretrained_dict_1 = {k: v for k, v in pretrained_dict.items() if k in net_state_dict}
然后,用预训练模型的参数字典 对 新模型的参数字典net_state_dict 进行更新:
net_state_dict.update(pretrained_dict_1)
最后,将更新了参数的字典 “放”回到网络中:
net.load_state_dict(net_state_dict)
不同层设置不同的学习率:
ignored_params = list(map(id, net.fc3.parameters())) # 返回的是parameters的 内存地址
base_params = filter(lambda p: id§ not in ignored_params, net.parameters())
optimizer = optim.SGD([
{‘params’: base_params},
{‘params’: net.fc3.parameters(), ‘lr’: 0.001*10}], 0.001, momentum=0.9, weight_decay=1e-4)