Vgg-Face模型运行触发ResourceExhaustedError问题求助

可行解决方案 #

1. 动态分配GPU显存 #

TensorFlow默认会占用全部GPU显存，可通过设置动态分配避免显存耗尽：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

也可以直接限制显存使用量（根据你的GPU显存调整数值，比如设为4096MB）：

tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)]
)

2. 分批加载输入数据 #

避免一次性将所有数据加载到内存，改用tf.data.Dataset分批处理：

# 转换为Dataset格式
train_dataset = tf.data.Dataset.from_tensor_slices((train_x, train_y)).batch(batch_size).prefetch(tf.data.AUTOTUNE)
val_dataset = tf.data.Dataset.from_tensor_slices((val_x, val_y)).batch(batch_size).prefetch(tf.data.AUTOTUNE)

# 训练时传入Dataset对象
score = race_model.fit(train_dataset, epochs=epochs, validation_data=val_dataset, callbacks=[checkpointer, early_stop])

3. 降低模型显存占用 #

Vgg-Face参数量较大，可通过以下方式精简：

• 将部分普通卷积层替换为深度可分离卷积（DepthwiseConv2D）
• 减少卷积层的滤波器数量
• 移除模型中冗余的全连接层

4. 清理内存碎片 #

训练前手动清理内存，避免显存累积：

import gc
gc.collect()
tf.keras.backend.clear_session()

5. 调整输入数据类型 #

确保输入数据为float32类型（比float64占用显存少一半）：

train_x = train_x.astype('float32')
val_x = val_x.astype('float32')

6. 关闭Eager执行模式（可选） #

若Eager模式导致显存管理问题，可切换到Graph模式：

tf.compat.v1.disable_eager_execution()

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