【踩坑】修复循环设置os.environ[‘CUDA_VISIBLE_DEVICES‘]无效

小锋学长生活大爆炸

发布于 2025-05-24 05:14:33

25900

代码可运行

文章被收录于专栏：小锋学长生活大爆炸小锋学长生活大爆炸

运行总次数：0

代码可运行

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn] 如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

问题示例

for gpus in ['0', '1',  '2',  '3',  '4',  '5',  '6',  '7']:
    os.environ['CUDA_VISIBLE_DEVICES'] = gpus
    print(torch.cuda.get_device_name(0))

始终将使用第一个GPU，即CUDA:0.

问题解析

在使用 os.environ['CUDA_VISIBLE_DEVICES'] 动态设置可见的 GPU 时，通常只能在程序开始运行之前进行设置，并且这种设置在程序运行后无法在一个进程中动态改变。这是因为深度学习框架（如 TensorFlow、PyTorch 等）在初始化时就会读取 CUDA_VISIBLE_DEVICES 环境变量，并在后续操作中使用这些设置。因此，尝试在运行时通过循环动态更改 CUDA_VISIBLE_DEVICES 是无效的。

解决方法

1. 子进程方法

通过在子进程中运行深度学习代码，每次子进程启动时重新设置 CUDA_VISIBLE_DEVICES 环境变量。在这种方法中，每次循环都会启动一个新的子进程 your_script.py，并为该子进程设置不同的 GPU。

import os
import subprocess

gpus = ['0', '1', '2']

for gpu in gpus:
    env = os.environ.copy()
    env['CUDA_VISIBLE_DEVICES'] = gpu
    subprocess.run(['python', 'your_script.py'], env=env)

2. 使用PyTorch的设备管理 (个人推荐)

在 PyTorch 中，你可以直接通过 torch.cuda.set_device(device) 动态设置设备，而无需更改 CUDA_VISIBLE_DEVICES。这种方法在单个进程中动态设置不同的 GPU。示例如下：

import torch

gpus = [0, 1, 2]

for gpu in gpus:
    torch.cuda.set_device(gpu)
    print(torch.cuda.get_device_name(torch.cuda.current_device()))
    # 你的深度学习代码
    model = your_model.to(torch.device(f'cuda:{gpu}'))
    # 继续进行训练或推理

3. 使用多进程方法

使用 multiprocessing 模块启动多个进程，每个进程设置不同的 GPU。这种方法启动多个独立进程，每个进程可以独立设置 CUDA_VISIBLE_DEVICES 并在其上运行代码。：

import os
from multiprocessing import Process

def run_on_gpu(gpu):
    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu)
    # 你的深度学习代码
    # 如训练、推理等

gpus = [0, 1, 2]
processes = []

for gpu in gpus:
    p = Process(target=run_on_gpu, args=(gpu,))
    p.start()
    processes.append(p)

for p in processes:
    p.join()