由于YOLOv5提供的官方镜像里面的pytorch版本是基于cuda 10.2的,英伟达30系显卡不支持该版本的cuda,因此要升级pytorch版本,才能在镜像中使用GPU。
#拉取镜像
#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,所以出现报错。
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
conda uninstall pytorch torchvision torchaudio cudatoolkit
查找合适的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容器镜像进行开发