Keras/TensorFlow下GPU、CPU等硬件高效利用问题求助

硬件配置 #

• 2018年Lambda机器：24核CPU、4块GPU、4TB SSD
• 2022年Dell机器：40核CPU、3块A6000 GPU、1TB NVMe SSD + 1TB SSD + 2块8TB HDD

训练现状 #

• 限制1核CPU+1块GPU训练4万参数模型，单epoch耗时9-10分钟；增加CPU/GPU数量后，耗时无变化
• 训练初期磁盘负载高，数据读入内存后无磁盘活动，确认数据可完全存入内存
• 当前batch size为500，调整大小后训练时间无明显改善甚至变慢
• 关闭GPU仅用CPU时，Dell机器40核CPU利用率20%-50%，偶尔单核负载超95%，单epoch耗时增至12-13分钟

资源监控结果 #

• CPU（htop）：1核时负载极高；2核时每核约60%；10核时每核约13%；20核时每核约7%
• GPU（nvtop/nvidia-smi）：仅GPU0利用率10-25%，其余GPU无内存写入、几乎未使用
• 磁盘（atop）：仅训练初期有高负载，后续无活动

模型与环境细节 #

• 环境：Python 3.8、TensorFlow 2.10
• 模型结构：Siamese双分支，每分支含17层3x3卷积（伸缩结构），输入69x69x3图像，后续卷积层16通道；分支后接3层全连接（25、20、10单元），经product layer合并为10输出，再通过全连接层得到1个输出；使用类二元交叉熵的自定义损失函数
• 数据集：约60万张图像
• 训练方式：Keras model.fit搭配自定义数据生成器，预期多CPU并行生成数据但实际CPU大多空闲

核心问题 #

当前CPU、GPU、磁盘均未充分利用，增加硬件资源无法提升训练速度，疑似存在未定位的瓶颈（如主内存/PCI总线带宽、数据生成器并行失效等）

排查建议 #

1. 数据生成器并行优化

   • 检查自定义数据生成器是否含不可序列化对象（如自定义类实例、未关闭的文件句柄），这类问题会导致workers参数失效
   • 显式设置model.fit参数：workers=N（N建议设为CPU核心数的1/2到1倍）、use_multiprocessing=True、max_queue_size=32（根据内存调整），强制开启多进程数据预处理
   • 单独测试生成器效率：脱离模型训练，统计生成器每秒输出的样本数，判断是否是生成器本身速度不足

2. GPU利用率低排查

   • 临时增加卷积通道数（如改为32/64），测试GPU利用率是否提升，排查是否因模型计算量偏小导致GPU无法满负载
   • 优化自定义损失函数：确保逻辑全向量化，避免使用Python循环（这类操作在CPU执行，会拖慢GPU等待速度），替换为TensorFlow原生算子
   • 启用多GPU策略：用tf.distribute.MirroredStrategy包裹模型，确保所有GPU参与训练，同时调整batch size适配分布式副本数

3. CPU负载异常排查

   • 拆解数据预处理步骤：若生成器存在串行操作（如全局变量依赖、单线程图像增强），需改为可并行逻辑；建议用tf.data.Dataset替代自定义生成器，通过map(..., num_parallel_calls=tf.data.AUTOTUNE)和prefetch(tf.data.AUTOTUNE)优化流水线
   • 对比tf.data.Dataset效果：将数据集转为该格式，测试训练速度与资源利用率变化

4. 其他潜在瓶颈

   • 测试内存带宽：用memtester工具检测内存读写带宽，排查是否存在内存传输瓶颈
   • 检查PCIe链路：通过nvidia-smi -q | grep Link查看GPU的PCIe链路宽度，确保运行在x16模式，避免数据传输带宽受限

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