TensorFlow循环创建模型内存泄漏：tf.keras.backend.clear_session失效求助

问题描述 #

循环中反复创建并丢弃大量神经网络模型时，旧模型会堆积在内存中，最终触发GPU内存不足崩溃。官方文档说明tf.keras.backend.clear_session()可清理冗余内容，但该命令在此场景下无效。

环境版本 #

• TensorFlow 2.8.0
• Keras 2.8.0

可复现最小示例代码 #

import tensorflow as tf
from tensorflow import keras

# 使用GPU
physical_devices = tf.config.list_physical_devices("GPU")
print("physical devices: ", physical_devices)
# 允许GPU内存动态增长
tf.config.experimental.set_memory_growth(physical_devices[0], True)

def create_nn_model():
    """初始化并返回神经网络模型"""

    Ndim = 100
    N_nodes_L1 = 1000
    N_nodes_L2 = 5000

    # 构建模型
    x_input = keras.Input(shape=[Ndim])
    L1 = keras.layers.Dense(N_nodes_L1, input_shape = [Ndim],
                                     activation="swish")(x_input)
    L2 = keras.layers.Dense(N_nodes_L2, input_shape=[N_nodes_L1],
                                     activation="swish")(L1)
    output = keras.layers.Dense(1, input_shape=[N_nodes_L2],
                                   activation = "linear")(L2)
    model = keras.Model(inputs=[x_input],
                        outputs = [output])
    # 绘制模型结构
    keras.utils.plot_model(model, "model.png", show_shapes=True)
    return(model)


for ii in range(1_000):
    print(f"Training model {ii+1} of 1,000")
    nn_model = create_nn_model()

    tf.keras.backend.clear_session()

报错信息 #

2023-06-21 18:35:33.887623: W tensorflow/core/common_runtime/bfc_allocator.cc:462] Allocator (GPU_0_bfc) ran out of memory trying to allocate 19.07MiB (rounded to 20000000)requested by op AddV2
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation.
Current allocation summary follows.
Current allocation summary follows.
2023-06-21 18:35:33.887904: I tensorflow/core/common_runtime/bfc_allocator.cc:1010] BFCAllocator dump for GPU_0_bfc
2023-06-21 18:35:33.888881: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (256):       Total Chunks: 66, Chunks in use: 56. 16.5KiB allocated for chunks. 14.0KiB in use in bin. 228B client-requested in use in bin.
2023-06-21 18:35:33.889243: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (512):       Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2023-06-21 18:35:33.889709: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (1024):      Total Chunks: 1, Chunks in use: 1. 1.2KiB allocated for chunks. 1.2KiB in use in bin. 1.0KiB client-requested in use in bin.
2023-06-21 18:35:33.889982: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (2048):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.
2023-06-21 18:35:33.890294: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (4096):      Total Chunks: 199, Chunks in use: 198. 920.0KiB allocated for chunks. 912.2KiB in use in bin. 773.4KiB client-requested in use in bin.
2023-06-21 18:35:33.890487: I tensorflow/core/common_runtime/bfc_allocator.cc:1017] Bin (8192):      Total Chunks: 0, Chunks in use: 0. 0B allocated for chunks. 0B in use in bin. 0B client-requested in use in bin.

...

2023-06-21 18:35:34.110374: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 2 Chunks of size 33554432 totalling 64.00MiB
2023-06-21 18:35:34.110547: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 33741824 totalling 32.18MiB
2023-06-21 18:35:34.110718: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 34217728 totalling 32.63MiB
2023-06-21 18:35:34.111092: I tensorflow/core/common_runtime/bfc_allocator.cc:1074] 1 Chunks of size 36870912 totalling 35.16MiB
2023-06-21 18:35:34.111263: I tensorflow/core/common_runtime/bfc_allocator.cc:1078] Sum Total of in-use chunks: 3.87GiB
2023-06-21 18:35:34.111450: I tensorflow/core/common_runtime/bfc_allocator.cc:1080] total_region_allocated_bytes_: 4162256896 memory_limit_: 4162256896 available bytes: 0 curr_region_allocation_bytes_: 4294967296
2023-06-21 18:35:34.111627: I tensorflow/core/common_runtime/bfc_allocator.cc:1086] Stats:
Limit:                      4162256896
InUse:                      4160154112
MaxInUse:                   4160154368
NumAllocs:                        2972
MaxAllocSize:                 36870912
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0

2023-06-21 18:35:34.111869: W tensorflow/core/common_runtime/bfc_allocator.cc:474] ****************************************************************************************************
2023-06-21 18:35:34.112024: W tensorflow/core/framework/op_kernel.cc:1733] RESOURCE_EXHAUSTED: failed to allocate memory

解决办法 #

1. 手动解除引用+强制垃圾回收 #

clear_session()无法立即回收Python侧的模型引用，需手动解除引用并触发垃圾回收：
修改循环代码如下：

import gc

for ii in range(1_000):
    print(f"Training model {ii+1} of 1,000")
    nn_model = create_nn_model()
    
    # 解除模型引用
    del nn_model
    # 强制触发Python垃圾回收
    gc.collect()
    
    tf.keras.backend.clear_session()

2. 复用模型结构（结构固定时） #

若每次创建的模型结构完全一致，无需重复构建，可先创建基准模型，后续通过克隆重置权重：

# 预先创建基准模型
base_model = create_nn_model()

for ii in range(1_000):
    print(f"Training model {ii+1} of 1,000")
    # 克隆模型并重置权重
    nn_model = tf.keras.models.clone_model(base_model)
    # 按需编译模型
    nn_model.compile(optimizer="adam", loss="mse")
    
    del nn_model
    gc.collect()
    tf.keras.backend.clear_session()

3. 调整GPU内存分配器 #

根据报错提示，设置环境变量缓解内存碎片：

• Linux/macOS：运行脚本前执行 export TF_GPU_ALLOCATOR=cuda_malloc_async
• Windows：命令行执行 set TF_GPU_ALLOCATOR=cuda_malloc_async 后再运行脚本

4. 升级TensorFlow版本 #

TensorFlow 2.8存在部分已知内存泄漏问题，升级到2.10+稳定版本可修复此类问题。

已知临时方案：使用multiprocess多进程模式，在子进程中创建训练模型，子进程结束后自动释放内存，但实现复杂度较高。

内容的提问来源于stack exchange，提问作者Atticus Beachy