PyTorch | 模型加载/参数初始化/Finetune

1、模型保存/加载

1.1 全部模型参数

训练过程当中，有时候会因为各类缘由中止训练，这时候咱们训练过程当中就须要注意将每一轮epoch的模型保存(通常保存最好模型与当前轮模型)。通常使用pytorch里面推荐的保存方法。该方法保存的是模型的参数。

#保存模型到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
....
....

1.2 部分模型参数

在不少时候，咱们加载的是已经训练好的模型，而训练好的模型可能与咱们定义的模型不彻底同样，而咱们只想使用同样的那些层的参数。

有几种解决方法：

（1）直接在训练好的模型开始搭建本身的模型，就是先加载训练好的模型，而后再它基础上定义本身的模型；

model_ft = models.resnet18(pretrained=use_pretrained)
self.conv1 = model_ft.conv1
self.bn = model_ft.bn
... ...

（2) 本身定义好模型，直接加载模型

#第一种方法：
mymodelB = TheModelBClass(*args, **kwargs)
# strict=False，设置为false,只保留键值相同的参数
mymodelB.load_state_dict(model_zoo.load_url(model_urls['resnet18']), strict=False)

#第二种方法：
# 加载模型
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)

# 方法2可能更直观一些

2、参数初始化

第二个问题是参数初始化问题，在不少代码里面都会使用到，毕竟不是全部的都是有预训练参数。这时就须要对不是与预训练参数进行初始化。pytorch里面的每一个Tensor实际上是对Variabl的封装，其包含data、grad等接口，所以能够用这些接口直接赋值。这里也提供了怎样把其余框架(caffe/tensorflow/mxnet/gluonCV等)训练好的模型参数直接赋值给pytorch.其实就是对data直接赋值。

pytorch提供了初始化参数的方法：

 def weight_init(m):
    if isinstance(m,nn.Conv2d):
        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        m.weight.data.normal_(0,math.sqrt(2./n))
    elif isinstance(m,nn.BatchNorm2d):
        m.weight.data.fill_(1)
        m.bias.data.zero_()

但通常若是没有很大需求初始化参数，也没有问题(不肯定性能是否有影响的状况下)，pytorch内部是有默认初始化参数的。

3、模型微调/Finetune

最后是微调，平时作实验，至少backbone是用预训练的模型，将其用做特征提取器，或者在它上面作微调。

用于特征提取的时候，要求特征提取部分参数不进行学习，而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)

 

参考：

https://zhuanlan.zhihu.com/p/48524007

https://zhuanlan.zhihu.com/p/38056115