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/