到目前为止,我们讨论了如何处理数据以及如何构建、训练和测试深度学习模型。然而,在某些时候,我们希望对学习的模型感到满意,我们希望保存结果以供以后在各种情况下使用(甚至可能在部署中进行预测)。此外,在运行较长的训练过程时,最佳做法是定期保存中间结果(检查点),以确保如果我们被服务器的电源线绊倒,我们不会损失几天的计算量。因此,是时候学习如何加载和存储单个权重向量和整个模型了。本节解决这两个问题。
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
6.6.1. 加载和保存张量
对于单个张量,我们可以直接调用load
和save
函数分别进行读写。这两个函数都需要我们提供一个名称,并且save
需要将要保存的变量作为输入。
我们现在可以将存储文件中的数据读回内存。
我们可以存储张量列表并将它们读回内存。
(tensor([0, 1, 2, 3]), tensor([0., 0., 0., 0.]))
(array([0., 1., 2., 3.]), array([0., 0., 0., 0.]))
(Array([0., 1., 2., 3.], dtype=float32),
Array([0., 0., 0., 0.], dtype=float32))
我们甚至可以编写和读取从字符串映射到张量的字典。当我们想要读取或写入模型中的所有权重时,这很方便。
{'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}
{'x': array([0., 1., 2., 3.]), 'y': array([0., 0., 0., 0.])}
array({'x': Array([0, 1, 2, 3], dtype=int32), 'y': Array([0., 0., 0., 0.], dtype=float32)},
dtype=object)
array({'x': <tf.Tensor: shape=(4,), dtype=int32, numpy=array([0, 1, 2, 3], dtype=int32)>, 'y': <tf.Tensor: shape=(4,), dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>},
dtype=object)
6.6.2. 加载和保存模型参数
保存单个权重向量(或其他张量)很有用,但如果我们想保存(并稍后加载)整个模型,它会变得非常乏味。毕竟,我们可能散布着数百个参数组。出于这个原因,深度学习框架提供了内置功能来加载和保存整个网络。需要注意的一个重要细节是,这会保存模型参数而不是整个模型。例如,如果我们有一个 3 层的 MLP,我们需要单独指定架构。这样做的原因是模型本身可以包含任意代码,因此它们不能自然地序列化。因此,为了恢复模型,我们需要用代码生成架构,然后从磁盘加载参数。让我们从我们熟悉的 MLP 开始。
class MLP(nn.Module):
def __init__(self):
super().__init__()
self.hidden = nn.LazyLinear(256)
self.output = nn.LazyLinear(10)
def forward(self, x):
return self.output(F.relu(self.hidden(x)))
net = MLP()
X = torch.randn(size=(2, 20))
Y = net(X)