深度学习——使用GPU

1.查看GPU版本的内容

nvidia-smi

2.计算设备

①查看CPU

import torch
from torch import nn

print(torch.device('cpu'))

输出：cpu

②查看GPU，cuda代表GPU

print(torch.cuda.device('cuda'))

输出

③查看GPU的数量

print(torch.cuda.device_count())

3.训练的时候选择CPU还是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')

查看GPU的列表

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()：返回可用的GPU

②try_gpu(10)：没有10块GPU所以返回CPU

①try_all_gpus()：返回所有GPU的列表

4.查询张量的所在设备

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

输出：cpu

5.①张量存储在GPU

# 存储在GPU
x = torch.ones(2, 3, device=try_gpu())
Y = torch.rand(2, 3, device=try_gpu())

②在GPU计算X+Y，保证X和Y要在同一块GPU不然报错

print(X+Y)

tensor([[1.8671, 1.9548, 1.6110],
        [1.1859, 1.3689, 1.8437]], device='cuda:0')

6.①用GPU训练神经网络:模型后面调用.to(device)

net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
print(net(X))

② 确认模型参数在GPU上

print(net[0].weight.data.device)

输出：cuda:0