深度学习模型的计算默认情况下是在cpu上进行的。
我们可以用torch.cuda.device_count()查看设备上GPU的数量：
print(torch.cuda.device_count())
输出：
设备上有一块显卡。
在pytorch中，cpu和gpu分别用torch.device('cpu')和torch.device('cuda')来表示。请注意，cpu设备意味着所有的物理cpu和内存，而gpu只代表一块显卡及其对应的显存。如果有多块gpu，我们使用torch.device(f'cuda:{i}')来表示第i块显卡，i从0开始。
现在我们定义了两个方便的函数， 这两个函数允许我们在不存在所需所有GPU的情况下运行代码。
def try_gpu(i=0): #@save """如果存在，则返回gpu(i)，否则返回cpu()""" if torch.cuda.device_count() >= i + 1: return torch.device(f'cuda:{i}') return torch.device('cpu') def try_all_gpus(): #@save """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]""" devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())] return devices if devices else [torch.device('cpu')] try_gpu(), try_gpu(10), try_all_gpus()
在GPU上计算张量
我们可以使用.device来查看张量所存储的设备：
X = torch.randn((4,16)) print(X.device)
输出：
请注意，我们执行任何计算，都要保证数据在同一块gpu上才可以，否则，系统将不知道在哪块gpu上进行计算，不知道计算结果要存在哪里。
要让张量的计算在gpu上进行，我们可以在创建张量的时候指定其存储在gpu上：
X = torch.randn((4,16),device=torch.device('cuda')) print(X.device)
输出：
类似地，神经网络模型也可以指定设备，我们创建一个神经网络模型：
net = nn.Sequential( nn.Linear(16,8), nn.Linear(8,1) )
可以用以下代码将其存储在gpu上：
net.to(device=torch.device(f'cuda:{0}'))
检查一下：
print(net[0].weight.device)
输出：
这里torch.device(f'cuda:{0}')等价于torch.device('cuda')。即，这个模型和上面的X存储在了同一块gpu上，即0号显卡。那么，net便可以对X进行计算，我们试一下：
print(net(X))
输出：
整体代码如下：
net = nn.Sequential( nn.Linear(16,8), nn.Linear(8,1) ) net.to(device=torch.device(f'cuda:{0}')) print(net[0].weight.device) X = torch.randn((4,16),device=torch.device('cuda')) print(net(X))
在上面的代码中，我们创建了一个具有两个全连接层的神经网络net，并将其放在了0号显卡上，打印确认了第一层参数所在的设备。然后创建了标准正态分布随机数张量X，也放在了0号显卡上，然后在0号显卡上计算了net(X)并输出结果。