Ubuntu中PyTorch与CUDA如何协同工作

在Ubuntu中，PyTorch与CUDA的协同工作主要依赖于正确安装和配置两者。以下是详细的步骤和说明：

安装显卡驱动

首先，需要安装NVIDIA的显卡驱动。没有安装驱动的话要先安装驱动，或者驱动版本太低的话要先进行升级。可以使用以下命令来安装推荐的显卡驱动：

sudo ubuntu-drivers autoinstall

或者，如果你想安装指定版本的驱动，可以使用：

sudo apt install nvidia-driver-555

安装CUDA Toolkit

安装完显卡驱动后，接下来安装CUDA Toolkit。这可以通过以下步骤完成：

1. 添加CUDA Toolkit的官方PPA（Personal Package Archive）：

sudo add-apt-repository ppa:graphics-drivers/ppa
sudo apt-get update

2. 安装CUDA Toolkit：

sudo apt install cuda

3. 验证CUDA安装是否成功，可以使用以下命令：

nvcc -V

安装cuDNN

最后，安装cuDNN库。这可以通过以下步骤完成：

1. 下载与CUDA版本兼容的cuDNN库。例如，对于CUDA 11.8，可以从NVIDIA官网下载对应的cuDNN版本。

2. 解压下载的cuDNN文件，并将解压后的文件夹路径添加到LD_LIBRARY_PATH环境变量中：

tar -xzvf cudnn-11.8.1-linux-x64-v8.6.0.70.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*
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

安装PyTorch的GPU版本

最后，安装PyTorch的GPU版本。这可以通过以下命令完成：

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

请将cu118替换为你安装的CUDA版本。

验证PyTorch与CUDA的协同工作

安装完成后，可以通过以下代码验证PyTorch是否能够使用GPU：

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

如果输出为True，则表示PyTorch已经成功配置为使用CUDA。

以上步骤应该能够帮助你在Ubuntu系统中成功配置PyTorch与CUDA的协同工作。如果在安装过程中遇到任何问题，可以参考PyTorch和CUDA的官方文档，或者在相关的技术社区寻求帮助。