前言
在我们对神经网络模型进行训练时,往往需要对模型进行初始化或者加载预训练模型。本文将对模型的权重初始化与加载预训练模型做一个学习记录,以便后续查询使用。
权重初始化
常见的初始化方法
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) # 设置学习率下降策略