saving and loading models

最新推荐文章于 2023-03-13 17:30:05 发布

JachinMa

最新推荐文章于 2023-03-13 17:30:05 发布

阅读量298

点赞数

本篇文档提供了有关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文件扩展名来保存这些检查点。

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()。通过这种方式，你拥有使用任何方式往任何设备上载入模型的灵活性。

JachinMa

关注

0
点赞
踩
0

收藏

觉得还不错? 一键收藏
0
评论
saving and loading models

本篇文档提供了有关PyTorch模型保存和加载的各种用例的解决方案。随意地阅读整篇文档，或跳到你需要的用例的代码上。在保存和加载模型时，有三个核心功能需要熟悉：torch.save:将序列化对象保存到磁盘。该函数使用Python的pickle模块中的utility进行序列化。可以使用该模块保存各种对象的模型、张量，以及词典。torch.load:使用pickle的unpickling工具将...
复制链接

扫一扫