PyTorch保存模型与加载模型+Finetune预训练模型使用

2021-12-12 22:18:58

参数初始化参

数的初始化其实就是对参数赋值。而我们需要学习的参数其实都是Variable，它其实是对Tensor的封装，同时提供了data，grad等借口，这就意味着我们可以直接对这些参数进行操作赋值了。这就是PyTorch简洁高效所在。所以我们可以进行如下操作进行初始化，当然其实有其他的方法，但是这种方法是PyTorch作者所推崇的：

def weight_init(m):

# 使用isinstance来判断m属于什么类型

    if isinstance(m, nn.Conv2d):

        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels

        m.weight.data.normal_(0, math.sqrt(2. / n))

    elif isinstance(m, nn.BatchNorm2d):

# m中的weight，bias其实都是Variable，为了能学习参数以及后向传播

        m.weight.data.fill_(1)

        m.bias.data.zero_()

Finetune

往往在加载了预训练模型的参数之后，我们需要finetune模型，可以使用不同的方式finetune。
局部微调：有时候我们加载了训练模型后，只想调节最后的几层，其他层不训练。其实不训练也就意味着不进行梯度计算，PyTorch中提供的requires_grad使得对训练的控制变得非常简单。

model = torchvision.models.resnet18(pretrained=True)
for param in model.parameters():
param.requires_grad = False
# 替换最后的全连接层，改为训练100类
# 新构造的模块的参数默认requires_grad为True
model.fc = nn.Linear(512, 100)

# 只优化最后的分类层
optimizer = optim.SGD(model.fc.parameters(), lr=1e-2, momentum=0.9)

全局微调：有时候我们需要对全局都进行finetune，只不过我们希望改换过的层和其他层的学习速率不一样，这时候我们可以把其他层和新层在optimizer中单独赋予不同的学习速率。比如：

ignored_params = list(map(id, model.fc.parameters()))
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())

optimizer = torch.optim.SGD([
            {'params': base_params},
            {'params': model.fc.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)
其中base_params使用1e-3来训练，model.fc.parameters使用1e-2来训练，momentum是二者共有的。

PyTorch预训练

加载部分预训练模型：其实大多数时候我们需要根据我们的任务调节我们的模型，所以很难保证模型和公开的模型完全一样，但是预训练模型的参数确实有助于提高训练的准确率，为了结合二者的优点，就需要我们加载部分预训练模型。

pretrained_dict = model_zoo.load_url(model_urls['resnet152']) 
model_dict = model.state_dict() 

# 将pretrained_dict里不属于model_dict的键剔除掉 
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} 
# 更新现有的model_dict 
model_dict.update(pretrained_dict) 
# 加载我们真正需要的state_dict 
model.load_state_dict(model_dict)

pytorch finetune模型

如果模型的key值和在大数据集上训练时的key值是一样的

我们可以通过下列算法进行读取模型

model_dict = model.state_dict()

pretrained_dict = torch.load(model_path)
# 1. filter out unnecessary keys
    diff = {k: v for k, v in model_dict.items() if \
            k in pretrained_dict and pretrained_dict[k].size() == v.size()}
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and model_dict[k].size() == v.size()}
    pretrained_dict.update(diff)
    # 2. overwrite entries in the existing state dict
    model_dict.update(pretrained_dict)
    # 3. load the new state dict
    model.load_state_dict(model_dict)

如果模型的key值和在大数据集上训练时的key值是不一样的，但是顺序是一样的

model_dict = model.state_dict()

pretrained_dict = torch.load(model_path)
    keys = []
    for k,v in pretrained_dict.items():
        keys.append(k)
    i = 0
    for k,v in model_dict.items():
        if v.size() == pretrained_dict[keys[i]].size():
            print(k, ',', keys[i])
            model_dict[k]=pretrained_dict[keys[i]]
        i = i + 1
    model.load_state_dict(model_dict)

如果模型的key值和在大数据集上训练时的key值是不一样的，但是顺序是也不一样的

自己找对应关系，一个key对应一个key的赋值

码农公寓

PyTorch保存模型与加载模型+Finetune预训练模型使用