memo: PyTorch | Save Model

保存模型与加载

使用的三个核心功能：

torch.save 将序列化对象保存到磁盘。此函数调用pickle模块把（模型、tensor、字典等）对象序列化。
torch.load 函数调用pickle的unpickling功能把文件反序列化到内存。
torch.nn.Module.load_state_dict 调用反序列化函数state_dict加载模型的参数字典。

1. 状态字典 state_dict

模型的state_dict包括：各层网络的可学习参数（权重和偏置），优化器的state_dict包括优化器的状态和超参数。
模型torch.nn.Module的可学习参数，用model.parameters()访问

2. 两种方式：

只保存网络中的参数（速度快，占空间少，推荐）

1
2
3
4
5
6
7


#保存参数
torch.save(net1.state_dict(), 'net_params.pt')   #或'.pth',或'pkl'

#加载
model = ModelClass(*args, **kwargs)  #先实例化一个模型对象
model = model.load_state_dict(torch.load(PATH))  #把文件反序列化成字典对象，把参数传给模型
model.eval() #设置dropuout 和 batch normalization层为评估模式，否则可能导致模型推断结果不一致。

保存整个网络的结构和参数与加载：

1
2
3
4
5
6


#保存
torch.save(net1, 'net.pkl')

#加载
newmodel = torch.load(PATH) #不需重构模型，直接load
newmodel.eval()

3. 保存和加载Checkpoint 用于推理/继续训练

保存训练状态：

1
2
3
4
5
6


torch.save({'epoch': epoch+1,     #保存当前的迭代次数
            'model_state_dict': model.state_dict(), #保存模型参数
            'optimizer_state_dict': optimizer.state_dict(), #保存优化器参数
            'loss': loss, #其余一些想保持的参数都可以添加进来
            ...,
            }, PATH)  #后缀可以用 '.pth.tar'或 '.pth'

加载：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12


model = ModelClass(*args, **kwargs)
optimizer = OptimizerClass(*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()

。。。。。部分内容需公众号验证

4. 在一个文件中保存多个模型

 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


# 保存
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()

当保存一个模型由多个torch.nn.Modules组成时，例如GAN(对抗生成网络)、sequence-to-sequence (序列到序列模型), 或者是多个模型融合, 可以采用与保存常规检查点相同的方法。换句话说，保存每个模型的 state_dict 的字典和相对应的优化器。如前所述，可以通过简单地将它们附加到字典的方式来保存任何其他项目，这样有助于恢复训练。

PyTorch 中常见的保存 checkpoint 是使用 .tar 文件扩展名。

要加载项目，首先需要初始化模型和优化器，然后使用torch.load()来加载本地字典。这里，你可以非常容易的通过简单查询字典来访问你所保存的项目。

请记住在运行推理之前，务必调用model.eval()去设置 dropout 和 batch normalization 为评估。如果不这样做，有可能得到不一致的推断结果。如果你想要恢复训练，请调用model.train()以确保这些层处于训练模式。

5. 使用在不同模型参数下的热启动模式

1
2
3
4
5


保存
torch.save(modelA.state_dict(), PATH)
加载
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

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

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

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

6.通过设备保存/加载模型

6.1 保存到 CPU、加载到 CPU

1
2
3
4
5
6


保存
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设备。

6.2 保存到 GPU、加载到 GPU

1
2
3
4
5
6
7
8


保存
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’))。

6.3 保存到 CPU，加载到 GPU

1
2
3
4
5
6
7
8


保存
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’))。

6.4 保存 torch.nn.DataParallel 模型

1
2
3
4


保存
torch.save(model.module.state_dict(), PATH)
加载
# 加载任何你想要的设备

torch.nn.DataParallel是一个模型封装，支持并行GPU使用。要普通保存 DataParallel 模型, 请保存model.module.state_dict()。这样，你就可以非常灵活地以任何方式加载模型到你想要的设备中。

Table of contents