Pytorch网络模型权重初始化、保存与加载模型、加载预训练模型、按需设置学习率

前言

在我们对神经网络模型进行训练时，往往需要对模型进行初始化或者加载预训练模型。本文将对模型的权重初始化与加载预训练模型做一个学习记录，以便后续查询使用。

权重初始化

常见的初始化方法

PyTorch 在 torch.nn.init 中提供了常用的初始化方法函数，这里主要简要介绍Xavier初始化与kaiming初始化。

Xavier初始化

Xavier 初始化方法，方法来源于2010年的一篇论文《Understanding the difficulty of training deep feedforward neural networks》

公式推导是从“方差一致性”出发，初始化的分布有均匀分布和正态分布两种。

Xavier 均匀分布

torch.nn.init.xavier_uniform_(tensor, gain=1.0)

该初始化方法服从均匀分布$U\sim (-a,a)$，其中a为：
$$a=\mathrm{gain}\times \sqrt{\frac{6}{\mathrm{f}\mathrm{a}\mathrm{n}\_\mathrm{i}\mathrm{n}+\mathrm{f}\mathrm{a}\mathrm{n}\_\mathrm{o}\mathrm{u}\mathrm{t}}}$$

该初始化方法中有一个参数 gain，增益的大小是依据激活函数类型来设定
eg：

nn.init.xavier_uniform_(w, gain=nn.init.calculate_gain(‘relu’))

PS：上述初始化方法，也称为 Glorot initialization

使用方法示例：

for m in model.modules():
    if isinstance(m, (nn.Conv2d, nn.Linear)):
        nn.init.xavier_uniform_(m.weight)

Xavier正态分布

torch.nn.init.xavier_normal_(tensor, gain=1.0)

该初始化方法服从正态分布$\mathcal{N}\left(0,{\mathrm{s}\mathrm{t}\mathrm{d}}^2\right)$：
$$\mathrm{std}=\mathrm{gain}\times \sqrt{\frac{2}{\mathrm{f}\mathrm{a}\mathrm{n}\_\mathrm{i}\mathrm{n}+\mathrm{f}\mathrm{a}\mathrm{n}\_\mathrm{o}\mathrm{u}\mathrm{t}}}$$
使用方法示例：

for m in model.modules():
    if isinstance(m, (nn.Conv2d, nn.Linear)):
        nn.init.xavier_normal_(m.weight)

kaiming初始化

kaiming初始化，方法来源于2015年的一篇论文《 Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification》

公式推导同样从“方差一致性”出法，kaiming是针对xavier初始化方法在relu这一类激活函数表现不佳而提出的改进，详细可以参看论文。

kaiming均匀分布

torch.nn.init.kaiming_uniform_(tensor, a=0, mode=‘fan_in’, nonlinearity=‘leaky_relu’)

该初始化方法服从均匀分布$\mathcal{U}($ -bound, bound $)$：
$$\text{ bound }=\mathrm{gain}\times \sqrt{\frac{3}{\text{ fan\_mode }}}$$
其中，a为激活函数的负半轴的斜率，mode可选为fan_in 或fan_out, fan_in使正向传播时，方差一致; fan_out使反向传播时，方差一致。
nonlinearity 建议选择 relu 和 leaky_relu ，默认值为 leaky_relu

kaiming正态分布

torch.nn.init.kaiming_normal_(tensor, a=0, mode=‘fan_in’, nonlinearity=‘leaky_relu’)

该初始化方法服从正态分布$\mathcal{N}\left(0,{\mathrm{s}\mathrm{t}\mathrm{d}}^2\right)$：
$$\mathrm{s}\mathrm{t}\mathrm{d}=\frac{\text{ gain }}{\sqrt{\text{ fan\_mode }}}$$
其中，a为激活函数的负半轴的斜率，mode可选为fan_in 或fan_out, fan_in使正向传播时，方差一致; fan_out使反向传播时，方差一致。
nonlinearity 建议选择 relu 和 leaky_relu ，默认值为 leaky_relu

模型权重初始化

# 定义权值初始化
def initialize_weights(self):
    for m in self.modules():
        if isinstance(m, nn.Conv2d):
            torch.nn.init.xavier_normal_(m.weight.data)
            if m.bias is not None:
                m.bias.data.zero_()
        elif isinstance(m, nn.BatchNorm2d):
            m.weight.data.fill_(1)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            torch.nn.init.normal_(m.weight.data, 0, 0.01)
            m.bias.data.zero_()

保存与加载模型

pytorch在保存模型时，可以保存整个神经网络的的结构信息和模型参数信息，save的对象是网络net；也可以只保存神经网络的训练模型参数，save的对象是net.state_dict()。

# 保存和加载整个模型  
torch.save(model_object, 'model.pth')  
model = torch.load('model.pth')  
     
# 仅保存和加载模型参数  
torch.save(model_object.state_dict(), 'params.pth')  
model_object.load_state_dict(torch.load('params.pth')) 

加载预训练模型

# load params, 这里加载的是模型的参数，不是整个模型
pretrained_dict = torch.load('net_params.pkl')
# 仅保存了整个模型, 需要使用以下语句
# pretrained_dict = torch.load('net_params.pkl').state_dict()

# 获取当前网络的dict
net_state_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.update(pretrained_dict_1)

# 将包含预训练模型参数的字典"放"到新模型中
net.load_state_dict(net_state_dict)

pytorch预训练模型的简单修改与使用

以resnet预训练模型举例，resnet源代码的pytorch官方实现。 resnet网络最后一层分类层fc是对1000种类型进行划分，如果自己的数据集只有6类，可以只对fc层进行修改：

#调用模型
model = torchvision.models.resnet50(pretrained=True)

#提取fc层中固定的参数
fc_features = model.fc.in_features

#修改类别
model.fc = nn.Linear(fc_features, 6)

按需设置学习率

# ================================= #
#         按需设置学习率
# ================================= #

# 将fc3层的参数从原始网络参数中剔除
ignored_params = list(map(id, net.fc3.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params, net.parameters())

# 为fc3层设置需要的学习率
optimizer = optim.SGD([
    {'params': base_params},
    {'params': net.fc3.parameters(), 'lr': lr_init*10}],  lr_init, momentum=0.9, weight_decay=1e-4)

criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)     # 设置学习率下降策略

参考资料

1. 《Pytorch模型训练实用教程》
2. 《Understanding the difficulty of training deep feedforward neural
   networks》
3. 《 Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification》
4. PyTorch Documentation