是的,PyTorch在Ubuntu上支持GPU。為了在Ubuntu上使用PyTorch的GPU功能,你需要安裝CUDA和cuDNN,并確保你的NVIDIA顯卡驅動已經正確安裝。以下是詳細步驟:
首先,你需要安裝與你的顯卡型號相匹配的NVIDIA顯卡驅動。你可以通過以下命令查看顯卡型號:
lspci | grep -i nvidia
然后,根據你的顯卡型號選擇合適的驅動版本進行安裝。
接下來,你需要安裝CUDA Toolkit。你可以從NVIDIA官網下載并安裝適合你Ubuntu版本的CUDA Toolkit。
在安裝了CUDA Toolkit之后,你還需要安裝與CUDA版本相匹配的cuDNN庫。你可以從NVIDIA官網下載并安裝cuDNN。
在完成了上述步驟之后,你可以使用pip或conda安裝PyTorch的GPU版本。以下是使用pip安裝PyTorch GPU版本的命令:
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
請將cu113替換為你安裝的CUDA版本號。
安裝完成后,你可以通過以下Python代碼驗證PyTorch是否支持GPU:
import torch
print(torch.__version__)
print("CUDA is available. You can use GPU.") if torch.cuda.is_available() else print("CUDA is not available. Using CPU.")
如果輸出顯示CUDA可用,那么你就成功地在Ubuntu上配置了支持GPU的PyTorch環境。