在Ubuntu下设置PyTorch GPU加速,可按以下步骤进行:
- 安装NVIDIA驱动:添加NVIDIA官方驱动PPA,更新软件包列表后安装适合GPU型号的驱动,安装完成后重启计算机。
- 安装CUDA Toolkit:从NVIDIA官网下载与GPU和PyTorch兼容的CUDA版本,通过运行安装脚本进行安装,安装过程中可勾选驱动程序(若未安装)。
- 安装cuDNN库:注册NVIDIA开发者账号,下载与CUDA版本兼容的cuDNN,解压后将其文件复制到CUDA安装目录,并设置相应权限。
- 安装PyTorch GPU版:使用pip或conda安装,根据CUDA版本从PyTorch官网获取对应安装命令,如
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
(CUDA 11.7为例)。
- 验证安装:运行Python代码
import torch
,再执行print(torch.cuda.is_available())
,若输出为True
,则表示配置成功。