pytorch 中注意cuda版本和gpu算力匹配
在安装torch时,一定要注意显卡的cuda版本问题。
比如,在 RTX2080上 同样的环境中 程序可以正常运行,而换到A100中,就会报错如下:
NVIDIA A100-PCIE-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_70 sm_75. If you want to use the NVIDIA A100-PCIE-40GB GPU with PyTorch, please check the instructions at Start Locally | PyTorch
大概意思就是: NVIDIA A100-PCIE-40GB 带有的CUDA算力是8.0,它和现有的PyTorch版本不匹配,现有的PyTorch版本支持的CUDA算力是 3.7,5.0,6.0,7.0,7.5。
支持的CUDA算力是与安装的cuda的版本有关的,cuda 10.2 仅仅支持 3.7,5.0,6.0,7.0算力,不支持8.0算力。而cuda11是支持8.0算力的。
目前安装的torch版本是1.7.0,所以,需要安装cuda11及其以上,并且和torch 1.7.0不冲突的版本。
进入 PyTorch官网:Previous PyTorch Versions | PyTorch
选择合适的CUDA版本, 也可以去 Previous PyTorch Versions 进行查看选择,
最终选择了 v1.7.1 CUDA 11.0的版本
# CUDA 11.0
pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
问题解决。
参考:https://zhuanlan.zhihu.com/p/427395039
来源:m0_46483236