如何将 Pytorch (+ cuda) 与 A100 GPU 一起使用？答案

【问题标题】：How does one use Pytorch (+ cuda) with an A100 GPU?如何将 Pytorch (+ cuda) 与 A100 GPU 一起使用？
【发布时间】：2021-07-03 15:36:31
【问题描述】：

我试图将我当前的代码与 A100 gpu 一起使用，但出现此错误：

---> 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?

【问题讨论】：

Pytorch 1.7.0 或更高版本与 CUDA 11.0 或更高版本 should work。或者你可以使用NGC
@RobertCrovella 如果你说的是真的，那么所需的命令是conda install -y pytorch==1.7.1 torchvision torchaudio cudatoolkit=11.0 -c pytorch -c conda-forge，如果有效，将很快尝试。

标签： python machine-learning neural-network pytorch

【解决方案1】：

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

【讨论】：

您是使用 cuda 11.0 还是 pytorch 1.7.x 安装 pytorch 1.8.0 的？
我尝试了 1.8.0 和 1.7.1，都在工作。

【解决方案2】：

我有一个 A100 并且成功了

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

现在在pytorch site 上也推荐

【讨论】：

【解决方案3】：

从@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

这解决了我的问题。

【讨论】：

对我来说，conda 安装不起作用，但是 pip 安装，不知道为什么