在Ubuntu下選擇PyTorch版本時,需要考慮以下幾個因素:
首先,你需要確定你的系統上安裝的CUDA版本。這可以通過在終端中運行以下命令來完成:
nvcc --version
這將返回CUDA Toolkit的版本信息。
PyTorch的每個版本都支持特定的CUDA版本。你可以在PyTorch的官方網站上找到詳細的兼容性表格。例如,如果你的CUDA版本是11.4,你可以選擇安裝PyTorch 1.10.0或更高版本。
conda install pytorch torchvision torchaudio cudatoolkit=X.X -c pytorch
將X.X替換為你的CUDA版本號。
pip install /path/to/downloaded/torch-X.X.X+cpu-cp36-cp36m-linux_x86_64.whl
將X.X.X替換為你的PyTorch版本號,cp36替換為你的Python版本號。
安裝完成后,你可以通過以下Python代碼驗證PyTorch是否正確安裝,并且是否能夠使用GPU:
import torch
print("PyTorch version:", torch.__version__)
print("CUDA is available:", torch.cuda.is_available())
如果torch.cuda.is_available()返回True,則表示PyTorch已經正確配置并且可以使用GPU。