本文主要涉及到3个函数:
torch.save: 使用Python的pickle实用程序将对象进行序列化,然后将序列化的对象保存到disk,可以保存各种对象,包括模型、张量和字典等。
torch.load: 使用pickle unpickle工具将pickle的对象文件反序列化为内存。
torch.nn.Module.load_state_dict: 用反序列化的state_dict来加载模型参数。
1 读写tensor
1.1 单个张量
import torch
x = torch.tensor([3.,4.])
torch.save(x, 'x.pt')
x1 = torch.load('x.pt')
print(x1)
输出:
tensor([3., 4.])
1.2 张量列表和张量词典
y = torch.ones((4,2))
torch.save([x,y],'xy.pt')
torch.save({'x':x, 'y':y}, 'xy_dict.pt')
xy = torch.load('xy.pt')
xy_dict = torch.load('xy_dict.pt')
print(xy)
print(xy_dict)
输出:
[tensor([3., 4.]), tensor([[1., 1.],
[1., 1.],
[1., 1.],
[1., 1.]])]
{'x': tensor([3., 4.]), 'y': tensor([[1., 1.],
[1., 1.],
[1., 1.],
[1., 1.]])}
2 保存和加载模型
2.1 state_dict
state_dict是一个从每一个层的名称映射到这个层的参数Tesnor的字典对象。
注意,只有具有可学习参数的层(卷积层、线性层等)和注册缓存(batchnorm’s running_mean)才有state_dict中的条目。优化器(torch.optim)也有一个state_dict,其中包含关于优化器状态以及所使用的超参数的信息。
from torch import nn
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)
net = MLP()
print(net.state_dict())
print('\n',net.state_dict()['output.weight'])
optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
print(optimizer.state_dict())
输出:
OrderedDict([('hidden.weight', tensor([[ 0.1493, -0.1645, 0.5165],
[-0.0773, 0.3487, 0.4961]])), ('hidden.bias', tensor([-0.4560, -0.5111])), ('output.weight', tensor([[0.3356, 0.4229]])), ('output.bias', tensor([-0.6374]))])
tensor([[0.3356, 0.4229]])
{'state': {},
'param_groups': [{'lr': 0.001,
'momentum': 0.9,
'dampening': 0,
'weight_decay': 0,
'nesterov': False,
'params': [139788808398168,
139788808397952,
139788808397736,
139788808398096]}]}
2.2 保存和加载
PyTorch中保存和加载训练模型有两种常见的方法:
仅保存和加载模型参数(state_dict);
保存和加载整个模型。
2.2.1 保存和加载state_dict(推荐方式)
torch.save(net.state_dict(), 'net_state_dict.pt')## 后缀名一般写为: .pt或.pth
net1 = MLP()
net1.load_state_dict(torch.load('net_state_dict.pt'))
输出:
<All keys matched successfully>
注意: load_state_dict() 接受一个词典对象,而不是一个指向对象的路径。因此你需要先使用torch.load()来反序列化。比如,你不能直接这么用model.load_state_dict(PATH)。
2.2.2 保存和读写整个模型
这个就相对来说比较简单了。
torch.save(net, 'net.pt')
net2 = torch.load('net.pt')
**注意:**以这种方式保存模型将使用Python的pickle模块保存整个模块。 这种方法的缺点是序列化的数据被绑定到特定的类,并且在保存模型时使用了确切的词典结构。 这样做的原因是因为pickle不会保存模型类本身。 而是将其保存到包含这个类的文件的路径,该路径在加载时使用。 因此,在其他项目中使用或重构后,您的代码可能会以各种方式中断。
2.2.3 保存和加载checkpiont
## Save
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': 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']
保存用于检查或继续训练的checkpiont时,你必须保存的不只是模型的state_dict。 保存优化器的state_dict也很重要,因为它包含随着模型训练而更新的缓冲区和参数。 你可能要保存的其他项目包括你未设置的时间段,最新记录的训练损失,外部torch.nn.Embedding层等。
2.2.4 在一个文件中保存多个模型
#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'])
常见的PyTorch约定是使用.tar文件扩展名保存这些检查点。
2.3 使用来自不同模型的参数进行模型热启动
这种方法一般用于迁移学习。利用经过训练的参数,即使只有少数几个可用的参数,也将有助于热启动训练过程,并希望与从头开始训练相比,可以更快地收敛模型。
torch.save(modelA.state_dict(), PATH)
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
无论是从缺少某些键的部分state_dict加载,还是加载比要加载的模型更多的key的state_dict,都可以在load_state_dict()函数中将strict参数设置为False,以忽略不匹配键。
如果你想要将一个层的参数加载到另一个层,但是一些keys不匹配,你只需改变你所加载的state_dict中的名称即可。
3 跨设备保存和加载模型
3.1 在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))
3.2 在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)
3.3 在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)
4 保存torch.nn.DataParallel的模型
torch.save(model.module.state_dict(), PATH)
# Load to whatever device you want,加载方法使用常规方式即可。
参考链接: