TensorFlow 2.10训练CNN时GPU利用率低、内存攀升问题咨询

核心疑问：GPU利用率从未超过10%（常低至2%-3%）但显存缓慢攀升，是否意味着GPU未正常工作？使用更大训练集时，首个epoch末尾变慢，验证测试阶段疑似冻结。

环境与背景 #

• 运行环境：Windows系统 + NVIDIA GeForce RTX 3070 + TensorFlow 2.10
• 自我判断：GPU应处于运行状态，代码采用逐批次从磁盘读取数据的方式
• 数据集情况：约140万张尺寸为(32,32,4)的训练切片，使用14.4万张小子集可完成20轮epoch训练

GPU信息日志 #

2022-12-01 13:58:25.441213: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1616] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 5472 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 3070, pci bus id: 0000:01:00.0, compute capability: 8.6

2022-12-01 13:58:28.187036: I tensorflow_io/core/kernels/cpu_check.cc:128] Your CPU supports instructions that this TensorFlow IO binary was not compiled to use: SSE SSE2 SSE3 SSE4.1 SSE4.2 AVX AVX2 FMA

数据加载代码 #

# 获取包含所有水体和陆地图像的文件路径
data_dir = folder + '/*/*'

# 创建数据集并打乱
ds = tf.data.Dataset.list_files(data_dir, shuffle=True)

# 读取图像和标签，设置批次大小
ds = ds.shuffle(len(ds)).map(process_image).batch(bs) # .repeat(epoch)

# 设置验证集大小
training_size = int(len(ds)*0.66)

# 划分训练集和验证集
Xy_train = ds.take(training_size)
Xy_test = ds.skip(training_size)

....
CNN.fit(Xy_train, epochs=epochs, validation_data=(Xy_test))

问题分析与优化建议 #

1. GPU利用率低的核心原因 #

GPU利用率低但显存持续攀升，说明GPU确实在工作，但CPU数据预处理/磁盘IO速度跟不上GPU计算速度，导致GPU大部分时间处于等待数据的闲置状态。大训练集时验证阶段疑似冻结，本质也是验证集的数据加载未优化，加上数据量更大，IO阻塞更严重。

2. 针对性优化方案 #

（1）优化tf.data.Dataset加载流程

• 预取数据：在数据集末尾添加.prefetch(tf.data.AUTOTUNE)，让GPU计算当前批次时，CPU提前准备下一批数据，避免等待。
• 并行映射：map操作添加num_parallel_calls=tf.data.AUTOTUNE，利用多CPU核心并行处理图像加载和预处理，提升数据准备速度。
• 缓存数据：内存足够时用.cache()将数据集缓存到内存；内存不足则用.cache("cache_file")缓存到磁盘，避免重复读取原始文件。
• 调整shuffle缓冲区：ds.shuffle(len(ds))会把140万条文件路径全部加载到内存，拖慢速度且占用资源，建议设置固定大小缓冲区（如ds.shuffle(10000)），平衡打乱效果与性能。

优化后数据集代码示例：

ds = tf.data.Dataset.list_files(data_dir, shuffle=True)
ds = ds.shuffle(10000) \
       .map(process_image, num_parallel_calls=tf.data.AUTOTUNE) \
       .cache() \
       .batch(bs) \
       .prefetch(tf.data.AUTOTUNE)

（2）调整批次大小

若当前批次bs过小，GPU无法充分利用计算资源。可逐步增大bs（如从32→64→128，直到显存接近饱和），提升GPU利用率。

（3）优化预处理函数process_image

检查函数中是否有非TensorFlow原生的耗时CPU操作，尽量替换为tf.image系列API，让预处理操作可被TensorFlow优化，甚至部分转移到GPU执行。

（4）CPU指令集提示说明

日志中提到的CPU指令集未利用问题，对性能有一定影响，但不是核心瓶颈。若需优化，可从源码编译TensorFlow IO或使用兼容更高指令集的预编译包，优先级低于数据加载优化。

3. 验证GPU正常工作的方法 #

运行一个简单的GPU密集型任务，观察GPU利用率是否飙升：

import tensorflow as tf
# 创建大矩阵并在GPU上计算
a = tf.random.normal((10000, 10000))
b = tf.random.normal((10000, 10000))
c = tf.matmul(a, b)

若此时GPU利用率能达到较高水平，说明GPU硬件和TensorFlow配置正常，问题确实出在数据加载环节。

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