FastestDet训练自己的数据集

 目录

前言

一、数据集构造

二、配置文件设置

1..yaml配置文件介绍

2.开始训练

3.模型评估

4.导出onnx进行测试

总结

---

前言

代码：https://github.com/dog-qiuqiu/FastestDet

算法介绍：https://zhuanlan.zhihu.com/p/536500269

训练环境：Win11

---

一、数据集构造

数据集格式与Darknet Yolo相同，每张图片对应一个txt标签文件。标签格式也是基于Darknet Yolo的数据集标签格式:“category cx cy wh”，其中category为类别下标，cx, cy为归一化标签框中心点的坐标，w, h为归一化标签框的宽度和高度，.txt标签文件内容示例如下:

2 0.2841796875 0.3815104166666667 0.068359375 0.257825
1 0.43505859375 0.55078125 0.0966796875 0.122395833333333

数据集文件结构如下：

 

 

 其中train.txt和val.txt的格式为：

E:\cv_code\FastestDet-main\datasets\images\train\aug_0_10.jpg
E:\cv_code\FastestDet-main\datasets\images\train\aug_0_100.jpg
E:\cv_code\FastestDet-main\datasets\images\train\aug_0_101.jpg

 生成train.txt和val.txt的代码示例：

import os
import glob

# Define the directory containing the dataset
train_dir = r'E:\cv_code\FastestDet-main\datasets\images\train'
val_dir = r'E:\cv_code\FastestDet-main\datasets\images\val'

# Get a list of all file paths in the dataset directory
train_files = glob.glob(os.path.join(train_dir, '*.jpg'))
val_files = glob.glob(os.path.join(val_dir, '*.jpg'))

# Write the file paths to a .txt file
with open('datasets/train.txt', 'w') as file:
    for path in train_files:
        file.write(path + '\n')

with open('datasets/val.txt', 'w') as file:
    for path in val_files:
        file.write(path + '\n')

二、配置文件设置

1..yaml配置文件介绍

以./configs/coco.yaml为例：

DATASET:
  TRAIN: 'E:\cv_code\FastestDet-main\datasets\train.txt'       # train.txt
  VAL: 'E:\cv_code\FastestDet-main\datasets\val.txt'           # val.txt
  NAMES: "configs/coco.names"                                  # label name
MODEL:
  NC: 4                                                        # number of classes
  INPUT_WIDTH: 352                                             # input width
  INPUT_HEIGHT: 352                                            # input height
TRAIN:
  LR: 0.001                                                    # learning rate
  THRESH: 0.25                                                 # threshold
  WARMUP: true                                                 # warmup
  BATCH_SIZE: 96                                               # batch size
  END_EPOCH: 10                                                # train epoch
  MILESTIONES:                                                 # Declining learning rate steps
    - 100
    - 200
    - 250

其中coco.names为你的数据集标签名

person
bicycle
car
motorbike
.......
toothbrush

2.开始训练

预训练模型为：shufflenetv2.pth

指定你的配置文件：

        parser.add_argument('--yaml', type=str, default='configs/test.yaml', help='.yaml config')

开始训练:

python3 train.py

3.模型评估

python3 eval.py --yaml configs/coco.yaml --weight weights/weight_AP05:0.253207_280-epoch.pth

4.导出onnx进行测试

python3 test.py --yaml configs/coco.yaml --weight weights/weight_AP05:0.253207_280-epoch.pth --img data/3.jpg --onnx

---

总结：

以上就是今天要讲的内容，本文仅仅简单介绍了FastestDet的使用，伙伴们快去实践吧。