计算设备

PyTorch可以指定用来存储和计算的设备，如使用内存的CPU或者使用显存的GPU。默认情况下，PyTorch会将数据创建在内存，然后利用CPU来计算。

讲真，到现在为止还没拿自己显卡训练过实际项目

  1. import torch
  2. from torch import nn
  4. print(torch.cuda.device_count())
  5. print(torch.cuda.current_device())
  6. print(torch.cuda.get_device_name(0))
  9. 结果：
  10. 1
  11. 0
  12. GeForce RTX 2060

Tensor的GPU计算

默认情况下，Tensor会被存在内存上。因此，之前我们每次打印Tensor的时候看不到GPU相关标识。使用.cuda()可以将CPU上的Tensor转换（复制）到GPU上。如果有多块GPU，我们用.cuda(i)来表示第 i 块GPU及相应的显存（i从0开始）且cuda(0)和cuda()等价。

  1. x = torch.tensor([1, 2, 3])
  2. x = x.cuda()
  3. print(x)
  5. # 也可以直接指定存储tensor的设备
  6. x = torch.tensor([1, 2, 3], device='cuda')
  7. print(x)
  9. # 通过divice属性查看该Tensor所在的设备
  10. print(x.device)
  13. 结果：
  14. tensor([1, 2, 3], device='cuda:0')
  15. tensor([1, 2, 3], device='cuda:0')
  16. cuda:0

我们还可以使用Tensor.to(device)方法，将Tensor转移到其他设备上：

  1. device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  2. x = torch.tensor([1, 2, 3], device=device)
  3. # or
  4. x = torch.tensor([1, 2, 3]).to(device)
  5. print(x)
  6. y = x**2
  7. print(y)
  10. 结果：
  11. tensor([1, 2, 3], device='cuda:0')
  12. tensor([1, 4, 9], device='cuda:0')

可以发现，如果对在GPU上的数据进行运算，那么结果还是存放在GPU上。而且要注意，存储在不同位置中的数据是不可以直接进行计算的。

模型的GPU运算

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

  1. net = nn.Linear(3, 1)
  2. print(list(net.parameters()))
  3. print(list(net.parameters())[0].device)
  6. 结果：
  7. tensor([[-0.0826,  0.4849,  0.1278]], requires_grad=True), Parameter containing:
  8. tensor([-0.5701], requires_grad=True)]
  9. cpu

同样的，我么需要保证模型输入的Tensor和模型都在同一设备上，否则会报错。转换到GPU上：

  1. net.cuda()
  2. print(list(net.parameters())[0].device)
  3. print(net(torch.rand(3, 3).cuda()))
  6. 结果：
  7. cuda:0