Pytorch：使用GPU运算

如果电脑的显卡是NVIDIA的，可以用GPU进行运算。

 

# 当GPU可用时,我们可以运行以下代码
# 我们将使用`torch.device`来将tensor移入和移出GPU
#CUDA是NVIDIA推出的并行计算架构，该架构使GPU能够解决复杂的计算问题。
device = torch.device(('cuda') if torch.cuda.is_available() else 'cpu')#如果没有GPU则用CPU计算
y = torch.ones_like(x, device=device)  # 直接在GPU上创建tensor
x = x.to(device)                       # 或者使用`.to("cuda")`方法
z = x + y
print(z)
print(z.to("cpu", torch.double))       # `.to`也能在移动时改变dtype

代码device = torch.device("cuda")中“cuda”也可以写成“cuda:0”。指定从哪个id的GPU开始，“cuda”默认从id为0的GPU开始。以下为输出

tensor([1.0445], device='cuda:0')
tensor([1.0445], dtype=torch.float64)

可以通过下述代码查看有多少个GPU

print(torch.cuda.device_count()) 

 未完待续......