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

您所在的位置：网站首页 › facerig人脸一直在初始化 › Pytorch网络模型权重初始化、保存与加载模型、加载预训练模型、按需设置学习率

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

2024-01-15 23:38| 来源: 网络整理| 查看: 265

前言

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

权重初始化常见的初始化方法

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 ∼ ( − a , a ) U\sim(-a,a) U∼(−a,a)，其中a为： a = gain ⁡ × 6 fan_in ⁡ + fan_out ⁡ a=\operatorname{gain} \times \sqrt{\frac{6}{\operatorname{fan\_in} +\operatorname{fan\_out}}} a=gain×fan_in+fan_out6

该初始化方法中有一个参数 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)

该初始化方法服从正态分布 N ( 0 , s t d 2 ) \mathcal{N}\left(0, \mathrm{std}^{2}\right) N(0,std2)： std ⁡ = gain ⁡ × 2 fan_in ⁡ + fan_out ⁡ \operatorname{std}=\operatorname{gain} \times \sqrt{\frac{2}{\operatorname{fan\_in} +\operatorname{fan\_out}}} std=gain×fan_in+fan_out2 使用方法示例：

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’)

该初始化方法服从均匀分布 U ( \mathcal{U}( U( -bound, bound ) ) )： bound = gain ⁡ × 3 fan_mode \text { bound }=\operatorname{gain} \times \sqrt{\frac{3}{\text { fan\_mode }}} bound =gain× fan_mode 3 其中，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’)

该初始化方法服从正态分布 N ( 0 , s t d 2 ) \mathcal{N}\left(0, \mathrm{std}^{2}\right) N(0,std2)： s t d = gain fan_mode \mathrm{std}=\frac{\text { gain }}{\sqrt{\text { fan\_mode }}} std= fan_mode gain 其中，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) # 设置学习率下降策略参考资料《Pytorch模型训练实用教程》《Understanding the difficulty of training deep feedforward neural networks》《 Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification》PyTorch Documentation

【本文地址】

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

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

今日新闻

推荐新闻