在PyTorch中使用CUDA，需要确保安装的PyTorch版本与你的CUDA版本兼容， 且正确安装了匹配GPU的CUDA Toolkit。以下是在PyTorch中使用CUDA的一般步骤：

1. 检查CUDA支持： 首先，确保你的GPU支持CUDA。你可以在官方CUDA支持列表上查找你的GPU型号。或者直接命令行 nvidia-smi 进行查看。

2. 安装CUDA Toolkit： 下载并安装与你的GPU型号匹配的CUDA Toolkit。你可以从NVIDIA官网下载。在安装期间，可以选择安装适用于你的系统的CUDNN库。

3. 安装cuDNN（可选）： cuDNN是NVIDIA的深度神经网络库，可以加速深度学习任务。在CUDNN下载页面下载适用于你的CUDA版本的cuDNN，并按照安装说明进行安装。

4. 安装PyTorch： 选择合适的PyTorch版本并使用pip或conda进行安装。

服务器重装系统后，需要重新安装显卡驱动、cuda及Pytorch等，有些步骤容易忘记，这里记录一下。这里我的服务器配置以及安装版本的情况如下：

• 服务器系统：Ubuntu 22.04
• 显卡驱动版本：nvidia-driver-535
• cuda版本：

1、安装显卡驱动

在终端命令行输入 nvidia-smi 命令，若正常显示显卡的基础信息，则已成功安装nvidia驱动。

nvidia-smi

如果报错，则按照以下步骤安装nvidia驱动，具体步骤如下：

1. 找到合适自己电脑显卡的驱动

ubuntu-drivers devices

会显示你的电脑上可用的nvidia驱动。只需要安装推荐的版本即可（后面有recommend字样)

这里，服务器显示 nvidia-driver-535 是推荐安装的版本。

2. 找到对应版本的显卡驱动

我们直接安装 nvidia-driver-535 即可：

sudo apt-get install nvidia-driver-535

整个安装过程需要花费一小段时间，耐心等待即可。

3. 重启电脑。
4. 再次输入 nvidia-smi 命令，已经可以正常显示显卡的基本信息。

注意：右上角显示 CUDA Version：12.2，这也是这个版本的显卡驱动最高支持的cuda-toolkit版本。

2、安装CUDA

接下来打开nvidia CUDA Toolkit Archive的官方网址：https://developer.nvidia.com/cuda-toolkit-archive

选择想要安装的版本，这里我们安装CUDA Toolkit 12.2.0，如下图所示选择：

根据我们的操作系统及系统版本选择，这里我们选择了基于 deb(network) 的方式来安装：

根据页面上生成的安装指令进行操作即可：

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update
sudo apt-get -y install cuda

查看目录是否存在：

cd /usr/local/cuda-12.2/bin

配置环境变量：安装成功后需要配置一下环境变量。这步很关键，如果不配置环境变量，则运行机器学习运算的时候会找不到GPU。

# 1.切换回home目录
cd ~  
# 2.修改.bashrc文件
vim .bashrc     
# 3.进入.bashrc文件后，拉到最下面，然后键盘按E键，进入编辑模式添加：
export PATH="/usr/local/cuda-12.2/bin:$PATH"
export LD_LIBRARY_PATH="/usr/local/cuda-12.2/lib64:$LD_LIBRARY_PATH"  
# 4.按ESC键退出编辑
:wq   # 退出保存
# 5. 应用生效
source .bashrc

查看安装的cuda版本：

nvcc -V

可以看到，已经成功安装了 cuda 12.2。

3、安装cuDNN

cuDNN（CUDA Deep Neural Network Library）是一个针对深度学习应用优化的加速库，它使用CUDA技术加速深度神经网络的训练和推断。这个库提供了一些高效的原语，例如卷积、池化、归一化等，可以帮助开发人员快速实现深度神经网络，并且提升训练和推断的速度。

CUDA和cuDNN是两个相互依赖的软件包，它们之间有以下关系：

• CUDA是必需的：cuDNN需要依赖于CUDA才能够工作。因为cuDNN是使用CUDA技术实现的，所以只有支持CUDA的硬件和系统才能够运行它。在安装cuDNN之前，必须先安装CUDA，并且在环境变量中设置CUDA的路径。
• cuDNN是可选的：尽管cuDNN可以提升深度学习应用的性能和效率，但它并不是必需的。如果不使用cuDNN，仍然可以使用CUDA进行深度学习计算。cuDNN只是一个可选的加速库，可以根据需要选择是否使用它。

cuDNN官网：https://developer.nvidia.com/rdp/cudnn-archive

首先在cuDNN官网找到跟此前安装的cuda版本兼容的版本，点击进入选择对应服务器系统类型的安装包即可。

4、安装pyTorch

pytorch官网：https://pytorch.org/get-started/locally/

先到pytorch官网找到在你的操作系统、包、CUDA版本、语言版本下对应的安装脚本。根据你的实际情况选择Pytorch安装包版本，然后复制页面自动生成的脚本进行安装。

这里，使用conda打开需要安装pyTorch的虚拟环境，这里我们的虚拟环境名称为pytorch。

执行安装命令：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

注意：新版本的pytorch已经自带python版本的cuda了，pytorch自带的cuda和前面装的cuda没有关系。

在 python 环境中检查是否安装完成：

import torch
torch.cuda.current_device()
torch.cuda.device_count()
torch.cuda.get_device_name(0)
torch.cuda.is_available()

5、卸载CUDA

参考官方文档：NVIDIA CUDA Installation Guide for Linux

参考资料

• 在PyTorch中使用CUDA, pytorch与cuda不同版本对应安装指南，查看CUDA版本，安装对应版本pytorch