在使用单机多卡的时候我们会经常遇到一些GPU的操作，例如查询当前有几块GPU卡，GPU卡名字，如何在torch中使用多卡等问题，下面我们先进行一些基础的学习，这里不再赘述如何安装torch，可以看我之前的笔记。

一些常用的方法

判断cuda是否可用

torch.cuda.is_available()

获取GPU卡的数量

torch.cuda.device_count()

获取当前设备的索引

torch.cuda.current_device()

获取GPU卡的名字

torch.cuda.get_device_name(0)

下面我们来上一些简单的代码查看下

import torch
print(torch.cuda.is_available())
print(torch.cuda.device_count())
print(torch.cuda.current_device())
gpus = [torch.cuda.get_device_name(g_index) for g_index in range(torch.cuda.device_count())]
print(gpus)

返回

我们可以通过os.environ来设置环境变量使用指定的GPU卡为主显卡

os.environ["CUDA_VISIBLE_DEVICES"] = "1,0"

这样我两张卡默认第二张卡为主显卡，在使用torch.nn.DataParallel指定运算显卡时，显卡的对应关系如下

实际显卡编号----->运算显卡编号
1       ----->0（主卡）
0       ----->1

那么上面我们学习了学习使用指定GPU和获取GPU卡信息，当使用显卡时，我们需要怎么操作？

存在多张卡的情况，我们使用所有GPU卡方法如下

net = torch.nn.DataParallel(model)

直接使用torch.nn.DataParallel将你的模型wrap一下即可，这时所有显卡都将会被用上

如果显卡很多，指定使用其中几张显卡使用，那么操作如下

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
#或者
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'

CUDA_VISIBLE_DEVICES 表示当前可以被python环境程序检测到的显卡

关于GPU还有一些其他的函数

查看GPU的容量

torch.cuda.get_device_capability(device)

清空程序占用的GPU资源

torch.cuda.empty_cache()

使用cuda进行检查计算测试代码

print(torch.rand(3,3).cuda())