PyTorch在Ubuntu上支持GPU加速吗

PyTorch在Ubuntu上支持GPU加速
PyTorch作为主流深度学习框架，天然支持GPU加速（尤其是NVIDIA显卡），借助CUDA和cuDNN库可显著提升计算效率。在Ubuntu系统上，通过正确安装驱动、CUDA Toolkit、cuDNN及对应PyTorch版本，即可实现GPU加速功能。

实现GPU加速的前提条件

1. NVIDIA GPU硬件：需配备支持CUDA的NVIDIA显卡（如GeForce GTX系列、Tesla系列、RTX系列等），可通过lspci | grep -i nvidia命令查看显卡型号。
2. NVIDIA显卡驱动：安装与显卡型号匹配的官方驱动（如GTX 1050 Ti推荐驱动版本470），确保驱动能被系统识别。
3. CUDA Toolkit：安装与PyTorch版本兼容的CUDA Toolkit（如PyTorch 2.0+通常支持CUDA 11.7/11.8），CUDA是连接PyTorch与GPU的桥梁。
4. cuDNN库：安装与CUDA版本匹配的cuDNN（如CUDA 11.7对应cuDNN 8.4+），用于优化深度神经网络的GPU计算。

在Ubuntu上配置PyTorch GPU环境的步骤

1. 安装NVIDIA驱动：
   添加NVIDIA官方驱动PPA并更新软件源，然后安装对应驱动版本（以驱动470为例）：

   sudo add-apt-repository ppa:graphics-drivers/ppa
   sudo apt update
   sudo apt install nvidia-driver-470
   sudo reboot  # 重启使驱动生效
   

   安装完成后，可通过nvidia-smi命令验证驱动是否正常（显示显卡型号及驱动版本）。

2. 安装CUDA Toolkit：
   从NVIDIA官网下载对应Ubuntu版本的CUDA Toolkit（如CUDA 11.7），运行本地安装脚本并选择“安装CUDA Toolkit”及“驱动程序”（若驱动未提前安装）：

   wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run
   sudo sh cuda_11.7.0_515.43.04_linux.run
   

   安装完成后，将CUDA添加至环境变量（编辑~/.bashrc）：

   echo 'export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}' >> ~/.bashrc
   echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
   source ~/.bashrc
   

   通过nvcc --version验证CUDA是否安装成功。

3. 安装cuDNN库：
   注册NVIDIA开发者账号并登录，下载与CUDA版本匹配的cuDNN（如CUDA 11.7对应cuDNN 8.4.1），解压后复制文件至CUDA目录：

   tar -xzvf cudnn-11.7-linux-x64-v8.4.1.50.tgz
   sudo cp cuda/include/cudnn*.h /usr/local/cuda/include
   sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
   sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
   

   cuDNN安装无需额外验证，正确复制文件即可。

4. 安装PyTorch GPU版本：
   推荐通过conda或pip安装（以conda为例，CUDA 11.7版本）：

   conda create -n pytorch_gpu_env python=3.8  # 创建虚拟环境（可选但推荐）
   conda activate pytorch_gpu_env
   conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch
   

   或使用pip安装（需指定CUDA版本对应的whl文件，如CUDA 11.7）：

   pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
   

   注意：PyTorch版本与CUDA版本需严格匹配（如PyTorch 2.0+支持CUDA 11.7/11.8），避免因版本不兼容导致无法使用GPU。

验证GPU加速是否生效

安装完成后，通过以下Python代码验证PyTorch是否能识别GPU：

import torch
print("PyTorch版本:", torch.__version__)
print("CUDA是否可用:", torch.cuda.is_available())  # 若为True则表示成功
print("当前GPU设备索引:", torch.cuda.current_device())
print("GPU设备名称:", torch.cuda.get_device_name(torch.cuda.current_device()))

若输出中torch.cuda.is_available()返回True，且显示正确的GPU名称（如"NVIDIA GeForce RTX 3060"），则说明PyTorch已成功配置GPU加速。