查看显卡是否有效、显卡个数
import torch
torch.cuda.is_available() #是否有效
torch.cuda.device_count() #个数
修改默认使用显卡(主卡)
1. 修改加载显卡顺序
一般情况下是默认第0张显卡为主卡,如果你想设置其他显卡为主卡的话,可以使用
os.environ["CUDA_VISIBLE_DEVICES"]
比如有三张显卡,设置1号显卡为主卡:
os.environ["CUDA_VISIBLE_DEVICES"] = "1,2,0"
在显卡上加载及运算时,就会按照1->2->0的顺序加载。
2. 修改可使用显卡数量
在第1点中你查看可使用显卡个数的话,显卡个数为3。可是如果你想禁用掉其他显卡,只使用第1张,可以设置:
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
运行torch.cuda.device_count() 结果为1。
注意:这种情况下,如果你要使用torch.device()来设置数据载入的显卡,这个1号显卡的名称为“cuda:0”,而不是“cuda:1”,因为它是可使用显卡中的第0位,如果设置成“cuda:1”就会发生“RuntimeError: CUDA error: invalid device ordinal”错误。