转载请注明出处：小锋学长生活大爆炸[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()
