1. CUDA 与 Pytorch 自动安装
一般而言, Pip 会自动安装 Pytorch 所需的依赖库(如果有 GPU 会自动安装 CUDA )。
pip install torch
但需注意,安装的 CUDA 仅包含运行时所需的动态链接库等,有时一些包的安装需要 CUDA 编译器,此时需要安装 cudatoolkit-dev 。
conda install cudatoolkit-dev
2. CUDA 与 Pytorch 手动安装
参考 https://pytorch.org/get-started/previous-versions/ 。
# CUDA 11.3 conda 安装
# 若使用国内源需要去除末尾的 -c pytorch
conda install pytorch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 cudatoolkit=11.3 -c pytorch
# CUDA 11.1 pip 安装
pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu111/torch_stable.html
注意 Pytorch 版本与 CUDA 版本的对应。安装完毕后在 Python 中运行如下代码确认安装成功。
import torch
# CUDA 是否可用
torch.cuda.is_available()
# GPU 数
torch.cuda.device_count()
3. CUDA 与驱动版本的对应
CUDA 版本不能高于驱动所支持的最高版本, CUDA 版本与驱动版本的对应关系可在官网查询。
https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#title-resolved-issues
4. CUDA 与 Compute Capability 的对应
部分型号 GPU 的 Compute Capability 如下(参考自 https://developer.nvidia.com/cuda-gpus )。
GPU | Compute Capability |
---|---|
NVIDIA A100 | 8.0 |
NVIDIA V100 | 7.0 |
Tesla P100 | 6.0 |
GeForce RTX 3090 Ti | 8.6 |
GeForce RTX 3090 | 8.6 |
GeForce RTX 3080 Ti | 8.6 |
GeForce RTX 3080 | 8.6 |
NVIDIA TITAN RTX | 7.5 |
Geforce RTX 2080 Ti | 7.5 |
Geforce RTX 2080 | 7.5 |
GeForce GTX 1080 Ti | 6.1 |
GeForce GTX 1080 | 6.1 |
CUDA 版本与 Compute Capability 对应如下(参考自 https://zh.wikipedia.org/zh-sg/CUDA )。
CUDA 版本 | 支持的计算能力 |
---|---|
9.0 - 9.2 | 3.0 - 7.2 |
10.0 - 10.2 | 3.0 - 7.5 |
11.0 - | 3.5 - 8.6 |
安装 Pytorch 后可运行如下 Python 代码进行验证。
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 及以后,该选项默认关闭,需要手动开启。开启代码如下。
# 是否允许矩阵乘法使用 TF32
torch.backends.cuda.matmul.allow_tf32 = True
# 是否允许 cuDNN 使用 TF32
torch.backends.cudnn.allow_tf32 = True