摘要
由于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,所以出现报错。