TensorFlow GPU内存使用机制及显存不足问题排查（GTX950）

一、TensorFlow的GPU数据处理机制 #

正常训练流程中，TensorFlow不会把整个数据集一次性载入GPU显存，而是按照你设置的批量大小（batch size），逐批将数据从CPU内存加载到GPU显存中进行训练。训练完一批数据后，会自动释放该批数据占用的GPU显存（除非计算图中的变量或中间张量被刻意保留）。

但如果你的代码是先把25GB数据集一次性加载到CPU内存，再转换成TensorFlow张量（比如直接用tf.convert_to_tensor或者Keras的fit(x=full_dataset)），此时CPU内存过载可能间接导致GPU显存分配异常；另外，如果模型本身的参数量+中间计算所需显存超过2GB，哪怕batch size=2也会触发OOM错误。

二、显存不足问题的解决办法 #

针对GTX 950（2GB显存）和25GB数据集的情况，可尝试以下方案：

1. 使用数据生成器按需加载数据 #

不要一次性加载整个数据集，改用TensorFlow的tf.data.Dataset或Keras的ImageDataGenerator（图像数据场景），实现按需加载、逐批预处理：

# 以tf.data处理图像数据集为例
import tensorflow as tf

def parse_data(file_path):
    # 单条数据的读取与预处理逻辑
    img = tf.io.read_file(file_path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, (224, 224))
    img = tf.cast(img, tf.float32) / 255.0
    # 此处省略标签读取逻辑，需根据实际数据集调整
    label = tf.constant(0, dtype=tf.int32)
    return img, label

# 获取所有数据文件路径
file_dataset = tf.data.Dataset.list_files("/path/to/your/dataset/*")
# 并行映射预处理逻辑，生成批量数据集
train_dataset = file_dataset.map(parse_data, num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.batch(2)

# 训练时直接传入生成的数据集
model.fit(train_dataset, epochs=10)

这种方式只会在需要时加载单批数据，既避免CPU内存过载，也能控制GPU的显存占用。

2. 减小模型参数量 #

如果你的DNN模型本身过大，哪怕batch size=2，模型参数+中间激活值的显存占用也可能超过2GB：

• 减少卷积层的滤波器数量（比如把Conv2D(256, ...)改为Conv2D(128, ...)）
• 移除冗余的全连接层，或缩小全连接层的神经元规模
• 替换为轻量化模型结构（比如用MobileNet、SqueezeNet替代VGG、ResNet50）

3. 启用混合精度训练 #

通过混合精度（FP16+FP32）减少显存占用，同时几乎不损失模型精度：

from tensorflow.keras.mixed_precision import set_global_policy

set_global_policy('mixed_float16')

启用后，模型的大部分参数和计算会采用FP16精度，显存占用可降低约50%。

4. 清理GPU冗余显存占用 #

在终端执行nvidia-smi命令，检查是否有其他进程（比如残留的TensorFlow会话、其他CUDA程序）占用GPU显存，关闭这些进程释放资源。

5. 优化显存分配配置 #

你已启用allow_growth=True，还可添加显存比例限制，预留部分显存给系统：

from keras.backend.tensorflow_backend import set_session
import tensorflow as tf

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.9  # 限制使用90%的GPU显存
sess = tf.Session(config=config)
set_session(sess)

6. 清理残留计算图 #

训练前手动清理之前的计算图与显存残留：

tf.keras.backend.clear_session()

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