基于YOLOv5训练自定义数据集

目录

• 基于YOLOv5训练自定义数据集
  • • • 1.环境配置
      • 2.数据集制作
      • 3.数据源配置
      • 4.模型训练
      • 5.模型的基本使用

基于YOLOv5训练自定义数据集

版本：YOLOv5，v7.0

1.环境配置

1. 下载YOLOv5源码
2. 更改anaconda下载包使用的源为国内的以及pip下载包使用的源为国内的

conda config --remove-key channels
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.bfsu.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
pip config set global.index-url https://mirrors.ustc.edu.cn/pypi/web/simple

3. 进入源码目录，打开Anaconda Powershell Prompt，创建虚拟环境。

// -n参数 指定虚拟环境名称
conda create -n yolov5 python==3.8.5

4. 激活新建的虚拟环境

conda activate yolov5

5. 安装torch的相关库，根据自身硬件配置选择安装什么版本。我的笔记本显卡为GTX 1050ti
   1. CPU版本的

   conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cpuonly
   

   2. GPU版本的
      1. 在CMD窗口下运行nvidia-smi查看显卡驱动的版本。
      2. 在https://pytorch.org/get-started/previous-versions/安装与显卡驱动兼容的pytorch版本。

       conda install pytorch==1.8.0 torchvision torchaudio cudatoolkit=10.2
      

      3. 验证GPU版本的pytorch是否安装成，在虚拟环境下执行

      python
      import torch
      print(torch.__version__)
      # 输出true表示安装成功
      print(torch.cuda.is_available()
      

6. 安装yolov5源码目录下requirements.txt所依赖的包，注意有些包安装过的可以注释掉。

1. 将requirements.txt文件中的下面两行注释掉
# torch>=1.8.0
# torchvision>=0.9.0
2. 安装依赖
 pip install -r requirements.txt

7. 测试：使用预训练好的权重进行测试

# 在yolov5源码目录下新建weights目录，将下载的预训练好的权重放置该目录下
 E:\yolov5> python detect.py --source data/images/bus.jpg --weights weights/yolov5s.pt

2.数据集制作

1. 方式1：使用labelimg，标注文件保存为YOLO格式
   1. 在anaconda虚拟环境下安装数据标注软件labelimg。

   pip install labelimg -i https://mirror.baidu.com/pypi/simple
   

   2. 将labelImg\data\predefined_classes.txt文件删除
   3. 在当前虚拟环境下安装labelimg运行需要的依赖

   conda install pyqt=5
   conda install -c anaconda lxml
   pyrcc5 -o libs/resources.py resources.qrc
   

   3. 启动labelimg：labelimg，然后进行数据标注，标注文件导出为YOLO格式
2. 方式2：进入在线网址https://www.makesense.ai/进行标注，推荐使用

3.数据源配置

1. 将制作好的数据集放置data目录下，数据集按照如下方式存放

├─images
│  ├─train
│  └─val
└─labels
    ├─train
    └─val

2. 修改数据集的配置文件，在data目录下新建xxx.yaml文件，该文件中指定训练集的位置、验证集的位置、数据集的类别数、类别的名称。例如：

# 指定训练集和验证集的位置
train: E:\ultralytics\data\images\train
val: E:\ultralytics\data\images\val
# 数据集的类别数
nc: 1
# 类别的名称
names: ['person']

4.模型训练

1. 在models目录下新建xxx.yaml作为模型的配置文件，例如person.yaml，其内容编辑为models目录下yolov5n.yaml配置文件的内容
2. 训练模型：在预训练好的模型权重的基础上训练

# 使用CPU进行训练
 python train.py --data data/person.yaml --cfg models/person.yaml --weights weights/yolov5n.pt --epoch 100 --batch-size 2 --device cpu
 # 使用单GPU进行训练
  python train.py --data data/person.yaml --cfg models/person.yaml --weights weights/yolov5n.pt --epoch 100 --batch-size 2 --device 0
  
上述命令参数的解释如下：
data参数： 指定存储训练集和验证集的位置的配置文件
cfg参数：指定存储模型的配置参数的配置文件
weights参数：训练时在预训练好的模型上进行微调，这个参数指定预训练模型的位置
epoch参数：训练过程中整个数据集迭代多少次
batch-size：一个epoch中输入多少图片进行权重的更新
device：指定训练时使用的设备，CPU还是GPU

3. 训练过程中中断了，继续训练：比如说需要训练100个epoch的，但是只训练到第50个epoch时候程序中断停止训练了，想从第50个epoch继续训练
   1. 方式1将train.py中的parser.add_argument('--resume', nargs='?', const=True, default=True, help='resume most recent training')的default参数值更改为True
   2. 方式2直接使用命令python .\train.py --resume 上一次终止训练保存的best.pt的路径
4. 训练完原有epoch，在训练完原有epoch的基础上继续训练：比如说已经训练完了200个epoch，但是模型还没有收敛还想使用训练了200轮的权重多训练几百轮
   1. 将train.py中的parser.add_argument('--resume', nargs='?', const=True, default=True, help='resume most recent training')的default参数值更改为True
   2. 将train.py中的parser.add_argument('--epochs', type=int, default=100, help='total training epochs')的default参数的值更改为想要继续训练到的epoch
   3. 将torch_utils.py下的smart_resume函数中的start_epoch变量的值更改为如下：

   # 值为从第几轮继续训练，比如说已经训练了200轮，则这个值为200
   ckpt['epoch'] = 200
   start_epoch = ckpt['epoch']
   

5. 使用tensorboard可视化训练：进入yolo源码根目录下，执行tensorboard --logdir=./runs

5.模型的基本使用

1. 使用detect.py进行检测

# 检测摄像头
python detect.py  --weights runs/train/exp/weights/best.pt --source 0  # webcam
# 检测图片文件
python detect.py  --weights runs/train/exp/weights/best.pt --source file.jpg  # image 
# 检测视频文件
python detect.py --weights runs/train/exp/weights/best.pt --source file.mp4  # video
# 检测一个目录下的文件
python detect.py --weights runs/train/exp/weights/best.pt path/  # directory
# 检测网络视频
python detect.py --weights runs/train/exp/weights/best.pt 'https://youtu.be/NUsoVlDFqZg'  # YouTube video
# 检测流媒体
python detect.py --weights runs/train/exp/weights/best.pt 'rtsp://example.com/media.mp4'  # RTSP, RTMP, HTTP stream