在你的train.py代码中加入:print(torch.cuda.is.available())后运行,如果得到的是false,证明你没有成功的调用cuda来加速你的模型训练。

经过排查,YOLOv8在安装依赖项的时候直接pip install ultralytics,全部集成在了ultralytics里面,但是torch的版本可能是最新的,而你的cuda如果在11.2以下就无法正常调用cuda了。

因此,我们直接借助一些ai工具得知cuda10.2(及以上)可以配适的torch版本,应该是1.10。

1.卸载原有的torch

pip uninstall torch torchvision

然后输入y表示yes

2.下载指定版本的1.10的torch(我这里还使用了清华镜像源进行加速)

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple/ torch==1.10.0 torchvision torchaudio -f https://download.pytorch.org/whl/cu116/torch_stable.html

问题解决!

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐