模型权重保存、加载、冻结(pytorch)

2024-07-11 06:58| 来源: 网络整理| 查看: 265

1. 保存整个网络 torch.save(net, PATH) model = torch.load(PATH) 2. 保存网络中的参数（速度快，占空间小） torch.save(net.state_dict(),PATH) model_dict = model.load_state_dict(torch.load(PATH))

model.state_dict函数会以有序字典OrderedDict形式返回模型训练过程中学习的权重weight和偏置bias参数，只有带有可学习参数的层（卷积层、全连接层等），以及注册的缓存（batchnorm的运行平均值）在state_dict 中才有记录。以下面的LeNet为例：

import torch.nn as nn import torch.nn.functional as F class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.conv1 = nn.Conv2d(3, 16, 5) self.pool1 = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(16, 32, 5) self.pool2 = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(32 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = F.relu(self.conv1(x)) # input(3, 32, 32) output(16, 28, 28) x = self.pool1(x) # output(16, 14, 14) x = F.relu(self.conv2(x)) # output(32, 10, 10) x = self.pool2(x) # output(32, 5, 5) x = x.view(-1, 32 * 5 * 5) # output(32*5*5) x = F.relu(self.fc1(x)) # output(120) x = F.relu(self.fc2(x)) # output(84) x = self.fc3(x) # output(10) return x net = LeNet() # 打印可学习层的参数 print(net.state_dict().keys())

上面的模型中，只有卷积层和全连接层具有可学习参数，所以net.state_dict()只会保存这两层的参数，而激活函数层的参数则不会保存。层的名字是上面实例化时确定的，如果是利用nn.Sequential定义多个层时，用层的位置索引表示每个层，如下所示：

示例：用nn.Sequential搭建模型时的state_dict

import torch.nn as nn import torch.nn.functional as F class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.feature = nn.Sequential( nn.Conv2d(3, 16, 5), nn.MaxPool2d(2, 2), nn.Conv2d(16, 32, 5), nn.MaxPool2d(2, 2)) self.fc1 = nn.Linear(32 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.feature(x) # input(3, 32, 32) x = x.view(-1, 32 * 5 * 5) # output(32*5*5) x = F.relu(self.fc1(x)) # output(120) x = F.relu(self.fc2(x)) # output(84) x = self.fc3(x) # output(10) return x net = LeNet() # 打印可学习层的参数 print(net.state_dict().keys())　

★模型加载

当我们对网络模型结构进行优化改进时，如果改进的部分不包含可学习的层，那么可以直接加载预训练权重。如：如果我们对上述lenet模型进行改进，将激活函数层改为nn.Hardswish()，因为不包含可学习的参数，所以改进的模型的state_dict()没有改变，仍然可以直接加载lenet模型的权重文件。当我们改进的部分改变了可学习的参数时，如果直接加载预训练权重就会发生不匹配的错误，比如：卷积的维度改变后会报错 size mismatch for conv.weight...（2）新增一些层后会出现 Unexpected key(s) in state_dict等

解决方案：遍历预训练文件的每一层参数，将能够匹配成功的参数提取出来，再进行加载。

import torch import torch.nn as nn import torch.nn.functional as F class LeNet_new(nn.Module): def __init__(self): super(LeNet_new, self).__init__() self.conv1 = nn.Conv2d(3, 16, 5) self.pool1 = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(16, 32, 5) self.pool2 = nn.MaxPool2d(2, 2) def forward(self, x): x = F.hardswish(self.conv1(x)) # input(3, 32, 32) output(16, 28, 28) x = self.pool1(x) # output(16, 14, 14) x = F.hardswish(self.conv2(x)) # output(32, 10, 10) x = self.pool2(x) # output(32, 5, 5) return x def intersect_dicts(da, db): return {k: v for k, v in da.items() if k in db and v.shape == db[k].shape} net = LeNet_new() state_dict = torch.load("Lenet.pth") # 加载预训练权重 print(state_dict.keys()) state_dict = intersect_dicts(state_dict, net.state_dict()) # 筛选权重参数 print(state_dict.keys()) net.load_state_dict(state_dict, strict=False) # 模型加载预训练权重中可用的权重

3. 保存网络参数，同时保存优化器参数、损失值等（方便追加训练）

如果还想保存某一次训练采用的优化器、epochs等信息，可将这些信息组合起来构成一个字典，然后将字典保存起来

# 保存 save_file = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "lr_scheduler": lr_scheduler.state_dict(), "epoch": epoch, "args": args} torch.save(save_file, "save_weights/model_{}.pth".format(epoch)) # 加载 checkpoint = torch.load(path, map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) args.start_epoch = checkpoint['epoch'] + 1 4. 冻结训练

在加载预训练权重后，可能需要固定一部分模型的参数，只更新另一部分参数。有两种思路实现这个目标，一个是设置不要更新参数的网络层为requires_grad = False，另一个就是在定义优化器时只传入要更新的参数。最优写法时：将不更新的参数的requires_grad设置为False，同时不将该参数传入optimizer。

示例：LeNet网络+MNIST手写识别+预训练模型加载+冻结训练

import torch from torch import nn from torchvision import datasets, transforms from torch.utils.data import DataLoader import torch.nn.functional as F from tqdm import tqdm transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) train_data = datasets.MNIST(root='../dataset', train=True, transform=transform, download=True) train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True) test_data = datasets.MNIST(root='../dataset', train=False, transform=transform, download=True) test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=False) class LeNet(nn.Module): def __init__(self): super(LeNet, self).__init__() self.feature = nn.Sequential( nn.Conv2d(1, 16, 5), nn.MaxPool2d(2, 2), nn.Conv2d(16, 32, 5), nn.MaxPool2d(2, 2)) self.fc1 = nn.Linear(32 * 4 * 4, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.feature(x) x = x.view(-1, 32 * 4 * 4) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x def train(epoch): loss_runtime = 0.0 for batch, data in enumerate(tqdm(train_loader, 0)): x, y = data x = x.to(device) y = y.to(device) y_pred = model(x) loss = criterion(y_pred, y) loss_runtime += loss.item() loss_runtime /= x.size(0) optimizer.zero_grad() loss.backward() optimizer.step() print("after %s epochs, loss is %.8f" % (epoch + 1, loss_runtime)) save_file = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch} torch.save(save_file, "model_{}.pth".format(epoch)) def test(): correct, total = 0, 0 with torch.no_grad(): for (x, y) in test_loader: x = x.to(device) y = y.to(device) y_pred = model(x) _, prediction = torch.max(y_pred.data, dim=1) correct += (prediction == y).sum().item() total += y.size(0) acc = correct / total print("accuracy on test set is :%5f" % acc) if __name__ == '__main__': start_epoch = 0 freeze_epoch = 0 resume = "model_5.pth" freeze = True model = LeNet() device = ("cuda:0" if torch.cuda.is_available() else "cpu") model = model.to(device) criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5) # 加载预训练权重 if resume: checkpoint = torch.load(resume, map_location='cpu') model.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] # 冻结训练 if freeze: freeze_epoch = 5 print("冻结前置特征提取网络权重，训练后面的全连接层") for param in model.feature.parameters(): param.requires_grad = False # 将不更新的参数的requires_grad设置为False，节省了计算这部分参数梯度的时间 optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01, momentum=0.5) for epoch in range(start_epoch, start_epoch + freeze_epoch): train(epoch) test() print("解冻前置特征提取网络权重，接着训练整个网络权重") for param in model.feature.parameters(): param.requires_grad = True optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01, momentum=0.5) for epoch in range(start_epoch + freeze_epoch, 100): train(epoch) test()

参考：

1. 加载预训练权重

【本文地址】

公司简介

联系我们