TensorFlow—OOM问题

TensorFlow—OOM 问题

记一个关于TensorFlow的小笔记：

使用model

在 TensorFlow 2.0 中每次使用完 keras 中的 model

import tensorflow as tf

model = tf.keras.models.load_model("model.h5")
prediction=model.predict()

一定要记着 清理释放内存，否则不断调用上面的model解决问题，一会儿内存就用完了，就会出现 OOM 问题。

此前网上的 OOM 问题大都是数据太大没法导入模型，这个 OOM 问题是由于用过的模型不及时清理，TensorFlow也不会帮你自动清理造成的。

清理内存

tf.keras.backend.clear_session()

或者使用 cuda 释放内存

from numba import cuda

cuda.select_device(0)
cuda.close()