PyTorch、CUDA Toolkit 及顯卡驅動的安裝流程

[JeffKko]

PyTorch、CUDA Toolkit 及顯卡驅動的安裝流程

Intro

安裝的 CUDA 版本必須符合 PyTorch 版本

CUDA 版又必須符合顯卡驅動程式的版本

所以寫個文章記錄一下

指令下

nvidia-smi

可以看到顯卡驅動的版本 (Driver Version)

顯卡驅動版本和 CUDA 對應版的可以看這個表

基本上都是版本高的可向下兼容

https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html

安裝 CUDA

到官網安裝查一下自己電腦適合安裝甚麼版本 CUDA
https://pytorch.org/get-started/locally/
進去後面後自動選的就是了

在到官網下載 CUDA

最新版 https://developer.nvidia.com/cuda-downloads
歷史版本 https://developer.nvidia.com/cuda-toolkit-archive

安裝好後可以下指令確認版本

nvcc -V

可以看到 release ${version}

安裝 PyTorch

CUDA 安裝完成後

回到 Pytron 官網可以看到對應的版本下有安裝指令

例如 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

這個指令會一併安裝 torchvision 和 torchaudio

有人說會安裝失敗可以改 pip3 改成 pip

測試

都安裝好後可以建立一個 test.py

import torch

print(torch.__version__)

tensor = torch.rand(3,4)
print(f"Device tensor is stored on: {tensor.device}")
# Device tensor is stored on: cpu

print(torch.cuda.is_available())
#True

tensor = tensor.to('cuda')
print(f"Device tensor is stored on: {tensor.device}")
# Device tensor is stored on: cuda:0

執行檔案測試

python test.py

2.0.0+cu117
Device tensor is stored on: cpu
True
Device tensor is stored on: cuda:0

看到這個輸出就代表成功了 (這個是 PyTorch 2.0.0 + cuda 11.7版本)

參考

• Linux下PyTorch、CUDA Toolkit 及显卡驱动版本对应关系（附详细安装步骤）
• 解决PyTorch与CUDA版本不匹配的问题