PyTorch多GPU脚本中第二块GPU未被利用的原因排查

核心问题分析 #

• 图像数据默认加载到CUDA:0：如果主进程中img_list的图像已被加载为CUDA:0上的张量，子进程会继承这些张量，但CUDA张量无法跨GPU直接迁移。即使子进程指定CUDA:1，计算仍会在CUDA:0执行，导致第二块GPU闲置，第一块GPU负载翻倍。
• 模型未正确迁移到指定GPU：若MyModel的__init__方法未将模型参数显式移动到传入的device，参数会默认留在CUDA:0，实际计算仍在第一块GPU运行。
• 进程启动方式的CUDA上下文继承：Linux/macOS下torch.multiprocessing默认用fork启动，会完整复制主进程的CUDA上下文，子进程即使指定CUDA:1，也会受继承的上下文影响优先使用CUDA:0资源。
• 代码笔误：np.range应为np.arange；args中的img变量未定义，属于无效参数，可能干扰逻辑。

修复方案 #

1. 主进程仅存图像路径，子进程内加载并指定GPU #

主进程不提前加载图像为张量，仅存储路径，子进程启动后按需加载并迁移到目标GPU：

if __name__ == '__main__':
    # 强制使用spawn启动方式，避免继承CUDA上下文
    mp.set_start_method('spawn')
    num_processes = 2
    processes = []

    # 主进程仅保留图像路径
    img_paths = [...]
    img_indices = np.arange(0, len(img_paths))

    for gpu_idx in range(num_processes):
        subindices = img_indices[gpu_idx::num_processes]
        # 移除无效的img参数
        p = mp.Process(target=my_single_gpu_optimization_func, args=(img_paths, subindices, gpu_idx))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

子进程函数实现：

def my_single_gpu_optimization_func(img_paths, subindices, gpu_idx):
    device = f'cuda:{gpu_idx}'
    # 加载图像并直接移到目标GPU
    for idx in subindices:
        img = load_image(img_paths[idx]).to(device)
        # 初始化模型并显式迁移到指定GPU
        model = MyModel()
        model = model.to(device)
        # 执行优化逻辑
        ...

2. 确保模型参数迁移到目标GPU #

如果MyModel的__init__未处理设备迁移，初始化后显式调用model.to(device)，确保参数全部移动到指定GPU。

3. 修正代码笔误 #

将np.range替换为np.arange，删除未定义的img参数。

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