1、準備條件：

• Anaconda3
• ubuntu18（其他版本應該也行，我電腦是ubuntu18）
• 已經(jīng)安裝了顯卡驅(qū)動，如果沒有，就用下面的代碼安裝

#安裝顯卡驅(qū)動 ubuntu-drivers devices #查看有哪些可用的顯卡驅(qū)動#方式一（推薦） sudo ubuntu-drivers autoinstall #安裝所有推薦驅(qū)動 #方式二：指定安裝版本 sudo apt install nvidia-440 #這里安裝440,自己看有哪些可用的驅(qū)動

安裝完顯卡驅(qū)動后記得重啟電腦

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

結(jié)果返回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')

結(jié)果：device='cuda:0'，說明成功了

總結(jié)

以上是生活随笔為你收集整理的ubuntu18安装pytorch-gpu，一行conda命令安装，不需要单独安装cuda的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。