Pytorch学习笔记——Tensor在GPU&CPU上操作

对复杂的神经网络和大规模的数据来说，我们一般将数据放在GPU上进行计算，下面介绍数据在GPU及CPU中相互切换的相关操作

1 查看张量所在设备并将其复制到CPU/GPU
.device .cpu() .cuda()

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

x = x.cuda(0) # 当设备GPU数>1时，（）中需要指定第几个GPU
print(x.device)

y = x.cpu()
print(y.device)
--------------
cpu
cuda:0
cpu

2 创建张量时放在GPU上

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

x = torch.tensor([1, 2, 3], device=device)
# or
# x = torch.tensor([1, 2, 3]).to(device) 另一种方式
x
----------------
tensor([1, 2, 3], device='cuda:0')

3 还可.to(device)

y = x.cpu()
z = x.to(device)  #  .to(device)方式未必需要在创建时
print(z.device)
-----------------
cuda:0