4.5 读取和存储

发表于 2025-02-17 更新于 2025-04-16 分类于深度学习，笔记

读写`Tensor`

可以直接使用save函数和load函数分别存储和读取Tensor。save使用Python的pickle实用程序将对象进行序列化，然后将序列化的对象保存到disk，使用save可以保存各种对象,包括模型、张量和字典等。而load使用pickle unpickle工具将pickle的对象文件反序列化为内存。

下面的例子创建了Tensor变量x，并将其存在文件名同为x.pt的文件里。

import torch
from torch import nn

x = torch.ones(3)
torch.save(x, 'x.pt')

然后将数据从存储的文件读回内存。

1 2	x2 = torch.load('x.pt') print(x2)

输出：

1	tensor([1., 1., 1.])

还可以存储一个Tensor列表并读回内存。

y = torch.zeros(4)
torch.save([x, y], 'xy.pt')
xy_list = torch.load('xy.pt')
print(xy_list)

输出：

1	[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]

存储并读取一个从字符串映射到Tensor的字典。

1
2
3

torch.save({'x': x, 'y': y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
print(xy)

输出：

1	{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}

读写模型

`state_dict`

在PyTorch中，Module的可学习参数(即权重和偏差)，模块模型包含在参数中(通过model.parameters()访问)。state_dict是一个从参数名称隐射到参数Tesnor的字典对象。

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()
net.state_dict()

输出：

OrderedDict([('hidden.weight', tensor([[ 0.2448,  0.1856, -0.5678],
                      [ 0.2030, -0.2073, -0.0104]])),
             ('hidden.bias', tensor([-0.3117, -0.4232])),
             ('output.weight', tensor([[-0.4556,  0.4084]])),
             ('output.bias', tensor([-0.3573]))])

注意，只有具有可学习参数的层(卷积层、线性层等)才有state_dict中的条目。优化器(optim)也有一个state_dict，其中包含关于优化器状态以及所使用的超参数的信息。

1 2	optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9) optimizer.state_dict()

输出：

{'param_groups': [{'dampening': 0,
   'lr': 0.001,
   'momentum': 0.9,
   'nesterov': False,
   'params': [4736167728, 4736166648, 4736167368, 4736165352],
   'weight_decay': 0}],
 'state': {}}

保存和加载模型

PyTorch中保存和加载训练模型有两种常见的方法:

仅保存和加载模型参数(state_dict)；
保存和加载整个模型。

无论是保存整个模型还是只保存模型参数，在加载模型时都需要在代码中定义出模型的结构。

1. 保存和加载`state_dict`(推荐方式)

保存：

1	torch.save(model.state_dict(), PATH) # 推荐的文件后缀名是pt或pth

加载：

1 2	model = TheModelClass(args, *kwargs) model.load_state_dict(torch.load(PATH))

2. 保存和加载整个模型

保存：

1	torch.save(model, PATH)

加载：

1	model = torch.load(PATH)

注意，模型类必须一致。保存整个模型时，模型类的定义会被序列化。因此，在加载模型时，模型类的定义必须与保存时完全一致。如果模型类的定义发生了变化（例如修改了模型的结构或方法或模型类的名称），可能会导致加载失败或行为异常。保存整个模型时，模型的结构、参数、优化器状态等都会被保存下来，适合需要完整恢复训练场景的情况。但是灵活性较低，由于保存的是整个模型对象，因此在加载时要求模型类的定义必须与保存时一致。如果模型类的定义发生变化，可能会导致加载失败。

采用方法一来实验一下:

X = torch.randn(2, 3)
Y = net(X)

PATH = "./net.pt"
torch.save(net.state_dict(), PATH)

net2 = MLP()
net2.load_state_dict(torch.load(PATH))
Y2 = net2(X)
Y2 == Y

输出：

1 2	tensor([[1], [1]], dtype=torch.uint8)

因为这net和net2都有同样的模型参数，那么对同一个输入X的计算结果将会是一样的。上面的输出也验证了这一点。

读写Tensor