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会被存储在内存上。因此 我们打印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上的数据进行运算 会报错。

和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和模型都在同一台设备上 否则会报错。