参考:pytorch学习笔记-UserWarning: CUDA initialization: CUDA unknown error - 简书
跑yolov5时,安装好pytorch之后,运行脚本
import torch
a=torch.cuda.is_available()
print(torch.__version__)
print(a)
出现了错误:
UserWarning: CUDA initialization: CUDA unknown error - this may be due to an incorrectly set up environment,
e.g. changing env variable CUDA_VISIBLE_DEVICES after program start. Setting the available devices to be zero.
(Triggered internally at /opt/conda/conda-bld/pytorch_1623448255797/work/c10/cuda/CUDAFunctions.cpp:115.)
return torch._C._cuda_getDeviceCount() > 0
首先检查显卡驱动,CUDA,cudnn,以及pytorch的版本是否匹配,如果不匹配,需要卸载之后重装对应的版本。
CUDA10.2 Python3.8 pytorch1.8 没有错。
如果版本都是正确的,需要设置一下环境变量,输入sudo vim ~/.bashrc
,在最后加上:
# 前三行是安装CUDA的时候需要设置的
export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export CUDA_HOME=/usr/local/cuda/bin
export CUDA_VISIBLE_DEVICES=0
保存退出,试一下是否能够使用CUDA。
如果还是不行,输入apt-get install nvidia-modprobe
,应该就没有问题了。
我是在安装了apt-get install nvidia-modprobe之后才OK的。