Pytorch：模型的保存与加载 torch.load()、torch.nn.Module.load_state_dict()

Pytorch：模型的保存与加载 torch.save()、torch.load()、torch.nn.Module.load_state_dict()

Pytorch 保存和加载模型后缀：.pt 和.pth

1 torch.save() [source]
保存一个序列化（serialized）的目标到磁盘。函数使用了Python的pickle程序用于序列化。模型（models），张量（tensors）和文件夹（dictionaries）都是可以用这个函数保存的目标类型。

torch.save(obj, f, pickle_module=<module '...'>, pickle_protocol=2)
参数 描述
obj 保存对象
f 类文件对象 (必须实现写和刷新)或一个保存文件名的字符串
pickle_modul 用于 pickling 元数据和对象的模块
 pickle_protocol 指定 pickle protocal 可以覆盖默认参数
栗子
保存整个模型：

torch.save(model,'save.pt')
只保存训练好的权重：

torch.save(model.state_dict(), 'save.pt')
2 torch.load() [source]
用来加载模型。torch.load() 使用 Python 的 解压工具（unpickling）来反序列化 pickled object 到对应存储设备上。首先在 CPU 上对压缩对象进行反序列化并且移动到它们保存的存储设备上，如果失败了（如：由于系统中没有相应的存储设备），就会抛出一个异常。用户可以通过 register_package 进行扩展，使用自己定义的标记和反序列化方法。

torch.load(f, map_location=None, pickle_module=<module 'pickle' from '...'>)
参数 描述
f 类文件对象 (返回文件描述符)或一个保存文件名的字符串
map_location 一个函数或字典规定如何映射存储设备
pickle_module 用于 unpickling 元数据和对象的模块 (必须匹配序列化文件时的 pickle_module )
栗子
torch.load('tensors.pt')

# Load all tensors onto the CPU
torch.load('tensors.pt', map_location=torch.device('cpu'))

# Load all tensors onto the CPU, using a function
torch.load('tensors.pt', map_location=lambda storage, loc: storage)

# Load all tensors onto GPU 1
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))

# Map tensors from GPU 1 to GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

# Load tensor from io.BytesIO object
with open('tensor.pt') as f:
buffer = io.BytesIO(f.read())
torch.load(buffer)
3 torch.nn.Module.load_state_dict(state_dict) [source]
使用 state_dict 反序列化模型参数字典。用来加载模型参数。将 state_dict 中的 parameters 和 buffers 复制到此 module 及其子节点中。

torch.nn.Module.load_state_dict(state_dict, strict=True)
参数 描述
state_dict 保存 parameters 和 persistent buffers 的字典
strict 可选，bool型。state_dict 中的 key 是否和 model.state_dict() 返回的 key 一致。
栗子
torch.save(model,'save.pt')
model.load_state_dict(torch.load("save.pt")) #model.load_state_dict()函数把加载的权重复制到模型的权重中去
3.1 什么是state_dict？
在PyTorch中，一个torch.nn.Module模型中的可学习参数（比如weights和biases），模型的参数通过model.parameters()获取。而state_dict就是一个简单的Python dictionary，其功能是将每层与层的参数张量之间一一映射。注意，只有包含了可学习参数（卷积层、线性层等）的层和已注册的命令（registered buffers，比如batchnorm的running_mean）才有模型的state_dict入口。优化方法目标（torch.optim）也有state_dict，其中包含的是关于优化器状态的信息和使用到的超参数。

因为state_dict目标是Python dictionaries，所以它们可以很轻松地实现保存、更新、变化和再存储，从而给PyTorch模型和优化器增加了大量的模块化（modularity）。

torch.nn.Module.state_dict

torch.nn.Module.state_dict(destination=None, prefix='', keep_vars=False)
返回一个包含模型状态信息的字典。包含参数（weighs and biases）和持续的缓冲值（如：观测值的平均值）。只有具有可更新参数的层才会被保存在模型的 state_dict 数据结构中。

栗子：

module.state_dict().keys()

# ['bias', 'weight']
torch.optim.Optimizer.state_dict

torch.optim.Optimizer.state_dict()
返回一个包含优化器状态信息的字典。包含两个 key：

state：字典，保存当前优化器的状态信息。不同优化器内容不同。
param_groups：字典，包含所有参数组（eg：超参数）。
栗子：
from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy


# 定义模型
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

# 初始化模型
model = TheModelClass()

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

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

# 打印优化器的 state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, " ", 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]}]
4 保存/加载
4.1 state_dict（推荐）
保存：
torch.save(model.state_dict(), PATH)
加载：
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
保存模型的推理过程的时候，只需要保存模型训练好的参数，使用torch.save()保存state_dict，能够方便模型的加载。因此推荐使用这种方式进行模型保存。

记住一定要使用model.eval()来固定dropout和归一化层，否则每次推理会生成不同的结果。

注意，load_state_dict()需要传入字典对象，因此需要先反序列化state_dict再传入load_state_dict()

4.2 整个模型
保存
torch.save(model, PATH)
加载
# 模型类必须在别的地方定义
model = torch.load(PATH)
model.eval()
这种保存/加载模型的过程使用了最直观的语法，所用代码量少。这使用Python的pickle保存所有模块。这种方法的缺点是，保存模型的时候，序列化的数据被绑定到了特定的类和确切的目录。这是因为pickle不保存模型类本身，而是保存这个类的路径，并且在加载的时候会使用。因此，当在其他项目里使用或者重构的时候，加载模型的时候会出错。

一般来说，PyTorch的模型以.pt或者.pth文件格式保存。

一定要记住在评估模式的时候调用model.eval()来固定dropout和批次归一化。否则会产生不一致的推理结果。

4.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()
# - 或者 -
model.train()
在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。保存优化器的state_dict也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时epoch数，最新的模型损失，额外的torch.nn.Embedding层等。

要保存多个组件，则将它们放到一个字典中，然后使用torch.save()序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。

加载各个组件，首先初始化模型和优化器，然后使用torch.load()加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。

同样，评估模型的时候一定不要忘了调用model.eval()。

4.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()
# - 或者 -
modelA.train()
modelB.train()
保存的模型包含多个torch.nn.Modules时，比如GAN，一个序列-序列模型，或者组合模型，使用与保存常规检查点的方式来保存模型。也就是说，保存每个模型的state_dict和对应的优化器到一个字典中。我们可以保存任何能帮助我们继续训练的东西到这个字典中。

4.5 使用其他模型来预热当前模型
保存
torch.save(modelA.state_dict(), PATH)
加载
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
在迁移学习或者训练新的复杂模型时，加载部分模型是很常见的。利用经过训练的参数，即使只有少数参数可用，也将有助于预热训练过程，并且使模型更快收敛。

在加载部分模型参数进行预训练的时候，很可能会碰到键不匹配的情况（模型权重都是按键值对的形式保存并加载回来的）。因此，无论是缺少键还是多出键的情况，都可以通过在load_state_dict()函数中设定strict参数为False来忽略不匹配的键。

如果想将某一层的参数加载到其他层，但是有些键不匹配，那么修改state_dict中参数的key可以解决这个问题。

4.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=torch.device('cpu')，此时，map_location动态地将tensors的底层存储重新映射到CPU设备上。

上述代码只有在模型是在一块GPU上训练时才有效，如果模型在多个GPU上训练，那么在CPU上加载时，会得到类似如下错误：

KeyError: ‘unexpected key “module.conv1.weight” in state_dict’
原因是在使用多GPU训练并保存模型时，模型的参数名都带上了module前缀，因此可以在加载模型时，把key中的这个前缀去掉：

# 原始通过DataParallel保存的文件
state_dict = torch.load('myfile.pth.tar')
# 创建一个不包含`module.`的新OrderedDict
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # 去掉 `module.`
new_state_dict[name] = v
# 加载参数
model.load_state_dict(new_state_dict)
GPU上保存，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")) # 选择希望使用的GPU
model.to(device)
在把GPU上训练的模型加载到GPU上时，只需要使用model.to(torch.devie('cuda'))将初始化的模型转换为CUDA优化模型。同时确保在模型所有的输入上使用.to(torch.device('cuda'))。注意，调用my_tensor.to(device)会返回一份在GPU上的my_tensor的拷贝。不会覆盖原本的my_tensor，因此要记得手动将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 训练保存的模型时，将 torch.load() 函数的 map_location 参数 设置为 cuda:device_id。这种方式将模型加载到指定设备。下一步，确保调用 model.to(torch.device(‘cuda’)) 将模型参数 tensor 转换为 cuda tensor。最后，确保模型输入使用 .to(torch.device(‘cuda’)) 为 cuda 优化模型准备数据。

注意：调用 my_tensor.to(device) 会在 GPU 上返回 my_tensor 的新副本，不会覆盖 my_tensor。因此，使用 my_tensor = my_tensor.to(torch.device(‘cuda’)) 手动覆盖。

4.7 保存torch.nn.DataParallel模型
保存
torch.save(model.state_dict(), PATH)
加载

# Load to whatever device you want
torch.nn.DataParallel 是支持模型使用 GPU 并行的封装器。要保存一个一般的 DataParallel 模型， 请保存 model.module.state_dict()。这种方式，可以灵活地以任何方式加载模型到任何设备上。

5 参考
https://blog.csdn.net/LXYTSOS/article/details/90639524
https://blog.csdn.net/cpongo3/article/details/93624288
————————————————
版权声明：本文为CSDN博主「宁静致远*」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/weixin_40522801/article/details/106563354

-----------------------------------------------------------------

PyTorch模型保存深入理解

前面写过一篇PyTorch保存模型的文章：Pytorch模型保存与加载，并在加载的模型基础上继续训练 ，简单介绍了PyTorch模型保存和加载函数的用法，足以快速上手，但对相关函数和参数的具体用法和代表的含义没有进行展开介绍，这篇文章用于记录之。

PyTorch保存模型的语句是这样的：

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

加载是这样的：

model.load_state_dict(torch.load(path))

下面我们将其拆开逐句介绍，深入理解。

1.torch.save()和torch.load()

顾名思义，save函数是PyTorch的存储函数，load函数则是读取函数。save函数可以将各种对象保存至磁盘，包括张量，列表，ndarray，字典，模型等；而相应地，load函数将保存在磁盘上的对象读取出来。

用法：

torch.save(保存对象, 保存路径)
torch.load(文件路径)

应用举例：

保存张量

In [3]: a = torch.ones(3)                                                       

In [4]: a                                                                       
Out[4]: tensor([1., 1., 1.])

In [5]: torch.save(a, './a.pth')          # 保存Tensor               

In [6]: a_load = torch.load('./a.pth')    # 读取Tensor

In [7]: a_load                                                                  
Out[7]: tensor([1., 1., 1.])

保存字典

In [11]: b = {k:v for v,k in enumerate('abc',1)}                                

In [12]: b                                                                      
Out[12]: {'a': 1, 'b': 2, 'c': 3}

In [13]: torch.save(b, './b.rar')                        

In [14]: torch.load('./b.rar')                           
Out[14]: {'a': 1, 'b': 2, 'c': 3}

可以看出，保存和读取非常方便。这里需要注意的是文件的命名，命名必须要有扩展名，扩展名可以为‘xxx.pt’，‘xxx.pth’，‘xxx.pkl’，‘xxx.rar’等形式。

2.model.state_dict()

在PyTorch中，state_dict是一个从参数名称隐射到参数Tesnor的字典对象。

In [15]: class MLP(nn.Module): 
    ...:     def __init__(self): 
    ...:         super(MLP, self).__init__() 
    ...:         self.hidden = nn.Linear(3, 2) 
    ...:         self.act = nn.ReLU() 
    ...:         self.output = nn.Linear(2, 1) 
    ...:  
    ...:     def forward(self, x): 
    ...:         a = self.act(self.hidden(x)) 
    ...:         return self.output(a) 
    ...:                                                                        

In [16]: net = MLP()                                                            

In [17]: net.state_dict()                                                       
Out[17]: 
OrderedDict([('hidden.weight', tensor([[ 0.4839,  0.0254,  0.5642],
                      [-0.5596,  0.2602, -0.5235]])),
             ('hidden.bias', tensor([-0.4986, -0.5426])),
             ('output.weight', tensor([[0.0967, 0.4980]])),
             ('output.bias', tensor([-0.4520]))])

可以看出，state_dict()返回的是一个有序字典，该字典的键即为模型定义中有可学习参数的层的名称+weight或+bias，值则对应相应的权重或偏差，无参数的层则不在其中。

state_dict()返回模型可学习参数的键值对，那么问题来了：model.parameters()不也是模型可学习参数的访问方式吗？它们有什么不同？

我们再来对比了一下常用的model.parameters()，在训练时，我们常常将此语句放在优化器中，表示要优化学习的模型参数。关于model.parameters(), model.named_parameters(), model.children(), model.named_children(), model.modules(), model.named_modules()请看我的这篇博客

In [18]: net.parameters()                                                       
Out[18]: <generator object Module.parameters at 0x7f009b78acf0>

In [19]: list(net.parameters())                                                 
Out[19]: 
[Parameter containing:
 tensor([[ 0.4839,  0.0254,  0.5642],
         [-0.5596,  0.2602, -0.5235]], requires_grad=True),
 Parameter containing:
 tensor([-0.4986, -0.5426], requires_grad=True),
 Parameter containing:
 tensor([[0.0967, 0.4980]], requires_grad=True),
 Parameter containing:
 tensor([-0.4520], requires_grad=True)]
#######################################################################
In [20]: net.named_parameters()                                                 
Out[20]: <generator object Module.named_parameters at 0x7f00a009a390>

In [21]: list(net.named_parameters())                                           
Out[21]: 
[('hidden.weight', Parameter containing:
  tensor([[ 0.4839,  0.0254,  0.5642],
          [-0.5596,  0.2602, -0.5235]], requires_grad=True)),
 ('hidden.bias', Parameter containing:
  tensor([-0.4986, -0.5426], requires_grad=True)),
 ('output.weight', Parameter containing:
  tensor([[0.0967, 0.4980]], requires_grad=True)),
 ('output.bias', Parameter containing:
  tensor([-0.4520], requires_grad=True))]

可以看出，model.parameters()是一个生成器，每个参数张量都是一个参数容器，它的对象是各个参数Tensor。这里简单的提一下model.named_parameters()，它和model.parameters很像，都返回一个可迭代对象，但从上例可以看出它多返回一个参数名称，这样有利于访问和初始化或修改参数。
我们在用优化器优化参数时，优化对象是纯参数，所以用model.parameters():

In [22]: optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  

In [23]: optimizer.state_dict()                                                 
Out[23]: 
{'state': {},
 'param_groups': [{'lr': 0.001,
   'momentum': 0.9,
   'dampening': 0,
   'weight_decay': 0,
   'nesterov': False,
   'params': [139638712861752,
    139640593827520,
    139640585217080,
    139640585217296]}]}

除了模型中可学习参数的层(卷积层、线性层等)有state_dict，优化器也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。

3.model.load_state_dict()

这是模型加载state_dict的语句，也就是说，它的输入是一个state_dict，也就是一个字典。模型定义好并且实例化后会自动进行初始化，上面的例子中我们定义的模型MLP在实例化以后显示的模型参数都是自动初始化后的随机数。
在训练模型或者迁移学习中我们会使用已经训练好的参数来加速训练过程，这时候就用load_state_dict()语句加载训练好的参数并将其覆盖在初始化参数上，也就是说执行过此语句后，加载的参数将代替原有的模型参数。
既然加载的是一个字典，那么需要注意的就是字典的键一定要相同才能进行覆盖，比如加载的字典中的'hidden.weight'只能覆盖当前模型的'hidden.weight'，如果键不同，则不能实现有效覆盖操作。键相同而值的shape不同，则会将新的键值对覆盖原来的键值对，这样在训练时会报错。所以我们在加载前一般会进行数据筛选，筛选是对字典的键进行对比来操作的：

pretrained_dict = torch.load(log_dir)  # 加载参数字典
model_state_dict = model.state_dict()  # 加载模型当前状态字典
pretrained_dict_1 = {k:v for k,v in pretrained_dict.items() if k in model_state_dict}  # 过滤出模型当前状态字典中没有的键值对
model_state_dict.update(pretrained_dict_1)  # 用筛选出的参数键值对更新model_state_dict变量
model.load_state_dict(model_state_dict)  # 将筛选出的参数键值对加载到模型当前状态字典中

以上代码简单的对预训练参数进行了过滤和筛选，主要是通过第3条语句粗略的过滤了键值对信息，进行筛选后要用Python更新字典的方法update()来对模型当前字典进行更新，update()方法将pretrained_dict_1中的键值对添加到model_state_dict中，若pretrained_dict_1中的键名和model_state_dict中的键名相同，则覆盖之；若不同，则作为新增键值对添加到model_state_dict中。显然，这里需要的是将pretrained_dict_1中的键值对覆盖model_state_dict的相应键值对，所以对应的键的名称必须相同，所以第3条语句中按键名称进行筛选，过滤出当前模型字典中没有的键值对。否则会报错。
如果想要细粒度过滤或更改某些参数的维度，如进行卷积核参数维度的调整，假如预训练参数里conv1有256个卷积核，而当前模型只需要200个卷积核，那么可以采用类似以下语句直接对字典进行更改：

pretrained_dict['conv1.weight'] = pretrained_dict['conv1.weight'][:200,:,:,:]   # 假设保留前200个卷积核

以上。

原文链接：https://www.jianshu.com/p/6c558300130f