CUDA查找与选用GPU设备

查看显卡是否有效、显卡个数

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”错误。