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));

Model model = Model.newInstance("mlp");
model.setBlock(original);

Predictor predictor = model.newPredictor(new NoopTranslator());

NDArray y = ((NDList) predictor.predict(new NDList(x))).singletonOrThrow();

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

// Save file
File mlpParamFile = new File("mlp.params");
DataOutputStream os = new DataOutputStream(new FileOutputStream(mlpParamFile));
original.saveParameters(os);

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

SequentialBlock clone = createMLP();
// Load Parameters
clone.loadParameters(manager, new DataInputStream(new FileInputStream(mlpParamFile)));

由于两个实例具有相同的模型参数,在输入相同的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++) {
    System.out.println(originalParams.valueAt(i).getArray().equals(loadedParams.valueAt(i).getArray()));
}
true
true
true
true

5.4.3. 小结

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

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

5.4.4. 练习

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

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