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