时间戳

首页 > 解决方案 > 使用多个 GPU 时如何保存模型？

python - 使用多个 GPU 时如何保存模型？

问题描述

当我在 1 个 GPU 上进行训练时，我会调用model.save('test.model')它，并且在训练结束时我会得到一个文件，test.model然后我可以在预测代码中使用它。

现在我在 2 个 GPU 上进行训练，当我打电话时model.save('test_2gpus.model')，我得到一个名为test_2gpus.model. 在此文件夹中还有 2 个名为的文件夹：assets和variables，以及另一个名为saved_model.pb. 较assets旧的是空的，variables文件夹中有 3 个文件：variables.data-00000-of-00002、variables.data-00001-of-00002和variables.index.

使用超过 1 个 GPU 时，我是否必须以不同方式保存模型？

我应该得到一个.model文件还是应该得到那个文件夹？如果我应该获取该文件夹，那么如何加载该文件夹？

*使用张量流版本 2.0.0

*使用import tensorflow.keras as keras

标签： pythontensorflowmachine-learningkeras

解决方案

使用超过 1 个 GPU 时，我是否必须以不同方式保存模型？

请参阅此处的教程https://www.tensorflow.org/tutorials/distribute/save_and_load

不，您可以使用 保存它model.save(path)，然后使用 加载tf.keras.models.load_model(path)。

我应该得到一个 .model 文件还是应该得到那个文件夹？如果我应该获取该文件夹，那么如何加载该文件夹？

您应该获得一个文件夹，然后通过引用该文件夹进行加载。

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    model = ... # build model
    model.compile()

folder = "test_2gpus.model"
model.save(folder)
model_loaded = tf.keras.models.load_model(folder)

如果你得到一个unknown loss错误，那是因为模型是在加载时编译的。tf.keras.models.load_model有一个默认的关键字compile参数True。要绕过错误，请设置compile=False，然后在model.compile通话中设置您的损失。请记住，如果您只想使用模型进行预测，则不需要编译。

推荐阅读