如果电脑的显卡是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())
未完待续......