目标
本文解决yolov5快速运行环境搭建和测试运行
yolov5介绍
https://github.com/ultralytics/YOLOv5
yolo5 docker
- 必须是要linux环境，不推荐windows
- 习惯用ubuntu，注意支持的版本最高是20.04。
推荐用rufus制作U盘启动，选择ubuntu 20.04光盘。部分笔记本安装ubuntu 20.04，网卡驱动无法识别，需要更新内核。下载新版内核保存，用sudo dpkg -i * 安装后重启就可以了。 - 安装Nvidia驱动,使用软件和更新->附加驱动，选择高版本驱动即可。不需要安装
- docker和nvidia-docker安装
使用带有 cuda 环境的 docker 容器，需要安装 nvidia-docker 组件。
sudo apt-get install curl
curl -fsSL get.docker.com -o get-docker.sh
sudo sh get-docker.sh --mirror Aliyun
sudo systemctl enable docker
sudo systemctl start docker
sudo usermod -aG docker $USER
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
&& curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \
&& curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
sudo apt-get update
sudo apt-get install nvidia-docker2
sudo systemctl restart docker
- 测试CUDA是否能够运行
sudo docker run --rm --gpus all nvidia/cuda:11.0-base nvidia-smi
- docker yolov5
从yolov5官网下载最新代码 https://github.com/ultralytics/YOLOv5
进入yolov5代码文件夹，备份Dockerfile文件，创建空的Dockerfile文件
touch Dockerfile
在Dockerfile中复制如下内容
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
# Start FROM Nvidia PyTorch image https://ngc.nvidia.com/catalog/containers/nvidia:pytorch
FROM nvcr.io/nvidia/pytorch:21.10-py3
# Install linux packages
RUN sed -i s@/archive.ubuntu.com/@/mirrors.aliyun.com/@g /etc/apt/sources.list
RUN apt update && apt install -y zip htop screen libgl1-mesa-glx
# Create working directory
WORKDIR /workspace
COPY ./ /workspace
# Install python dependencies
#RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && python -m pip install --upgrade pip Pillow torch torchtext torchvision numpy
RUN pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && python -m pip install --upgrade pip Pillow
#RUN pip uninstall -y nvidia-tensorboard nvidia-tensorboard-plugin-dlprof
RUN pip install --no-cache -r requirements.txt coremltools onnx gsutil notebook wandb>=0.12.2
#RUN pip install --no-cache -U torch torchvision numpy Pillow
# RUN pip install --no-cache torch==1.10.0+cu113 torchvision==0.11.1+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
# Downloads to user config dir
ADD https://ultralytics.com/assets/Arial.ttf /root/.config/Ultralytics/
# Set environment variables
# ENV HOME=/usr/src/app
# Usage Examples -------------------------------------------------------------------------------------------------------
# Build and Push
# t=ultralytics/yolov5:latest && sudo docker build -t $t . && sudo docker push $t
# Pull and Run
# t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all $t
# Pull and Run with local directory access
# t=ultralytics/yolov5:latest && sudo docker pull $t && sudo docker run -it --ipc=host --gpus all -v "$(pwd)"/datasets:/usr/src/datasets $t
# Kill all
# sudo docker kill $(sudo docker ps -q)
# Kill all image-based
# sudo docker kill $(sudo docker ps -qa --filter ancestor=ultralytics/yolov5:latest)
# Bash into running container
# sudo docker exec -it 5a9b5863d93d bash
# Bash into stopped container
# id=$(sudo docker ps -qa) && sudo docker start $id && sudo docker exec -it $id bash
# Clean up
# docker system prune -a --volumes
# Update Ubuntu drivers
# https://www.maketecheasier.com/install-nvidia-drivers-ubuntu/
# DDP test
# python -m torch.distributed.run --nproc_per_node 2 --master_port 1 train.py --epochs 3
sudo docker build -t yolov5_docker .
sudo docker container run -p 7777:8888 --ipc=host -it --rm --gpus "device=0" --name yolov5 -v /home/yakeworld/yolov5:/workspace -v /home/yakeworld/datasets:/datasets yolov5_docker /bin/bash
本地数据准备
下载口罩数据
链接：https://pan.baidu.com/s/1oP1luvtpR2x07GCT6yfQDA
提取码：ajus
参考coco128数据结构，依样画葫芦即可。
训练模型
python train.py --batch 16 --epochs 100 --data mask.yaml --weights yolov5s.pt
预测模型
python detect.py --weights runs/train/exp/weights/best.pt --source data/images/
结果
