CUDA 与 PyTorch 安装

1. CUDA 安装

参考 https://blog.csdn.net/weixin_43929198/article/details/142524298 。

# 从 conda 安装 CUDA ，包括 cuda-toolkit 和 nvcc 等
# 版本可根据需求替换，参考 https://developer.nvidia.com/cuda-toolkit-archive
conda install nvidia/label/cuda-12.1.0::cuda
# 设置 CUDA_HOME
export CUDA_HOME=/path/to/miniconda3/envs/cuda_env/targets/x86_64-linux

2. CUDA 与 PyTorch 手动安装

参考 https://pytorch.org/get-started/previous-versions/ 。

# conda 安装（ CUDA 12.1 ）
conda install pytorch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 pytorch-cuda=12.1 -c pytorch -c nvidia

# pip 安装（ CUDA 12.1 ）
pip install torch==2.5.0 torchvision==0.20.0 torchaudio==2.5.0 --index-url https://download.pytorch.org/whl/cu121

注意 PyTorch 版本与 CUDA 版本的对应。安装完毕后在 Python 中运行如下代码确认安装成功。

import torch
# CUDA 是否可用
torch.cuda.is_available()
# GPU 数
torch.cuda.device_count()

3. CUDA 与驱动版本的对应

CUDA 版本不能高于驱动所支持的最高版本（ CUDA 与驱动版本对应 ）。

4. CUDA 与 Compute Capability 的对应

部分型号 GPU 的 Compute Capability 如下（ https://developer.nvidia.com/cuda-gpus ）。

GPU	Compute Capability
NVIDIA H100	9.0
NVIDIA A100	8.0
NVIDIA V100	7.0
GeForce RTX 3090 Ti	8.6
GeForce RTX 3090	8.6

CUDA 版本与 Compute Capability 需对应，参考 https://developer.nvidia.com/cuda-toolkit-archive ，也可安装 PyTorch 后进行验证。

import torch
torch.zeros(1).cuda()

若输出类似下方警告，则说明 CUDA 版本与 GPU 的 Compute Capability 不对应，需要重新安装。本例子为 RTX 3090 显卡安装 CUDA 版本为 10.2 的 PyTorch 库时的警告。注意这个问题不能通过 torch.cuda.is_available() 判断，因为版本不对应最常体现为计算速度偏慢，或 PyTorch 某些代码执行时异常卡住，但可能不会报错。

.../miniconda3/envs/new/lib/python3.9/site-packages/torch/cuda/__init__.py:146: UserWarning:
NVIDIA GeForce RTX 3090 with CUDA capability sm_86 is not compatible with the current PyTorch installation.
The current PyTorch install supports CUDA capabilities sm_37 sm_50 sm_60 sm_61 sm_70 sm_75 compute_37.
If you want to use the NVIDIA GeForce RTX 3090 GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

5. 关于 PyTorch 的 TensorFloat-32 性能提示

新的 Nvidia GPU 有 TF32 Tensor Cores ，可大幅加速浮点矩阵乘法和卷积运算效率，但精度略有损失。在 PyTorch 1.7 至 PyTorch 1.11 中，该选项默认开启，但在 PyTorch 1.12 及以后，该选项默认关闭，需要手动开启（ https://pytorch.org/docs/main/notes/cuda.html#tf32-on-ampere ）。

# 是否允许矩阵乘法使用 TF32
torch.backends.cuda.matmul.allow_tf32 = True
# 是否允许 CUDNN 使用 TF32
torch.backends.cudnn.allow_tf32 = True