pytorch和GPU有关操作（CUDA）

使用GPU前，我们首先需要查看GPU信息，这可以通过如下命令实现：

watch nvidia-smi

输出如下：

可以看到总共有两块3090显卡。

一、设定计算设备

默认情况下Pytorch将数据创建在内存，然后利用CPU进行计算，所以我们我们需要手动设定GPU信息。接下来介绍几个相关指令

查看GPU是否可用：torch.cuda.is_available()

查看GPU数量：torch.cuda.device_count()

查看当前GPU索引号（从0开始）：torch.cuda.current_device()

根据索引号查看GPU名字：torch.cuda.get_device_name(index_number:int)

二、TENSOR的GPU计算

默认情况下，Tensor会被存储在内存上。因此，我们打印Tensor时候看不到GPU相关标识：

x=torch.tensor([1,2,3])
print(x)

输出：

tensor([1,2,3])

使用.cuda()可以内存中的Tensor转换到GPU上。如果有多块GPU，可以使用.cuda(i)来表示第i块GPU所对应的显存（从0开始），注意cuda(0)和cuda()等价：

x=x.cuda(0)
print(x)

输出：

tensor([1, 2, 3], device='cuda:0')

我们可以通过Tensor的device属性来查看该Tensor所在的设备：

print(x.device)

输出：

cuda:0

我们可以在创建的时候就指定设备：

device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x=torch.tensor([1,2,3],device=device)
# 或者通过如下的方式指定
x=torch.tensor([1,2,3]).to(device)

print(x)

输出：

tensor([1, 2, 3], device='cuda:0')

如果对GPU上的数据进行运算，那么结果还是存放在GPU上：

y=x**2
print(y)

输出：

tensor([1, 4, 9], device='cuda:0')

注意存储在CPU上的数据无法直接与存放在GPU上的数据进行运算，会报错。

三、模型的GPU计算

和Tensor类似，Pytorch模型也可以通过.cuda转换到GPU上，我们可以通过检查模型参数的device属性来查看存放模型的设备：

net=nn.Linear(3,1)
print(list(net.parameters())[0].device)

输出：

cpu

此时可见模型在CPU上，我们将其转换到GPU上：

net.cuda()
print(list(net.parameters())[0].device)

输出：

cuda:0

注意我们需要保证模型输入的Tensor和模型都在同一台设备上，否则会报错。