【PyTorch】计算设备

1. 介绍

• CPU设备意味着所有物理CPU和内存， 这意味着PyTorch的计算将尝试使用所有CPU核心。可以用以下方式表示：

  torch.device('cpu') 
  

• GPU设备只代表一个GPU和相应的显存。

  torch.device('cuda')
  

  如果有多个GPU，我们使用以下方式表示第$i$块GPU（从0开始）

  torch.device(f'cuda:{
         i}') 
  

  另外，cuda:0和cuda是等价的，都是指第1块显卡。

• 深度学习框架要求计算的所有输入数据都在同一设备上，无论是CPU还是GPU。

• 不经意地移动数据可能会显著降低性能，框架通常会使用自己的线程来执行GPU计算，如果我们需要频繁地切换到Python的主线程，那么这就可能会触发全局解释器锁（GIL），导致GPU的计算被阻塞。

2. 查询和使用

• 查询可用GPU的数量

  torch.cuda.device_count()
  

• 处理不存在GPU的情况

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

• 查询GPU列表

  devices = [torch.device(f'cuda:{
         i}') 
  	for i in range(torch.cuda.device_count())]
  devices = if devices else [torch.device('cpu')]
  

• 查询张量所在的设备。 默认情况下，张量是在CPU上创建的。

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

• 创建张量时指定存储设备

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

• 数据复制到其他设备
  • CPU到GPU

  x = torch.tensor([1, 2, 3])
  x.cuda()
  # x.to(torch.device('cuda'))
  

  • GPU到CPU

  x = torch.tensor([1, 2, 3], device=torch.device('cuda'))
  x.cpu()
  # x.to(torch.device('cpu'))
  

  • 多GPU的情况：指定要复制到的GPU的序号

  x = torch.tensor([1, 2, 3])
  x.cuda(i)
  # x.to(torch.device(f'cuda:{i}'))
  

• 模型复制到其他设备：与数据复制到其他设备的方法相同。

  net = nn.Linear(2, 1)
  net.cuda()
  # net.to(torch.device('cuda'))