pytorch自学笔记——GPU

 1、查看GPU可用的数量

torch.cuda.device_count()

2、两个实用的函数——如果没有GPU返回CPU来执行代码

def try_gpu(i=0):  
    """如果存在，则返回gpu(i)，否则返回cpu()。"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():  
    """返回所有可用的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')]

3、将tensor从CPU上迁移到GPU上（这里就用到了上述两个实用函数）

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

4、查看tensor所在位置是GPU还是CPU

x.device

5、神经网络于GPU

#一般在cpu上把权重初始化
net = nn.Sequential(nn.Linear(3, 1))

#将所有的参数在0号gpu上copy一份
net = net.to(device=try_gpu(0))

必须将神经网络中的参数也移植到GPU中去否则会报错

6、确认模型参数存储在同一个GPU上

net[0].weight.data.device

查看网络是否已经移植到GPU上去了

环境搭建：参考博客（要改对应的版本外）

思路：

1、显卡型号

2、对应的cuda版本

3、cudnn

4、cuda对应的pytorch