轻松yolo5 自己的训练集

目标

本文解决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/