ubuntu

Ubuntu下PyTorch GPU加速如何设置

小樊
35
2025-08-28 00:53:25
栏目: 智能运维

在Ubuntu下设置PyTorch GPU加速,可按以下步骤进行:

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

0
看了该问题的人还看了