第二十五篇—加载预训练权重
YouOnly_LiveOnce
编辑于 2023年04月21日 06:09
收录于文集
共28篇

在Pytorch中，只有可学习参数的层(卷积层、线性层、BN层等)才有state_dict，model.state_dict()会以有序字典OrderedDict形式返回模型训练过程中学习的权重weight和偏置bias参数（参考网页链接​），如下述代码所示：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn 
from torch.nn import init 

# 定义原模型
class Conv(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.LeakyReLU()
        self.my_conv2=nn.Conv2d(out_ch,out_ch,k,s,p)
        self.my_bn2=nn.BatchNorm2d(out_ch)
        self.my_act2=nn.LeakyReLU()
    def forward(self,x):
        y=self.my_act(self.my_bn(self.my_conv(x)))
        return self.my_act2(self.my_bn2(self.my_conv2(y)))

# 创建模型
net=Conv(in_ch=1,out_ch=1,k=3)
# 模型参数初始化
for name,param in net.named_parameters():
    if "weight" in name:
        init.constant_(param,val=6)
    if "bias" in name:
        init.constant_(param,val=0)
# 打印可学习层的参数
print(net.state_dict())
# 保存可学习层的参数
torch.save(net.state_dict(),"Conv.pt")
复制成功

上述代码定义的模型中，只有卷积层和BN层具有可学习参数，所以net.state_dict()只会保存这两层的参数，而激活函数层的参数则不会保存。BN层除了权重weight和偏置bias参数，还会保存训练阶段统计的均值(running_mean)、训练阶段统计的方差(running_val)、训练阶段的batch数目(num_batches_tracked)，其中，weight和bias属于可学习参数，需要进行训练，而running_mean、running_val和num_batches_tracked三个参数则不需要训练，只是训练阶段的统计值。，如下图所示：

当我们对网络模型结构进行优化改进时，如果改进的部分不包含可学习的层，那么可以直接加载预训练权重。如：如果我们对上述代码的Conv模型进行改进，将激活函数层改为nn.Hardswish()，因为不包含可学习的参数，所以改进的模型的state_dict()没有改变，仍然可以直接加载Conv模型的权重文件，如下代码所示：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn 
from torch.nn import init 

# 定义原模型
class Conv(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.LeakyReLU()
        self.my_conv2=nn.Conv2d(out_ch,out_ch,k,s,p)
        self.my_bn2=nn.BatchNorm2d(out_ch)
        self.my_act2=nn.LeakyReLU()
    def forward(self,x):
        y=self.my_act(self.my_bn(self.my_conv(x)))
        return self.my_act2(self.my_bn2(self.my_conv2(y)))

# 创建模型
net=Conv(in_ch=1,out_ch=1,k=3)
# 模型参数初始化
for name,param in net.named_parameters():
    if "weight" in name:
        init.constant_(param,val=6)
    if "bias" in name:
        init.constant_(param,val=0)
# 保存可学习的参数
torch.save(net.state_dict(),"Conv.pt")

# 定义改进的模型
class Conv_improve(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv_improve,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.Hardswish()	# 改进1
        self.my_conv2=nn.Conv2d(out_ch,out_ch,k,s,p)
        self.my_bn2=nn.BatchNorm2d(out_ch)
        self.my_act2=nn.Hardswish()	# 改进2
    def forward(self,x):
        y=self.my_act(self.my_bn(self.my_conv(x)))
        return self.my_act2(self.my_bn2(self.my_conv2(y)))
# 创建模型
net2=Conv_improve(in_ch=1,out_ch=1,k=3)
# 加载预训练权重
ckpt=torch.load("Conv.pt")
net2.load_state_dict(ckpt)
print(net2.state_dict())
复制成功

结果如下：

当我们改进的部分改变了可学习的参数时，如果直接加载预训练权重就会发生不匹配的错误，如下代码所示：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn 
from torch.nn import init 

# 定义原模型
class Conv(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.LeakyReLU()
        self.my_conv2=nn.Conv2d(out_ch,out_ch,k,s,p)
        self.my_bn2=nn.BatchNorm2d(out_ch)
        self.my_act2=nn.LeakyReLU()
    def forward(self,x):
        y=self.my_act(self.my_bn(self.my_conv(x)))
        return self.my_act2(self.my_bn2(self.my_conv2(y)))

# 创建模型
net=Conv(in_ch=1,out_ch=1,k=3)
# 模型参数初始化
for name,param in net.named_parameters():
    if "weight" in name:
        init.constant_(param,val=6)
    if "bias" in name:
        init.constant_(param,val=0)
# 保存可学习的参数
torch.save(net.state_dict(),"Conv.pt")

# 定义改进的模型
class Conv_improve(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv_improve,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.LeakyReLU()
    def forward(self,x):
        return self.my_act(self.my_bn(self.my_conv(x)))
# 创建模型
net2=Conv_improve(in_ch=1,out_ch=1,k=3)
# 加载预训练权重
ckpt=torch.load("Conv.pt")
net2.load_state_dict(ckpt)
复制成功

结果如下：

这时我们需要遍历预训练文件的每一层参数，将能够匹配成功的参数提取出来，再进行加载就可以了，如下代码所示：

代码块
Python
自动换行
复制代码
import torch 
import torch.nn as nn 
from torch.nn import init 

# 定义原模型
class Conv(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.LeakyReLU()
        self.my_conv2=nn.Conv2d(out_ch,out_ch,k,s,p)
        self.my_bn2=nn.BatchNorm2d(out_ch)
        self.my_act2=nn.LeakyReLU()
    def forward(self,x):
        y=self.my_act(self.my_bn(self.my_conv(x)))
        return self.my_act2(self.my_bn2(self.my_conv2(y)))

# 创建模型
net=Conv(in_ch=1,out_ch=1,k=3)
# 模型参数初始化
for name,param in net.named_parameters():
    if "weight" in name:
        init.constant_(param,val=6)
    if "bias" in name:
        init.constant_(param,val=0)
# 保存可学习的参数
torch.save(net.state_dict(),"Conv.pt")

# 定义改进的模型
class Conv_improve(nn.Module):
    def __init__(self,in_ch,out_ch,k=1,s=1,p=0):
        super(Conv_improve,self).__init__()
        self.my_conv=nn.Conv2d(in_ch,out_ch,k,s,p)
        self.my_bn=nn.BatchNorm2d(out_ch)
        self.my_act=nn.LeakyReLU()
    def forward(self,x):
        return self.my_act(self.my_bn(self.my_conv(x)))
# 创建模型
net2=Conv_improve(in_ch=1,out_ch=1,k=3)
# 加载预训练权重
ckpt=torch.load("Conv.pt")

# 定义参数提取函数
def intersect_dicts(da,db):
    '''
    da:预训练权重;
    db:模型权重;
    匹配条件: key相同; value的shape相同
    '''
    return {k:v for k,v in da.items() if k in db and v.shape==db[k].shape}
# 参数筛选
ckpt=intersect_dicts(ckpt,net2.state_dict())
# 加载可匹配的参数
net2.load_state_dict(ckpt)
# 打印参数
print(net2.state_dict())
复制成功

结果如下：

神经网络 深度学习 pytorch 迁移学习 预训练权重
本文为我原创，未经授权禁止转载
cv23188158