Ubuntu上PyTorch如何加速

1. 安装NVIDIA GPU驱动
确保系统安装了适配GPU型号的最新NVIDIA驱动（可通过nvidia-smi命令验证驱动是否识别GPU）。若未安装，可通过以下步骤操作：

• 添加NVIDIA官方PPA：sudo add-apt-repository ppa:graphics-drivers/ppa
• 更新软件包列表：sudo apt update
• 安装推荐驱动（如nvidia-driver-470及以上）：sudo ubuntu-drivers autoinstall
• 重启计算机使驱动生效。

2. 安装CUDA Toolkit
CUDA是PyTorch GPU加速的核心依赖，需匹配PyTorch支持的版本（如CUDA 11.7/11.8）。安装步骤：

• 从NVIDIA官网下载对应Ubuntu版本（如20.04/22.04）的CUDA Toolkit安装包（.run或.deb格式）。
• 若使用.deb包，执行：sudo dpkg -i cuda-repo-ubuntu2004-11-7-local_11.7.0-515.43.04-1_amd64.deb；若使用.run包，执行：sudo sh cuda_11.7.0_515.43.04_linux.run（安装时勾选“CUDA Toolkit”组件）。
• 添加CUDA环境变量至~/.bashrc：

  export PATH=/usr/local/cuda/bin:$PATH  
  export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH  
  

• 生效配置：source ~/.bashrc。

3. 安装cuDNN库
cuDNN是针对深度学习的GPU加速库，需与CUDA版本严格匹配（如cuDNN 8.4对应CUDA 11.7）。安装步骤：

• 注册NVIDIA开发者账号，下载cuDNN安装包（如cudnn-11.7-linux-x64-v8.4.1.50.tgz）。
• 解压并复制文件至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*  
  

4. 安装支持GPU的PyTorch版本
通过PyTorch官网获取与CUDA版本匹配的安装命令（如CUDA 11.7）：

• pip安装：pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
• conda安装：conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch
  安装后，通过以下代码验证GPU可用性：

import torch  
print(torch.cuda.is_available())  # 应输出True  
print(torch.cuda.get_device_name(0))  # 显示GPU型号  

5. 优化代码实现加速

• 将模型与数据移至GPU：

  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
  model = model.to(device)  
  data = data.to(device)  # 确保输入数据也在GPU上  
  

• 使用混合精度训练（AMP）：通过torch.cuda.amp模块减少内存占用、提升计算速度：

  from torch.cuda.amp import autocast, GradScaler  
  scaler = GradScaler()  
  for data, target in dataloader:  
      optimizer.zero_grad()  
      with autocast():  # 自动选择float16/float32计算  
          output = model(data)  
          loss = criterion(output, target)  
      scaler.scale(loss).backward()  # 缩放梯度防止溢出  
      scaler.step(optimizer)  
      scaler.update()  # 更新缩放因子  
  

• 优化数据加载：在DataLoader中设置num_workers（如num_workers=4）启用多进程并行加载，避免数据预处理成为瓶颈。

6. 其他加速技巧

• 分布式训练：若有多个GPU，可使用torch.distributed模块的DistributedDataParallel（DDP）实现多卡并行，显著提升大规模模型训练速度。
• 调整批量大小：适当增大batch_size（如从32增至64）可提高GPU利用率，但需避免超过GPU显存上限（可通过nvidia-smi监控显存使用情况）。
• 使用最新版本：保持PyTorch、CUDA、cuDNN为最新版本，新版本通常包含性能优化（如PyTorch 2.0+对CUDA 11.8的优化）。