其实之前笔者写代码的时候用到模型的保存和加载，需要用的时候就去度娘搜一下大致代码，现在有时间就来整理下整个pytorch模型的保存和加载，开始学习把~
pytorch的模型和参数是分开的，可以分别保存或加载模型和参数。所以pytorch的保存和加载对应存在两种方式：
1. 直接保存加载模型
（1）保存和加载整个模型
|
1
2
3
4
|
# 保存模型torch.save(model, 'model.pth\pkl\pt') #一般形式torch.save(net, PATH) # 加载模型model = torch.load('model.pth\pkl\pt') #一般形式为model_dict=torch.load(PATH) |
（2）仅保存和加载模型参数(推荐使用,需要提前手动构建模型)
速度快，占空间少
|
1
2
3
4
|
# 保存模型参数torch.save(model.state_dict(), 'model.pth\pkl\pt') #一般形式为torch.save(net.state_dict(),PATH)# 加载模型参数model.load_state_dict(torch.load('model.pth\pkl\pt') #一般形式为model_dict=model.load_state_dict(torch.load(PATH)) |
state_dict（） 是一个Python字典，将每一层映射成它的参数张量。注意只有带有可学习参数的层（卷积层、全连接层等），以及注册的缓存（batchnorm的运行平均值）在state_dict 中才有记录。state_dict同样包含优化器对象，存储了优化器的状态，所使用到的超参数。
然而，在实验中往往需要保存更多的信息，比如优化器的参数，那么可以采取下面的方法保存：
|
1
2
3
|
torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN, 'optimizer': optimizer.state_dict(),'alpha': loss.alpha, 'gamma': loss.gamma}, checkpoint_path + '/m-' + launchTimestamp + '-' + str("%.4f" % lossMIN) + '.pth.tar') |
如下一个完整的使用model.state_dict()和optimizer.state_dict()例子：
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
|
# 定义模型class TheModelClass(nn.Module): #定义一个神经网络模型 TheModelClass def __init__(self): super(TheModelClass, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = x.view(-1, 16 * 5 * 5) x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x# 初始化模型model = TheModelClass()# 初始化优化器optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)# 打印模型的 state_dictprint("Model's state_dict:")for param_tensor in model.state_dict(): # param_tensor 为参数名称 print(param_tensor, "\t", model.state_dict()[param_tensor].size())# 打印优化器的 state_dictprint("Optimizer's state_dict:")for var_name in optimizer.state_dict(): print(var_name, "\t", optimizer.state_dict()[var_name]) |
输出结果：
Model's state_dict:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]
（3）load提供了很多重载的功能，其可以把在GPU上训练的权重加载到CPU上跑
|
1
2
3
4
5
6
7
|
torch.load('tensors.pt')# 强制所有GPU张量加载到CPU中torch.load('tensors.pt', map_location=lambda storage, loc: storage) #或者model.load_state_dict(torch.load('model.pth', map_location='cpu'))# 把所有的张量加载到GPU 1中torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))# 把张量从GPU 1 移动到 GPU 0torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'}) |
上述代码只有在模型在一个GPU上训练时才起作用。如果我在多个GPU上训练模型并且保存它，然后尝试在CPU上加载，会得到错误：KeyError: ‘unexpected key “module.conv1.weight” in state_dict' 如何解决？
因为此时已经使用模型保存了模型nn.DataParallel，该模型将模型存储在该模型中module，而现在您正试图加载模型DataParallel。您可以nn.DataParallel在网络中暂时添加一个加载目的，也可以加载权重文件，创建一个没有module前缀的新的有序字典，然后加载它。看第（4）点
（4）通过DataParalle使用多GPU时的保存和加载
|
1
2
3
|
odel=DataParalle(model)#保存参数torch.save(model.module.state_dict(), 'model.pth') |
由此看出多个GPU时多了一个该模型中module，加载再cpu时，创建一个没有module前缀的新的有序字典，然后加载它。
补充：一般来说，PyTorch的模型以.pt或者.pth文件格式保存。
2. 保存加载用于推理的常规Checkpoint/或继续训练**
checkpoint检查点：不仅保存模型的参数，优化器参数，还有loss，epoch等（相当于一个保存模型的文件夹）
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
|
if (epoch+1) % checkpoint_interval == 0: checkpoint = {"model_state_dict": net.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "epoch": epoch} path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch) torch.save(checkpoint, path_checkpoint)#或者#保存torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, ... }, PATH)#加载model = TheModelClass(*args, **kwargs)optimizer = TheOptimizerClass(*args, **kwargs)checkpoint = torch.load(PATH)model.load_state_dict(checkpoint['model_state_dict'])optimizer.load_state_dict(checkpoint['optimizer_state_dict'])epoch = checkpoint['epoch']loss = checkpoint['loss']model.eval()# - 或者 -model.train() |
注意：
在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。保存优化器的state_dict也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时epoch数，最新的模型损失，额外的torch.nn.Embedding层等。
要保存多个组件，则将它们放到一个字典中，然后使用torch.save()序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。
加载各个组件，首先初始化模型和优化器，然后使用torch.load()加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。
同样，评估模型的时候一定不要忘了调用model.eval()。
是不是很简单！！以上为个人经验，希望能给大家一个参考，也希望大家多多支持服务器之家。
原文链接：https://blog.csdn.net/qq_38765642/article/details/109784913
