os.environ[“CUDA_VISIBLE_DEVICES”]
前提:测试设备共有4块显卡
测试:
import os
# os.environ["CUDA_VISIBLE_DEVICES"]设置程序可见的显卡号
# 设置可见显卡为1,2,3号, 0号显卡不可见
os.environ["CUDA_VISIBLE_DEVICES"] = "1, 2, 3"
import torch
x = torch.randn(5)
# 测试
y = x.to("cuda:0")
y = x.to("cuda:1")
y = x.to("cuda:2")
y = x.to("cuda:3")
运行结果, 第8行报错
结论
os.environ[“CUDA_VISIBLE_DEVICES”]设置程序可见的显卡
无论设置可见的显卡号是多少,在程序中可见的显卡号都是从0开始