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'))
输出
<torch.cuda.device object at 0x000001E1A09CBC70>
③查看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