在Ubuntu上安装PyTorch可以分为几个步骤,具体如下:
安装Python和pip:
sudo apt-get update
sudo apt-get install python3 python3-pip
安装CUDA和cuDNN(如果使用GPU版本):
创建虚拟环境(推荐):
python3 -m venv pytorch_env
source pytorch_env/bin/activate
安装CPU版本:
pip3 install torch torchvision
安装GPU版本:
pip3 install torch torchvision torchaudio cudatoolkit=<cuda_version> -f https://download.pytorch.org/whl/cu<cuda_version>/torch_stable.html
将<cuda_version>
替换为你安装的CUDA版本号,例如11.4
。
在Python交互环境中导入PyTorch并验证安装是否成功:
python3
>>> import torch
>>> print(torch.__version__)
>>> torch.cuda.is_available()
如果torch.cuda.is_available()
返回True
,则表示CUDA可用,你正在使用GPU版本的PyTorch。