使用多个 GPU 时如何保存我的模型？

Question

当我在 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文件还是我应该得到那个文件夹？ 如果我应该获取该文件夹，那么如何加载该文件夹？

*使用 tensorflow 2.0.0 版

*使用import tensorflow.keras as keras

Answer 1

使用超过 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调用中设置您的损失。 请记住，如果您只想使用模型进行预测，则无需编译。