如何在A100图形处理器上使用Pytorch (+ cuda)？

Question

我试图在A100图形处理器上使用我当前的代码，但是我得到了这个错误：

---> backend='nccl'
/home/miranda9/miniconda3/envs/metalearningpy1.7.1c10.2/lib/python3.8/site-packages/torch/cuda/__init__.py:104: UserWarning: 
A100-SXM4-40GB with CUDA capability sm_80 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 A100-SXM4-40GB GPU with PyTorch, please check the instructions at https://pytorch.org/get-started/locally/

这是令人困惑的，因为它指向通常的pytorch安装，但没有告诉我使用哪个pytorch版本+ cuda版本的组合用于我的特定硬件(A100)。为A100安装pytorch的正确方法是什么？

以下是我尝试过的一些版本：

# conda install -y pytorch==1.8.0 torchvision cudatoolkit=10.2 -c pytorch
# conda install -y pytorch torchvision cudatoolkit=10.2 -c pytorch
#conda install -y pytorch==1.7.1 torchvision torchaudio cudatoolkit=10.2 -c pytorch -c conda-forge
# conda install -y pytorch==1.6.0 torchvision cudatoolkit=10.2 -c pytorch
#conda install -y pytorch==1.7.1 torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge

# conda install pytorch torchvision torchaudio cudatoolkit=11.0 -c pytorch
# conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge
# conda install -y pytorch torchvision cudatoolkit=9.2 -c pytorch # For Nano, CC
# conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge

请注意，这可能是微妙的，因为我在过去的机器+ pytorch版本中遇到过这个错误：

How to solve the famous unhandled cuda error, NCCL version 2.7.8 error?

Answer 1

在@SimonB答案中的链接pytorch site中，我这样做了：

pip3 install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html

这为我解决了这个问题。

Answer 2

我有一个A100，并且已经成功地使用

conda install pytorch torchvision torchaudio cudatoolkit=11.1 -c pytorch -c nvidia

现在也建议在pytorch site上使用该选项

Answer 3

我也有同样的问题。您需要安装CUDA 11.0而不是10.2，并为此CUDA版本重新安装PyTorch。