TensorFlow MirroredStrategy下GPU未利用的问题求助

一、先确认ROCm环境与GPU识别 #

1. 验证节点ROCm环境
   在作业脚本中先执行rocm-smi，确认GPU状态正常、无其他进程占用；同时运行以下代码验证TensorFlow能识别GPU：

   import tensorflow as tf
   print("可用GPU设备:", tf.config.list_physical_devices('GPU'))
   

   若输出为空，说明TF未绑定ROCm，需检查：

   • 作业脚本是否正确加载集群的ROCm模块（如module load rocm/5.7 tensorflow-rocm/2.16），确保使用的是专为AMD GPU编译的TensorFlow-ROCm版本，而非普通CUDA版。
   • 集群是否分配了GPU资源：Slurm脚本需包含--gres=gpu:2，确保作业拿到2个GPU的独占权限。

2. 显式指定GPU设备
   HPC集群中TF可能不会自动选择GPU，需在MirroredStrategy初始化时显式指定设备：

   strategy = tf.distribute.MirroredStrategy(devices=["/GPU:0", "/GPU:1"])
   

   避免依赖自动检测导致TF fallback到CPU。

二、检查MirroredStrategy的使用规范 #

1. 所有模型相关操作必须在策略作用域内
   模型构建、编译、优化器初始化都要放在strategy.scope()下，包括自定义层/回调中涉及优化器的逻辑：

   with strategy.scope():
       model = build_your_cnn_model()
       optimizer = tf.keras.optimizers.Adam()
       model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy')
   

   若优化器或模型在作用域外初始化，会导致分布式训练失效，计算默认跑在CPU。

2. 自定义回调的分布式兼容
   切换优化器的自定义回调中，修改优化器参数的操作需通过strategy.run()执行，确保所有GPU上的优化器同步更新：

   def on_epoch_end(self, epoch, logs=None):
       def update_optimizer():
           self.model.optimizer.lr.assign(new_lr)
       strategy.run(update_optimizer)
   

   避免直接修改优化器导致仅CPU端参数更新，GPU端无变化。

三、排查数据集性能瓶颈（GPU饥饿的常见原因） #

1. 优化数据加载与预处理
   自定义EMGDataset类需基于tf.data.Dataset实现，并启用并行预处理与预取：

   dataset = EMGDataset(...)
   dataset = dataset.batch(batch_size)
   dataset = dataset.prefetch(tf.data.AUTOTUNE)  # 让TF提前准备下一批数据
   dataset = dataset.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)  # 并行预处理
   

   若数据从磁盘读取，建议转成TFRecord格式减少IO开销，避免CPU单线程加载跟不上GPU计算速度。

2. 确认输入张量形状与设备位置
   EMG输入是2D数组（通道×时间），需转成TF CNN期望的格式（如(batch_size, time_steps, num_channels)），并确保张量被分配到GPU：

   def preprocess_fn(data):
       # 转置通道与时间维度，添加batch维度
       data = tf.transpose(data, perm=[1, 0])
       data = tf.expand_dims(data, axis=0)
       return data, label
   

   可在预处理函数中打印张量设备，确认数据已进入GPU：

   print("数据张量设备:", data.device)
   

四、HPC作业配置优化 #

1. 调整CPU资源分配
   当前仅用3个CPU，可能不足以支撑数据预处理。在Slurm脚本中增加CPU核心数（如--cpus-per-task=8），让tf.data能充分利用多CPU并行加载数据，避免GPU等待。

2. 禁用CPU绑定（可选）
   部分集群的CPU绑定策略会限制TF的数据并行效率，可在作业脚本中添加--cpu-bind=none尝试解除绑定。

五、调试与验证步骤 #

1. 先验证单GPU训练
   注释掉MirroredStrategy代码，用单GPU训练，运行rocm-smi监控GPU使用率。若单GPU能正常利用，说明问题出在多分布式策略的配置；若单GPU也无使用率，回到环境与模型输入的排查。

2. 用TF Profiler定位瓶颈
   添加Profiler代码，记录训练过程中的设备使用情况：

   tf.profiler.experimental.server.start(6009)
   # 训练代码
   tf.profiler.experimental.stop()
   

   通过Profiler可直观看到是数据加载慢、模型计算未到GPU，还是其他环节阻塞。

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