摘要

由于YOLOv5提供的官方镜像里面的pytorch版本是基于cuda 10.2的，英伟达30系显卡不支持该版本的cuda，因此要升级pytorch版本，才能在镜像中使用GPU。

按照YOLO官方文档进行操作

拉取镜像，启动容器并加载镜像

#拉取镜像
#sudo docker pull ultralytics/yolov5:latest
sudo docker pull ultralytics/yolov5:v6.0
#可以指定端口映射进行启动，官方指令如下
# sudo docker run --ipc=host --gpus all -it ultralytics/yolov5:latest
sudo docker run --gpus all  -v /home/lbh/ubuntu-home/datasets:/usr/src/datasets --name yolov5_6.0_gpu -p 16023:23  -it ultralytics/yolov5:v6.0
#指定端口映射、指定GPU模式进行启动

运行推理代码

python detect.py --weights yolov5s.pt --source data/images/bus.jpg --device 0

由于30系显卡普遍不支持CUDA10，所以出现报错。

执行GPU推理时报错

查看pytorch版本

重新安装pytorch

配置国内conda源

conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/
#下面两个非常重要，从清华源下载速度非一般
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/
#查看是否添加成功
conda config --show-sources

卸载pytorch

conda uninstall pytorch torchvision torchaudio cudatoolkit

安装合适版本的pytorch

查找合适的pytorch版本
https://pytorch.org/get-started/locally/

#官方安装脚本
#conda install pytorch torchvision torchaudio cudatoolkit=11.6 -c pytorch -c conda-forge
#官方安装脚本有两个问题：
#1、需要去掉 c pytorch -c conda-forge 才能从清华源下载
#2、版本太新了，目前开源的YOLOv5_v6.0有部分代码不兼容
#可以尝试通过以下方式安装较老的版本
conda install pytorch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1 cudatoolkit=11.3
#删除cpuonly，否则会默认使用cpu版本的pytorch
conda uninstall cpuonly #此项待验证。。。。。

最后通过pycharm连接docker可以进行开发
pycharm连接docker容器镜像进行开发