多次加载 Keras 模型时 GPU 内存泄漏答案

【问题标题】：GPU memory leakage when I load Keras models multi times多次加载 Keras 模型时 GPU 内存泄漏
【发布时间】：2020-07-23 00:24:08
【问题描述】：

我有一个文件train.py通过tensorflow.keras训练神经网络模型，最好的模型会根据train_loss保存为best_train.h5。培训时间约为 7 天。

我有另一个文件test.py 来测试train.py 运行期间保存的模型。在test.py 模型中，我每小时加载一次best_train.h5 以查看测试性能。代码如下：

for i in range(7*24):
  time.sleep(1*60*60)
  model = tf.keras.models.load_model('best_train.h5')
  model.predict(test_data)

发现每次加载best-train.h5，占用的GPU内存都会增加。而在大约 200 次迭代之后，GPU 内存就会耗尽。会出现 OOM 错误。

我知道tf.clear_session() 可以释放 GPU 内存。但是此命令将清除所有会话和图形。我不是我想要的。因为在test.py还有我持有的其他型号。

【问题讨论】：

标签： python tensorflow keras

【解决方案1】：

因为 Keras 在所有模型之间共享全局会话。您可以创建新图表并分配会话仅用于预测：

self.graph = tf.Graph()
with self.graph.as_default():
   self.session = tf.Session(graph=self.graph)
   with self.session.as_default(): 
       # Load your model and preform prediction

您的 GPU 内存应在您完成预测后立即释放。

【讨论】：