多GPU上运行TensorFlow出现异常行为的技术问询

看了你贴的代码，能发现核心问题：你只是启动了多个训练进程，但没有给每个进程绑定对应的GPU，这样TensorFlow会默认抢占所有GPU资源，最后所有进程挤在一块GPU上，完全达不到你“每块GPU跑一个训练任务”的目标。我给你调整一下代码，补上GPU绑定和进程调度的关键逻辑：

1. 编写带GPU绑定的训练函数 #

import os
import time
from multiprocessing import Process
import tensorflow as tf

def run_training(gpu_idx, param):
    # 第一步：绑定当前进程到指定GPU
    # 方式1：用环境变量限制可见GPU（推荐，简单直接）
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_idx)
    
    # 方式2：用TensorFlow API配置（适合更精细的内存管理）
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        try:
            # 只让当前进程看到指定的GPU
            tf.config.set_visible_devices(gpus[gpu_idx], 'GPU')
            # 开启内存按需分配，避免一次性占满GPU内存
            tf.config.experimental.set_memory_growth(gpus[gpu_idx], True)
        except RuntimeError as e:
            print(f"GPU配置出错: {e}")

    # 这里写你的实际训练逻辑
    print(f"训练进程启动 | 使用GPU {gpu_idx} | 参数: {param}")
    # 模拟训练过程（替换成你的训练代码）
    time.sleep(30)
    print(f"训练进程结束 | 使用GPU {gpu_idx}")

2. 主进程调度：按GPU数量控制并发 #

if __name__ == "__main__":
    # 先检测可用GPU数量
    gpu_count = len(tf.config.list_physical_devices('GPU'))
    if gpu_count == 0:
        print("未检测到可用GPU，程序退出")
        exit()
    print(f"检测到 {gpu_count} 块可用GPU")

    # 假设你的训练参数列表（每个元素对应一个训练任务的参数）
    params = [("task_1_params",), ("task_2_params",), ("task_3_params",), ("task_4_params",)]
    pool = []
    current_gpu = 0

    for param in params:
        # 给当前训练任务分配GPU
        p = Process(target=run_training, args=(current_gpu, param))
        p.start()
        pool.append(p)
        # 循环分配GPU（如果任务数超过GPU数，就等空闲GPU）
        current_gpu = (current_gpu + 1) % gpu_count

        # 当并发进程数等于GPU数时，等待有进程结束再启动新的
        while len(pool) >= gpu_count:
            time.sleep(5)
            # 清理已经结束的进程，避免列表堆积
            pool = [p for p in pool if p.is_alive()]

    # 等待所有剩余进程完成
    for p in pool:
        p.join()
    print("所有训练任务已完成！")

关键细节说明 #

• GPU绑定：通过CUDA_VISIBLE_DEVICES或者TensorFlow的API，让每个进程只能看到自己被分配的GPU，彻底避免资源争抢。
• 内存管理：开启set_memory_growth非常重要，不然单进程会占满整个GPU的显存，导致其他进程无法启动。
• 进程调度：用current_gpu % gpu_count循环分配GPU，当任务数多于GPU数时，自动等待空闲GPU，不会让系统负载过高。
• 进程清理：定期过滤掉已结束的进程，保证pool列表里只存活跃进程，避免判断逻辑出错。

额外注意事项 #

• 每个进程的训练数据最好在run_training函数内部加载，不要用全局变量共享数据，多进程下容易出现数据混乱或者内存泄漏。
• 如果你的任务是相同模型不同数据的训练，这个方案完全够用；如果是分布式训练（比如多GPU训练同一个模型），那你需要用TensorFlow的分布式策略，而不是这种单进程单GPU的方式。

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