Pytorch 模型的保存与加载

Pytorch 模型的保存与加载

想先在一部分数据跑出来几个模型，然后选择最优的加载 在另一部分数据上继续跑，加载时遇到了问题，from_pretrained()显示state_dict中的都参数没加载成功，是因为都多了一层encoder: encoder.encoder.self.attention…所以查找了pytorch保存与加载模型的相关知识，在这里记录一下~

来源：

https://pytorch.org/tutorials/beginner/saving_loading_models.html#

首先需要熟悉三个核心函数：

1. torch.save: 保存序列化对象到磁盘中。该函数使用pickle做序列化，可以保存模型，张量和字典等对象
2. torch.load: 使用pickle做反序列化到内存中。
3. torch.nn.Module.load_state_dict:加载模型的参数字典

什么是state_dict？

pytorch里，torch.nn.Module可学习的参数可以通过model.parameters() 得到，而state_dict 就是将这些参数映射到一个python字典里,方便存储和更新。除了卷积层、线性层等参数外，Optimizer(torch.optim)也有state_dict, 包含Optimizer的信息和一些超参数。

state_dict的例子：

# Define model
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.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

# Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
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
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])
    
'''    
Output:

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]}]
'''

保存并加载模型

1. 保存、加载state_dict(推荐)

保存：

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

加载：

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

保存state_dict 会更灵活，方便恢复模型。

模型保存后的扩展名为.pt或.pth

记住在用模型推理之前要先调用model.eval()设置dropout和batch normalization层，否则会产生不连续的结果

（这里没调用好像也没产生啥后果）

注意：

这里有个很容易出错的地方，经常会误写成model.load_state_dict(PATH)，错在哪里呢？

load_state_dict()参数必须是个字典对象，而不是保存字典的对象。所以先反序列化得到字典对象state_dict= torch.load(PATH), 然后再load。

2. 保存、加载整个模型

保存：

torch.save(model,PATH)

加载

model = torch.load(PATH)
model.eval()

这种方法最符合我们的直觉，用pickle模块保存整个模型，保存了模型涉及到的具体类的结构，所以当重构模型时可能会出错。

扩展名同上,为 .pt或.pth

3. 保存、加载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()
# - or -
model.train()

当保存checkpoint用于推理或重训练的时候，可能需要epoch/loss等信息，所以需要保存成字典形式再save。

扩展名一般是.tar

4. 一个文件里保存多个模型

保存：

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'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

其实和上个方法相似，不同的模型保存在同一个字典，然后再保存到文件里。

扩展名同样为.tar

5. 只使用模型的一部分参数

保存：

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

加载：

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

当做迁移学习或加载部分参数来warm_start模型时，参数会不完全匹配，出现 missing keys或 unexpected keys，这时把load_state_dict的strict设为False即可。

如果出现需要加载的模型参数名和另一个模型不匹配，可以更改 state_dict的keys，文章开始提到的那个问题，每个参数前面多加了encoder.，所以把每个参数的key去掉了’ encoder.’ 这部分，使用

# 这里加list是因为，orderDict更改key会报错：
# OrderedDict mutated during iteration
# 变成list就可以了
for param_key in list(state_dict.keys())
    # 我这里是在最开始多了'encoder.
	if param_key.find('encoder.') == 0:
        tmp_key = param_key[8:]
        state_dict[tmp_key] = state_dict.pop(param_key)

6. 跨机器保存、加载模型

在GPU上保存，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.load()需要增加一个map_location参数，重映射到cpu

在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)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

在GPU上加载一个GPU上训练得到的模型时，需要用model.to(torch.device('cuda'))。

另外，模型的输入也需要使用.to(torch.device('cuda')), 注意，my_tensor.to(device)会返回my_tensor的一个拷贝，所以需要进行覆盖：

my_tensor = my_tensor.to(torch.device('cuda'))

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)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

在GPU上加载CPU保存的模型时，需要设置map_location参数为 cuda:device_id，然后 model.to(torch.device('cuda'))，最后注意所有的输入也需要.to(device)

保存torch.nn.DataParallel模型

保存:

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

加载：

# Load to whatever device you want

torch.nn.DataParallel设置GPU并行使用，使用model.module.state_dict()保存模型，然后加载的灵活性很强，可以加载到任意device上