Pytorch安装-2

安装流程

Ubuntu

NVIDA

测试驱动有无安装成功。

安装cuda，测试有无成功。
将cuda添加到环境变量

安装cudnn

Pytorch安装细节

我安装的时CUDA9.2

检测pytorch是否能用GPU

import torch
a = torch.cuda.is_available()
print(a)

ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())