Run this notebook online:Binder or Colab: Colab

5.4. 读写文件

到目前为止,我们讨论了如何处理数据,以及如何构建、训练和测试深度学习模型。然而,有时我们对所学的模型足够满意,我们希望保存训练的模型以备将来在各种环境中使用(可能部署进行预测)。此外,当运行一个耗时较长的训练过程时,最佳实践是定期保存中间结果(检查点),以确保在服务器电源被不小心断掉时不会损失几天的计算结果。因此,现在是时候学习如何加载和存储权重向量和整个模型。本节将讨论这些问题。

5.4.1. 加载和保存 NDArray

对于单个 NDArray,我们可以直接调用 encode()decode() 函数分别读写它们。为我们可以使用这两个函数结合 FileInputStream, FileOutputStream 来实现对文件的读写。

%load ../utils/djl-imports
NDManager manager = NDManager.newBaseManager();

NDArray x = manager.arange(4);
try (FileOutputStream fos = new FileOutputStream("x-file")) {
    fos.write(x.encode());
}
x
ND: (4) gpu(0) int32
[ 0,  1,  2,  3]

我们现在可以将存储在文件中的数据读回内存。

NDArray x2;
try (FileInputStream fis = new FileInputStream("x-file")) {
    // We use the `Utils` method `toByteArray()` to read
    // from a `FileInputStream` and return it as a `byte[]`.
    x2 = NDArray.decode(manager, Utils.toByteArray(fis));
}
x2
ND: (4) gpu(0) int32
[ 0,  1,  2,  3]

我们也可以存储一个 NDList,然后把它们读回内存。

NDList list = new NDList(x, x2);
try (FileOutputStream fos = new FileOutputStream("x-file")) {
    fos.write(list.encode());
}
try (FileInputStream fis = new FileInputStream("x-file")) {
    list = NDList.decode(manager, Utils.toByteArray(fis));
}
list
NDList size: 2
0 : (4) int32
1 : (4) int32

5.4.2. 加载和保存模型参数

保存单个权重向量(或其他NDArray)确实是有用的,但是如果我们想保存整个模型,并在以后加载它们。单独保存每个向量则会变得很麻烦。毕竟,我们可能有数百个参数散布在各处。因此,深度学习框架提供了内置函数来保存和加载整个网络。需要注意的一个重要细节是,这将保存模型的参数而不是保存整个模型。例如,如果我们有一个3层多层感知机,我们需要单独指定结构。因为模型本身可以包含任意代码,所以模型本身难以序列化。因此,为了恢复模型,我们需要用代码生成结构,然后从磁盘加载参数。让我们从熟悉的多层感知机开始尝试一下。

public SequentialBlock createMLP() {
    SequentialBlock mlp = new SequentialBlock();
    mlp.add(Linear.builder().setUnits(256).build());
    mlp.add(Activation.reluBlock());
    mlp.add(Linear.builder().setUnits(10).build());
    return mlp;
}

SequentialBlock original = createMLP();

NDArray x = manager.randomUniform(0, 1, new Shape(2, 5));

original.initialize(manager, DataType.FLOAT32, x.getShape());

ParameterStore ps = new ParameterStore(manager, false);
NDArray y = original.forward(ps, new NDList(x), false).singletonOrThrow();

y
ND: (2, 10) gpu(0) float32
[[-1.0524,  0.4173,  0.9115, -0.9705, -1.5318, -0.6548,  0.0033, -0.6443, -0.5181, -0.0511],
 [-0.4358,  0.0545,  0.5989, -0.5398, -1.0074, -0.3614, -0.1044, -0.4891, -0.2387,  0.1797],
]

接下来,我们将模型的参数存储为一个叫做“mlp.params”的文件:

// Save file
File mlpParamFile = new File("mlp.param");
DataOutputStream os = new DataOutputStream(Files.newOutputStream(mlpParamFile.toPath()));
original.saveParameters(os);

为了恢复模型,我们实例化了原始多层感知机模型。我们没有随机初始化模型参数,而是直接读取文件中存储的参数。

// Create duplicate of network architecture
SequentialBlock clone = createMLP();
// Load Parameters
clone.loadParameters(manager, new DataInputStream(Files.newInputStream(mlpParamFile.toPath())));

由于两个实例具有相同的模型参数,在输入相同的X时,两个实例的计算结果应该相同。让我们来验证一下。

// Original model's parameters
PairList<String, Parameter> originalParams = original.getParameters();
// Loaded model's parameters
PairList<String, Parameter> loadedParams = clone.getParameters();

for (int i = 0; i < originalParams.size(); i++) {
    if (originalParams.valueAt(i).getArray().equals(loadedParams.valueAt(i).getArray())) {
        System.out.printf("True ");
    } else {
        System.out.printf("False ");
    }
}
True True True True
NDArray yClone = clone.forward(ps, new NDList(x), false).singletonOrThrow();

y.equals(yClone);
true

5.4.3. 小结

  • encode()decode() 函数可用于 NDArrayNDList 对象的文件读写。

  • 我们可以保存和加载模型文件。

5.4.4. 练习

  1. 即使不需要将经过训练的模型部署到不同的设备上,存储模型参数还有什么实际的好处?

  2. 假设我们只想复用网络的一部分,以将其合并到不同的网络结构中。比如说,如果你想在一个新的网络中使用之前网络的前两层,你该怎么做?