Conda 环境下 Pytorch 安装避坑指南
1. CUDA 与 Pytorch 安装
参考 https://pytorch.org/get-started/previous-versions/ 。
# 若使用国内源需要去除末尾的 -c pytorch
conda install pytorch==1.10.1 torchvision==0.11.2 torchaudio==0.10.1 cudatoolkit=11.3 -c pytorch注意 Pytorch 版本与 CUDA 版本的对应。安装完毕后在 Python 中运行如下代码确认安装成功。
import torch
# CUDA 是否可用
torch.cuda.is_available()
# GPU 数
torch.cuda.device_count()2. CUDA 与驱动版本的对应
CUDA 版本不能高于驱动所支持的最高版本， CUDA 版本与驱动版本的对应关系可在官网查询。
https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html#title-resolved-issues
3. 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 |
| GeForce RTX 3070 Ti | 8.6 |
| GeForce RTX 3070 | 8.6 |
| Geforce RTX 3060 Ti | 8.6 |
| Geforce RTX 3060 | 8.6 |
| GeForce GTX 1650 Ti | 7.5 |
| NVIDIA TITAN RTX | 7.5 |
| Geforce RTX 2080 Ti | 7.5 |
| Geforce RTX 2080 | 7.5 |
| Geforce RTX 2070 | 7.5 |
| Geforce RTX 2060 | 7.5 |
| GeForce GTX 1080 Ti | 6.1 |
| GeForce GTX 1080 | 6.1 |
| GeForce GTX 1070 Ti | 6.1 |
| GeForce GTX 1070 | 6.1 |
| GeForce GTX 1060 | 6.1 |
| GeForce GTX 1050 | 6.1 |
CUDA 版本与 Compute Capability 对应如下（参考自 https://zh.wikipedia.org/zh-sg/CUDA ）。
| CUDA 版本 | 支持的计算能力 |
|---|---|
| 1.0 | 1.0 - 1.1 |
| 1.1 | 1.0 - 1.1+x |
| 2.0 | 1.0 - 1.1+x |
| 2.1 - 2.3.1 | 1.0 - 1.3 |
| 3.0 - 3.1 | 1.0 - 2.0 |
| 3.2 | 1.0 - 2.1 |
| 4.0 - 4.2 | 1.0 - 2.1+x |
| 5.0 - 5.5 | 1.0 - 3.5 |
| 6.0 | 1.0 - 3.5 |
| 6.5 | 1.1 - 5.x |
| 7.0 - 7.5 | 2.0 - 5.x |
| 8.0 | 2.0 - 6.x |
| 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/