目录
模型保存/加载所有模型参数
第一种方法
#保存模型到checkpoint.pth.tar
torch.save(model.module.state_dict(), ‘checkpoint.pth.tar’)
对应的加载模型方法为(这种方法需要先反序列化模型获取参数字典,因此必须先load模型,再load_state_dict):
mymodel.load_state_dict(torch.load(‘checkpoint.pth.tar’))
例子:说明如何在inference AND/OR resume train使用
#保存模型的状态,可以设置一些参数,后续可以使用
state = {'epoch': epoch + 1,#保存的当前轮数
'state_dict': mymodel.state_dict(),#训练好的参数
'optimizer': optimizer.state_dict(),#优化器参数,为了后续的resume
'best_pred': best_pred#当前最好的精度
,....,...}
#保存模型到checkpoint.pth.tar
torch.save(state, ‘checkpoint.pth.tar’)
#如果是best,则复制过去
if is_best:
shutil.copyfile(filename, directory + 'model_best.pth.tar')
checkpoint = torch.load('model_best.pth.tar')
model.load_state_dict(checkpoint['state_dict'])#模型参数
optimizer.load_state_dict(checkpoint['optimizer'])#优化参数
epoch = checkpoint['epoch']#epoch,可以用于更新学习率等
#有了以上的东西,就可以继续重新训练了,也就不需要担心停止程序重新训练。
train/eval
第二种方法(不推荐)
#保存
torch.save(mymodel,‘checkpoint.pth.tar’)
#加载
mymodel = torch.load(‘checkpoint.pth.tar’)
模型保存/加载部分模型参数
第一种方法:
(1)直接在训练好的模型开始搭建自己的模型,就是先加载训练好的模型,然后再它基础上定义自己的模型;
model_ft = models.resnet18(pretrained=use_pretrained)
self.conv1 = model_ft.conv1
self.bn = model_ft.bn
... ...
#又比如修改Alexnet最后一层数据集的类别数
alexnet.classifier._modules['6'] = nn.Linear(4096,250)
第二种方法:
(2) 自己定义好模型,直接加载模型
# 加载模型
model_pretrained = models.resnet18(pretrained=use_pretrained)
# mymodel's state_dict,
# 如: conv1.weight
# conv1.bias
mymodelB_dict = mymodelB.state_dict()
# 将model_pretrained的键与自定义模型的键进行比较,剔除不同的
pretrained_dict = {k: v for k, v in model_pretrained.items() if k in mymodelB_dict}
# 更新现有的model_dict
mymodelB_dict.update(pretrained_dict)
# 加载我们真正需要的state_dict
mymodelB.load_state_dict(mymodelB_dict)
Finetune
用于特征提取的时候,要求特征提取部分参数不进行学习,而pytorch提供了requires_grad参数用于确定是否进去梯度计算,也即是否更新参数。以下以minist为例,用resnet18作特征提取:
#加载预训练模型
model = torchvision.models.resnet18(pretrained=True)
#遍历每一个参数,将其设置为不更新参数,即不学习
for param in model.parameters():
param.requires_grad = False
# 将全连接层改为mnist所需的10类,注意:这样更改后requires_grad默认为True
model.fc = nn.Linear(512, 10)
# 优化
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)
用于全局精调时,我们一般对不同的层需要设置不同的学习率,预训练的层学习率小一点,其他层大一点。这要怎么做呢?
# 加载预训练模型
model = torchvision.models.resnet18(pretrained=True)
model.fc = nn.Linear(512, 10)
# 参考:https://blog.csdn.net/u012759136/article/details/65634477
ignored_params = list(map(id, model.fc.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, model.parameters())
# 对不同参数设置不同的学习率
params_list = [{'params': base_params, 'lr': 0.001},]
params_list.append({'params': model.fc.parameters(), 'lr': 0.01})
optimizer = torch.optim.SGD(params_list,
0.001,
momentum=args.momentum,
weight_decay=args.weight_decay)
目前Pytorch模型
(1)torchvision
https://github.com/pytorch/vision/tree/master/torchvision/models
包含了alexnet/densenet各种版本(densenet121/densenet169/densenet201/densenet161)/inception_v3/resnet各种版本(resnet18', 'resnet34', 'resnet50', 'resnet101','resnet152')/SqueezeNet各种版本( 'squeezenet1_0', 'squeezenet1_1')/VGG各种版本( 'vgg11', 'vgg11_bn', 'vgg13', 'vgg13_bn', 'vgg16', 'vgg16_bn','vgg19_bn', 'vgg19')
(2)其他预训练好的模型,如,SENet/NASNet等。
Cadene/pretrained-models.pytorchgithub.com
(3)gluonCV转pytorch的模型,包括,分类网络,分割网络等,这里的精度均比其他框架高几个百分点。