本篇文档提供了有关PyTorch模型保存和加载的各种用例的解决方案。随意地阅读整篇文档,或跳到你需要的用例的代码上。
在保存和加载模型时,有三个核心功能需要熟悉:
- torch.save:将序列化对象保存到磁盘。该函数使用Python的pickle模块中的utility进行序列化。可以使用该模块保存各种对象的模型、张量,以及词典。
- torch.load:使用pickle的unpickling工具将pickle了的对象反序列化到内存中。该函数还能够帮助数据跨设备传递。
- torch.nn.Module.load_state_dict:使用反序列的state_dict来加载模型的参数词典。更多细节参阅“state_dict是什么?”节。
state_dict是什么?
在PyTorch中,torch.nn.Module中可学习的参数(如权重和偏置)被包含在模型的parameters中(通过model.parameters访问)。state_dict是一个简单的将每个层映射到其参数张量上的Python词典对象。注意只有拥有可学习参数的层(卷积层,线性层等)以及已注册的缓冲区(batchnorm的running_mean)的层在模型的state_dict上有条目。optimizer对象(torch.optim)也有state_dict,它包含了optimizer的状态信息,以及使用的超参数。
因为state_dict对象为Python词典,它们可以被简单地保存、更新、修改和恢复,为PyTorch模型和optimizer添加了大量的模块化。
Example:
让我们从Training a classifier教程中使用的一个简单的模型来看看state_dict的作用。
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)
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])
其结果为:
Saving&Loading Model for Inference
Save/Load state_dict(推荐)
Save:
torch.save(model.state_dict(),PATH)
Load:
model = TheModelClass(*args,**kwargs)
model.load_state_dict(torch.load(PATH))
## 将模型设置为evaluation模式,有些函数在training/evaluation模式下的表现不同
## 见 https://blog.csdn.net/Strive_For_Future/article/details/83240553
model.eval()
当保存用于inference的模型时,只需要保存训练模型的学习参数即可。使用torch.save()函数保存模型的state_dict能够给你最大的在之后恢复模型的灵活性,这就是为什么它是保存模型的推荐方法。
常见的PyTorch约定为使用.pt或.pth文件扩展名来保存模型。
记住你必须在运行inference前,调用eval()来将dropout和批量归一化层设置为evaluation模式。如果不这样做,将会导致不一致的inference结果。
Note:
注意load_state_dict()函数将词典对象作为输入,而不是保存的对象的路径。这意味着你必须在将保存的state_dict传递给load_state_dict()之前,必须将其反序列化。比如说,你不能使用model.load_state_dict(PATH)进行载入。
Save/Load Entire Model
Save:
torch.save(model,PATH)
Load:
model = TheModelClass(*args,**kwargs)
model = torch.load(PATH)
该save/load过程使用最直观的语法,涉及了最少的代码。使用这种方式保存模型将会使用Python的pickle模块保存整个模块。该方法的缺点在于,当模型被保存后,序列化的数据将会被绑定到特定的类和确切的目录结构中。这一结果的原因是pickle并不保存模型类本身。它会保存包含类的文件的路径,该文件将会在加载时被使用。由于这个原因,你的代码可能会在其它项目中或在重构之后被使用时,因为各种方式中断。
常见的PyTorch约定为使用.pt或.pth文件扩展名来保存模型。
记住你必须在运行inference前,调用eval()来将dropout和批量归一化层设置为evaluation模式。如果不这样做,将会导致不一致的inference结果。
Saving&Loading a General Checkpoint for Inference and/or Resuming Training
Save:
torch.save({
'epoch':epoch,
'model_state_dict':model.state_dict(),
'optimizer':optimizer.state_dict(),
'loss':loss,
...
},PATH)
Load:
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()
当保存通用的用于推理或恢复训练的检查点时,你要保存的不仅仅是模型的state_dict。保存optimizer的state_dict也很重要,因为它包含了模型训练时更新的缓冲区和参数。其它你可能想保存的项有你剔除的epoch,最新记录的损失,外部的torch.nn.Embedding层等。
为了保存多个组件,请在词典中组织它们,并使用torch.save()来序列化该词典。常见的PyTorch约定为使用.tar文件扩展名来保存这些检查点。
想加载这些项,首先要做的是初始化模型和optimizer,然后使用torch.load()从本地载入词典。使用这种方式,你可以简单地通过查询词典来访问你想要的,且已保存的项。
记住你必须在运行inference前,调用eval()来将dropout和批量归一化层设置为evaluation模式。如果不这样做,将会导致不一致的inference结果。如果你想要恢复训练,则调用model.train()来保证这些层处于训练模式。
Saving Multiple Models in One File
Save:
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)
Load:
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.Module(如GAN,seq2seq)组成的模型,或多个模型的组合时,你可以使用和保存常规检查点(general checkpoint)时同样的方法。换句话说,保存由每个模型的state_dict以及对应的optimizer构成的词典。
常见的PyTorch约定为使用.tar文件扩展名来保存这些检查点。
想加载这些项,首先要做的是初始化模型和optimizer,然后使用torch.load()从本地载入词典。使用这种方式,你可以简单地通过查询词典来访问你想要的,且已保存的项。
记住你必须在运行inference前,调用eval()来将dropout和批量归一化层设置为evaluation模式。如果不这样做,将会导致不一致的inference结果。如果你想要恢复训练,则调用model.train()来保证这些层处于训练模式。
Warmstarting Model Using Parameters from a Different Model
Save:
torch.save(modelA.state_dict(),PATH)
Load
modelB = TheModelBClass(*args,**kwargs)
modelB.load_state_dict(torch.load(PATH),strict=False)
部分地载入模型,或载入部分的模型,是迁移学习或训练新的复杂模型时常见的情况。利用训练好的参数,即使其中只有一小部分可用,也有助于热启动训练过程,且很可能能够让你的模型比从头开始训练更快地收敛。
不论你是从迷失了一些关键字的部分的state_dict中,还是从比当前的模型拥有更多键的state_dict中加载,你可以把load_state_dict()中的strict参数设为False来忽略那些不匹配的键。
如果你想要把某一层的参数载入另一层,但是其中的一些键并不匹配,你只需要简单地把前者的state_dict中参数的键的名字修改得与后者的模型中的键相匹配即可。
Saving&Loading Model Across Devices
Save on GPU,load on CPU
Save:
torch.save(model.state_dict(),PATH)
Load:
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设备中。
Save on GPU,load on GPU
Save:
torch.save(model.state_dict(),PATH)
Load:
device = torch.device("cuda")
model = TheModelClass(*args,**kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
## 确保你同时将该模型上的所有张量也移动到device上。
当在GPU上载入已经在GPU上训练和保存的模型时,只需要使用model.to(torch.device(‘cuda’))将已经初始化的模型转换为CUDA优化的模型。同时,请务必在模型的所有输入上使用.to(torch.device(‘cuda’))函数来为模型准备数据。注意调用my_tensor.to(device)后会返回一个my_tensor在GPU上的新的复制。它不会覆盖my_tensor。因此,请记得手动覆盖张量my_tensors = my_tensor.to(torch.device(‘cuda’))。
Save on CPU,load on GPU
Save:
torch.save(model.state_dict(),PATH)
Load:
device = torch.device('cuda')
model = TheModelClass(*args,**kwargs)
model.load_state_dict(torch.load(PATH,map_location='cuda:0'))
model.to(device)
当在GPU上载入在CPU上训练和保存的模型时,将torch.load()函数上的map_location参数设为’cuda:device_id’。这将模型载入指定的GPU。接下来,确保调用 model .to(torch.device(‘cuda’))来将模型的参数张量转换为CUDA张量。最后,务必在模型的所有输入上使用 .to(torch.device(‘cuda’))函数来为被CUDA优化的模型准备数据。注意,调用my_tensor.to(device)会返回一个my_tensor在GPU上的新的复制。它不会覆盖my_tensor。因此,请记得手动覆盖张量my_tensors = my_tensor.to(torch.device(‘cuda’))。
Save torch.nn.DataParallel Models
Save:
torch.save(model.module.state_dict(),PATH)
Load:
## 随便找个你喜欢的设备载入即可
torch.nn.DataParallel是一个支持并行利用GPU的模型包装器。要一般性地保存DataParallel模型,请保存model.Module.state_dict()。通过这种方式,你拥有使用任何方式往任何设备上载入模型的灵活性。