【PyTorch】计算设备

文章目录

  • 1. 介绍
  • 2. 查询和使用

1. 介绍

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

    torch.device('cpu') 
    
  • GPU设备只代表一个GPU和相应的显存

    torch.device('cuda')
    

    如果有多个GPU,我们使用以下方式表示第 i i i块GPU(从0开始)

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

    另外,cuda:0cuda是等价的,都是指第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'))
    

你可能感兴趣的:(深度学习,pytorch)