Ubuntu PyTorch如何配置CUDA

在Ubuntu上配置PyTorch以使用CUDA可以显著提高深度学习任务的性能。以下是详细的步骤指南：

1. 安装前的准备工作

• 更新系统：

  sudo apt update
  sudo apt upgrade
  

• 安装NVIDIA显卡驱动： 确保你的显卡支持CUDA，并安装相应的驱动程序。可以通过以下命令安装：

  sudo apt install nvidia-driver-xxx
  

  其中xxx是你的显卡支持的驱动版本。

• 检查CUDA版本： 使用以下命令来检查CUDA版本：

  nvcc --version
  

• 安装必要的依赖：

  sudo apt install build-essential cmake git pkg-config libgtk-3-dev libavcodec-dev libavformat-dev libswscale-dev python3-dev python3-numpy libtbb2 libtbb-dev libdc1394-22-dev
  

2. 安装CUDA Toolkit

• 访问NVIDIA CUDA Toolkit下载页面，根据你的显卡型号和系统选择合适的CUDA版本进行下载。

• 例如，对于Ubuntu 20.04，你可以下载CUDA 11.7：

  wget https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda-repo-ubuntu2004-11-7-local_11.7.1-450.51.06-1_amd64.deb
  sudo dpkg -i cuda-repo-ubuntu2004-11-7-local_11.7.1-450.51.06-1_amd64.deb
  sudo apt-key add /var/cuda-repo-ubuntu2004-11-7-local/7fa2af80.pub
  sudo apt update
  sudo apt install cuda-toolkit-11.7
  

3. 安装cuDNN

• 访问NVIDIA cuDNN下载页面，下载与你的CUDA版本兼容的cuDNN库。

• 例如，对于CUDA 11.7，你可以下载cuDNN 8.9.7：

  wget https://developer.download.nvidia.com/compute/redist/cudnn/v8.9.7/local_installers/cudnn-ubuntu2004-v8.9.7-450.51.06-1+cudnn8.9.7-1_amd64.deb
  sudo dpkg -i cudnn-ubuntu2004-v8.9.7-450.51.06-1+cudnn8.9.7-1_amd64.deb
  sudo cp /var/cuda-repo-ubuntu2004-11-7-local/cuda-*-keyring.gpg /usr/share/keyrings/
  sudo apt update
  sudo apt install cudnn
  

4. 安装PyTorch

• 使用Conda安装PyTorch（推荐）：

  • 下载并安装Miniconda：

    wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
    bash Miniconda3-latest-Linux-x86_64.sh
    source ~/.bashrc
    

  • 创建并激活一个新的Conda环境：

    conda create -n pytorch_env python=3.8
    conda activate pytorch_env
    

  • 根据你的CUDA版本，使用以下命令安装PyTorch：

    conda install pytorch torchvision torchaudio cudatoolkit=11.7 -c pytorch -c conda-forge
    

  • 或者，使用pip安装PyTorch：

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

• 验证安装： 打开Python解释器并尝试导入PyTorch：

  import torch
  print(torch.__version__)
  print(torch.cuda.is_available())
  

  如果torch.cuda.is_available()返回True，则表示PyTorch已成功安装并可以使用GPU。