Pytorch保存和加载模型完全指南: 关于使用Pytorch读写模型的一切方法

本文是一篇关于如何用Pytorch保存和加载模型的指南。

本文主要涉及到3个函数：

1. torch.save: 使用Python的pickle实用程序将对象进行序列化，然后将序列化的对象保存到disk，可以保存各种对象,包括模型、张量和字典等。
2. torch.load: 使用pickle unpickle工具将pickle的对象文件反序列化为内存。
3. 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,