李沐《动手学》-使用GPU

1.查看gpu

nvidia-smi

watch -n 1 nvidia-smi#一秒刷新一次

import torch
from torch import nn

torch.device('cpu'), torch.cuda.device('cuda'), torch.cuda.device('cuda:1')

#查询gpu数量
torch.cuda.device_count()

自定义的两个简洁函数

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()

2.张量与gpu

• 张量需要在同一块设备上，否则会报错的。
  RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!
• cuda和cuda:0是等价的

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

X = torch.ones(2,3,device=try_gpu())
Y = torch.randn(2,3,device=try_gpu(1))

Z = X.cuda(1)
Y+Z

#Z已经在cuda:1上了，返回的就是Z
Z.cuda(1) is Z #True

3.神经网络与gpu

• 将网络指定到某块gpu上（其实就是模型参数放到该gpu上）net.to()
• 模型参数和输入数据都放到同一设备上，防止出错（深层原因：避免变量的复制

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

net(X)

net[0].weight.data.device