到目前为止,我们讨论了如何处理数据以及如何构建、训练和测试深度学习模型。然而,在某些时候,我们希望对学习的模型感到满意,我们希望保存结果以供以后在各种情况下使用(甚至可能在部署中进行预测)。此外,在运行较长的训练过程时,最佳做法是定期保存中间结果(检查点),以确保如果我们被服务器的电源线绊倒,我们不会损失几天的计算量。因此,是时候学习如何加载和存储单个权重向量和整个模型了。本节解决这两个问题。
import torch
from torch import nn
from torch.nn import functional as F
import flax
import jax
from flax import linen as nn
from flax.training import checkpoints
from jax import numpy as jnp
from d2l import jax as d2l
import numpy as np
import tensorflow as tf
6.6.1. 加载和保存张量
对于单个张量,我们可以直接调用load
和save
函数分别进行读写。这两个函数都需要我们提供一个名称,并且save
需要将要保存的变量作为输入。
我们现在可以将存储文件中的数据读回内存。
tensor([0, 1, 2, 3])
x2 = jnp.load('x-file.npy', allow_pickle=True)
x2
Array([0, 1, 2, 3], dtype=int32)
我们可以存储张量列表并将它们读回内存。
y = torch.zeros(4)
torch.save([x, y],'x-files')
x2, y2 = torch.load('x-files')
(x2, y2)
(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)
评论
查看更多