ubuntu18安装pytorch-gpu，一行conda命令安装，不需要单独安装cuda

1、准备条件：

• Anaconda3
• ubuntu18（其他版本应该也行，我电脑是ubuntu18）
• 已经安装了显卡驱动，如果没有，就用下面的代码安装

#安装显卡驱动
ubuntu-drivers devices  #查看有哪些可用的显卡驱动


#方式一（推荐）
sudo ubuntu-drivers autoinstall  #安装所有推荐驱动  

#方式二：指定安装版本
sudo apt install nvidia-440 #这里安装440,自己看有哪些可用的驱动

安装完显卡驱动后记得重启电脑

2、安装pytorch-gpu

用conda命令

 conda install pytorch-gpu torchvision

这就会自动安装pytorch-gpu和所需的cuda，cudnn。

测试 1

>>> print(torch.cuda.is_avilable())
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'torch.cuda' has no attribute 'is_avilable'
>>> print(torch.cuda.is_available())
True

结果返回True，说明安装gpu版成功

 

测试2，将torch放到gpu上。

>>> torch.rand(3,4).to('cuda')
tensor([[0.2400, 0.9294, 0.4117, 0.0969],
        [0.6468, 0.1129, 0.8767, 0.6132],
        [0.1432, 0.6973, 0.4343, 0.5509]], device='cuda:0')

结果：device='cuda:0'，说明成功了