在Ubuntu上配置PyTorch的GPU加速,可按以下步骤进行:
- 安装NVIDIA驱动:添加NVIDIA官方PPA,更新软件包列表后安装适合GPU型号的驱动,安装完成后重启计算机。
- 安装CUDA Toolkit:从NVIDIA官网下载对应Ubuntu版本的安装包,解压后运行安装程序,安装过程中可去掉驱动安装选项,安装完成后配置环境变量。
- 安装cuDNN库:注册NVIDIA开发者账号,下载与CUDA版本兼容的cuDNN库,解压后将文件复制到CUDA安装目录。
- 安装PyTorch:使用pip或conda安装支持GPU的PyTorch版本,安装时需指定与CUDA版本对应的PyTorch包。
- 验证安装:运行Python代码
import torch
,再使用print(torch.cuda.is_available())
检查GPU是否可用,若输出为True,则配置成功。