Pytorch存储权重以及如何加载

Pytorch存储权重以及如何加载

关于Pytorch如何使用，必然是官方教程写的好。

一、Pytorch如何保存权重

1.torch.save()

Pytorch保存权重使用的是torch.save()函数。该函数将序列化对象保存到磁盘。此函数使用Python的pickle模块进行序列化。使用此函数可以保存如模型、tensor、字典等各种对象。

常见的保存模型命令如下：

torch.save(model.state_dict(), PATH) # 保存state_dict()
torch.save(model, PATH) # 保存完整模型
torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH) # 保存Checkpoint，有了checkpoint再也不怕程序突然中途断了要重新跑了
torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH) # 一个文件保存多个模型

在 PyTorch 中最常见的模型保存使用‘.pt’或者是‘.pth’作为模型文件扩展名。

可以看的后两个输入参数都是字典，四种方式中都涉及到了state_dict()，那么state_dict()是啥东西呢？要搞懂Pytorch如何存储和加载模型这东西就必须搞懂。

2.state_dict()

在PyTorch中，torch.nn.Module模型的可学习参数（即权重weight和偏差bias，这个bias我们在设计网络时往往设置为默认参数，容易忽视其存在）包含在模型的参数中，（使用model.parameters()可以进行访问）。 state_dict是Python字典对象，它将每一层映射到其参数张量。注意，只有具有可学习参数的层（如卷积层，线性层等）的模型 才具有state_dict这一项。目标优化torch.optim也有state_dict属性，它包含有关优化器的状态信息，以及使用的超参数。

因为state_dict的对象是Python字典，所以它们可以很容易的保存、更新、修改和恢复，为PyTorch模型和优化器添加了大量模块。

下面通过从简单模型训练一个分类器中来了解一下state_dict的使用。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
# 定义模型
class TheModelClass(nn.Module):
    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.fc = nn.Sequential(nn.Linear(16 * 5 * 5, 120),
                                nn.ReLU(),
                                nn.Linear(120, 84),
                                nn.ReLU(),
                                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 = self.fc(x)
        return x

# 初始化模型
model = TheModelClass()

# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 打印模型的状态字典
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# 打印优化器的状态字典
print("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])
fc.0.weight 	 torch.Size([120, 400])
fc.0.bias 	 torch.Size([120])
fc.2.weight 	 torch.Size([84, 120])
fc.2.bias 	 torch.Size([84])
fc.4.weight 	 torch.Size([10, 84])
fc.4.bias 	 torch.Size([10])
Optimizer's state_dict:
state 	 {}
param_groups 	 [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'maximize': False, 'foreach': None, 'differentiable': False, 'params': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]}]

可以看到对于模型state_dict()存储的东西仅于我们设计的模型的__init__有关，和forward()函数无关。字典的键是我们自己定义的，其中利用nn.Sequential()产生的添加了后缀数字如fc.2.weight。知道了这一点，及其方便我们加载部分模型结构的权重。

二、Pytorch如何加载权重

1.torch.load()和model.load_state_dict()

常见用法如下：

# 保存/加载state_dict（推荐使用）
torch.save(model.state_dict(), PATH)
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))

load_state_dict()函数只接受字典对象，而不是保存对象的路径。这就意味着在你传给load_state_dict()函数之前，你必须反序列化你保存的state_dict，即使用torch.load()函数反序列化权重。例如，你无法通过 model.load_state_dict(PATH)来加载模型。

# 保存/加载完整模型
torch.save(model, PATH)
# 模型类必须在此之前被定义
model = torch.load(PATH)
model.eval()

# 保存和加载 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']

# 在一个文件中保存多个模型
torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)
modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

2.仅加载部分模型权重

在加载模型权重时，我们往往仅需要获取加载模型权重的部分网络。这个时候该咋办？了解了state_dict()后我们就知道该咋办了，就按照字典处理就行了。

checkpoint_path = PATH
checkpoint = torch.load(checkpoint_path)

model_dict = model.state_dict()
state_dict = {k:v for k,v in checkpoint_matching.items() if k in model_dict.keys()}
model_dict.update(state_dict)
model.load_state_dict(model_dict)

或者可以使用如下：

model.load_state_dict(torch.load(PATH), strict=False)

在迁移学习或训练新的复杂模型时，部分加载模型或加载部分模型是常见的情况。利用训练好的参数，有助于热启动训练过程，并希望帮助你的模型比从头开始训练能够更快地收敛。

无论是从缺少某些键的 state_dict 加载还是从键的数目多于加载模型的 state_dict , 都可以通过在load_state_dict()函数中将strict参数设置为 False 来忽略非匹配键的函数。

如果要将参数从一个层加载到另一个层，但是某些键不匹配，主要修改正在加载的 state_dict 中的参数键的名称以匹配要在加载到模型中的键即可。

3.torch.load(PATH, map_loaction=device)

主要是当

3.1保存到 GPU、加载到 GPU

• 保存

torch.save(model.state_dict(), PATH)

• 加载

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
# 确保在你提供给模型的任何输入张量上调用input = input.to(device)

当在GPU上训练并把模型保存在GPU，只需要使用model.to(torch.device('cuda'))，将初始化的 model 转换为 CUDA 优化模型。另外，请 务必在所有模型输入上使用.to(torch.device('cuda'))函数来为模型准备数据。请注意，调用my_tensor.to(device)会在GPU上返回my_tensor的副本。 因此，请记住手动覆盖张量：my_tensor= my_tensor.to(torch.device('cuda'))。

3.2 保存到 CPU，加载到 GPU

• 保存

torch.save(model.state_dict(), PATH)

• 加载

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)
# 确保在你提供给模型的任何输入张量上调用input = input.to(device)

在CPU上训练好并保存的模型加载到GPU时，将torch.load()函数中的map_location参数设置为cuda:device_id。这会将模型加载到 指定的GPU设备。接下来，请务必调用model.to(torch.device('cuda'))将模型的参数张量转换为 CUDA 张量。最后，确保在所有模型输入上使用 .to(torch.device('cuda'))函数来为CUDA优化模型。请注意，调用my_tensor.to(device)会在GPU上返回my_tensor的新副本。它不会覆盖my_tensor。 因此， 请手动覆盖张量my_tensor = my_tensor.to(torch.device('cuda'))。

3.3 保存到 CPU、加载到 CPU

• 保存

torch.save(model.state_dict(), PATH)

• 加载

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

当从CPU上加载模型在GPU上训练时, 将torch.device('cpu')传递给torch.load()函数中的map_location参数.在这种情况下，使用 map_location参数将张量下的存储器动态的重新映射到CPU设备。

三、多卡上的保存和使用

关于torch.nn.DataParallel模型的保存，因为涉及到多卡和单卡使用，和上面介绍还是有点区别的。接下来有时间就把这部分更新一下

torch.save(model.module.state_dict(), PATH)