张量流中的内存泄漏问题

Question

我的 TensorFlow 1.14 出现内存泄漏。我参考了 TensorFlow 的各种 GitHub 问题和内存泄漏来解决我的问题，我遵循了答案的建议，这似乎解决了问题。但是，它在这里不起作用。我什至将代码移植到 Tensorflow 2.1 和 2.3，但仍然无法解决问题。

每当我加载模型时，就会出现内存泄漏。我尝试在加载模型并使用垃圾收集 API 后清除会话，但泄漏仍然存在。

为了重现内存泄漏，我创建了一个简单的示例。我使用下面的函数来检查 python 进程使用的内存。

def memory_usage_func():
    import os
    import psutil
    process = psutil.Process(os.getpid())
    mem_used = process.memory_info()[0] >> 20
    print("Memory used:", mem_used)
    return mem_used

下面是加载模型和检查内存使用情况的函数：

for i in range(100):
    model = load_model('./model_example.h5', compile=False)
    del model
    memory_usage_func()

在上面的代码中，内存泄漏问题仍然存在。此外，我尝试进行预测。为此，我创建了一个会话，加载模型并运行 predict()。我也面临同样的内存泄漏问题。我使用tf.keras.backend.clear_session()和gc.collect()模型加载后。但是，它无法清除会话并释放内存。

Answer 1

当我尝试使用预先训练的嵌入模型来生成嵌入作为输入特征集时，我遇到了类似的问题。使用universal-sentence-encoder-4时，用于生成嵌入的内存不会被释放。既tf.keras.backend.clear_session()没有gc.collect()帮助。

我最终使用multiprocessing模块启动一个“子进程”来重复

1. 用多进程做某事
2. 杀死进程（释放内存）

import multiprocessing as mp
import time

def get_embedding(model_dir, text, q):
    print("      loading encoder")
    encoder = hub.load(model_dir)
    print("      generate embedding")
    encoded = encoder(text).numpy()
    print("      returning embedding")
    q.put(encoded)


manual_wait_time = 20

## for iterate
q = mp.Queue()
p = mp.Process(target=get_embedding, 
args=(encoder_dir, some_text, q, ))
p.start()
## p.join() ## hangs forever see reference
time.sleep(manual_wait_time)
print("manually wait %s seconds"%manual_wait_time)
embedded_mini_batch = q.get()
q.close()
p.terminate()

参考：

• 如何使用子进程强制Python释放内存？
• https://sopython.com/canon/82/programs-using-multiprocessing-hang-deadlock-and-never-complete/
• Process.join() 和队列不适用于大量数据