联系方式:
e-mail: FesianXu@163.com
QQ: 973926198
github: https://github.com/FesianXu

文章目录

  • 基模型参数加载
       从持久化模型开始
       加载模型
       部分加载模

  • 模型Fine-Tune

  • 给每一层或者每个模型设置不同的学习率

  • Pytorch内置的模型

  • Reference

在使用pytorch的时候,经常有需要使用一些通用的模型模块作为子模块,比如著名的resnet,densenet,alexnet,inception等等,在使用这些模型的时候,通常希望可以加载该模型在别的数据集(如ImageNet)上进行训练后的权值参数,以便加速整个模型训练过程[1]。在此为了简便,称之为这些预训练模型为基模型,加载基模型参数这个过程按照需求大致可以分为两类:

  1. 整个模型的预训练参数加载
  2. 部分模型的预训练参数加载
    在**对模型的参数进行fine-tune[**2-3]的时候,按照需求也可以大致分为两类:
  3. 固定整个基模型的参数,调节其他模型的参数
  4. 固定部分基模型的参数,调节其他模型的参数

1 基模型参数加载

1.1 从持久化模型开始

在pytorch中,保存一个模型的参数特别容易,用torch.save()即可,例如:

model = CNNNet(params)	# 作者自定义的模型,没提供源码
opt = torch.optim.Adam(model.parameters(), lr=1e-4)
model.train()
# here we train the models, skip these codes
saved_dict = {
	'model': model.state_dict(),
	'opt': opt.state_dict()
}
torch.save(saved_dict, './model.pth.tar')

我们发现,torch.save()保存的是一个字典,其中的keys可以自定义。这里有一点要注意的是,如果你用的优化器是例如Adam优化器[5-6]这类内部有参数需要持久化的,最好也将其保存下来

1.2 加载模型

如果是需要加载整个模型,直接用torch.load()和model.load_state_dict()即可,如:

model = CNNNet(params)
opt = torch.optim.Adam(model.parameters(), lr=1e-4)
# yes you also need to define the model and optimizer
checkpoint = torch.load('./model.pth.tar')
# here, checkpoint is a dict with the keys you defined before
model.load_state_dict(checkpoint['model'])
opt.load_state_dict(checkpoint['opt'])

这个过程中torch.load()只是负责读取模型参数,而用model.load_state_dict()进行加载,这个加载是按照名字进行索引的,如果名字对不上或者是参数的形状,类型对不上,就会报错。我们可以打印出其名字进行观察,如:

for name in checkpoint ['model'].keys():
    print(name)

输出如:

# stgcn是模型的名字,具体到不同的模型名字可能不同
# weight_model可能是个调用了子模块A的另一个模块B的属性
# conv_models是子模块A中调用了卷积与BN操作的属性
stgcn.weight_model.conv_models.0.conv_layer.weight
stgcn.weight_model.conv_models.0.conv_layer.bias
stgcn.weight_model.conv_models.1.conv_layer.weight
stgcn.weight_model.conv_models.1.conv_layer.bias
stgcn.weight_model.conv_models.1.batch_norm.weight
stgcn.weight_model.conv_models.1.batch_norm.bias

如果定义的模型(网络结构)和持久化的模型的参数名,形状,类型(网络参数)都能完全符合,就能正确加载。同时还注意到,变量的名字是由pytorch自行命名的,其命名根据就是你的变量名字,比如:

import torch.nn as nn
class model_A(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10,10)
        
class model_B(nn.Module):
    def __init__(self):
        super().__init__()
        self.sub_model = model_A()
        self.fc = nn.Linear(10,1)
model = model_B()

那么如果你打印prin(model_B),你就会发现子模块的名字为

sub_model.fc.weight
sub_model.fc.bias
fc.weight
fc.bias

我们观察到名字是以你的变量标识符命名的,这点和TensorFlow的命名机制完全不同,请注意。同时我们还观察到,只要是权值weight其命名后缀都是weight,同样偏置bias的后缀是bias,因此根据此,可以单独对权值进行L2正则[7],具体过程见[8]。

1.3 部分加载模型

根据上面的分析,我们便发现只要过滤掉不需要加载的模型的名字,即可实现部分模型加载了,例子如:

model = CNNNet()
checkpoint = torch.load('./model.pth.tar')
for name, params in model.stgcn.st_gcn_networks.named_parameters():
	 params_name = 'stgcn.st_gcn_networks.'+name
    if params_name in model.state_dict():
    	model.state_dict()[params_name].copy_(checkpoint['model'][params_name])

我们发现,通过这个代码,我们可以仅对model.stgcn.st_gcn_networks这个子模块的参数进行加载,而其他的参数保持初始化情况不变。

2 模型Fine-Tune

在模型的Fine-Tune(微调)或者联合调试过程中,我们经常需要固定某个模型的参数,而去调整其他模型的参数,主要方法有:

  • 通过切断某个模块的梯度流,但是这个会导致该模型前面的所有模型也没有梯度。
  • 通过设置某个模块的所有变量的requires_grad=False
  • 在优化器内设置需要进行梯度更新的变量。
    笔者在实践过程中最常用的是第三种方法,暂时只介绍第三种方法。代码很简单,例子如下:
trainable_vars = list(model.stgcn.weight_model.parameters())+ \
                 list(model.stgcn.fcn.parameters())+ \
                 list(model.stgcn.data_bn.parameters())+ \
                 list(model.stgcn.dim_map.parameters())+ \
                 list(model.aux_cls.parameters())                
opt = torch.optim.SGD(trainable_vars , lr=1e-4, momentum=0.9)

简单粗暴,但是其实很好用,当需要训练的变量很多,而需要固定的变量很少的时候,可以用对整个模型参数求补的方式求得,这里不多介绍了。
当需要对整个模型进行微调时,只需要:

opt = torch.optim.SGD(model.parameters(), lr=1e-6, momentum=0.9)

给每一层或者每个模型设置不同的学习率

在模型训练过程中,有些模块,比如对抗生成网络GAN[9]的生成器和判别器经常需要设置不同的学习率,以求得更好的效果或者不同模型之间的平衡。详细内容见我以前文章[8]中所述,这里不再累述。

Pytorch内置的模型

pytorch内置有一些经常使用的模型和其在大规模数据集上的预训练参数,只需要安装torchvision便可轻松使用,模型有:

  • resnet: resnet18,resnet34,resnet50,resnet101,resnet152
  • vgg: vgg11,vgg13, vgg16, vgg19
  • alexnet
  • densenet
  • inception
  • squeezenet

具体模型定义见:Github click me
使用方法很简单,如:

import torchvision.models as models
resnet18 = models.resnet18(pretrained=False)

在这里如果指定pretrained=True可以联网加载预训练模型,但是由于因为你懂得的原因,所以需要你懂得的辅助工具,建议读者自行去下载模型的文件后手动加载。模型文件的地址可以在模型定义文件中找到,如[https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py]中的resnet模型:

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}

Reference

[1]. Kaiming He, Ross Girshick, Piotr Dollár. Rethinking ImageNet Pre-training[J]. arXiv preprint, https://arxiv.org/abs/1811.08883
[2]. 迁移学习与fine-tuning有什么区别?
[3]. Fine tuning
[4]. pytorch
[5]. Adam 算法
[6]. Kingma D P, Ba J. Adam: A method for stochastic optimization[J]. arXiv preprint arXiv:1412.6980, 2014.
[7]. 曲线拟合问题与L2正则
[8]. pytorch中的L2和L1正则化,自定义优化器设置等操作
[9]. Goodfellow I, Pouget-Abadie J, Mirza M, et al. Generative adversarial nets[C]//Advances in neural information processing systems. 2014: 2672-2680.

Logo

开放原子开发者工作坊旨在鼓励更多人参与开源活动,与志同道合的开发者们相互交流开发经验、分享开发心得、获取前沿技术趋势。工作坊有多种形式的开发者活动,如meetup、训练营等,主打技术交流,干货满满,真诚地邀请各位开发者共同参与!

更多推荐