机器人小学期(4):将YOLO应用到自己的数据集上

YOLO主页上提供了很详细的教程,讲解如何使用作者提供的YOLO实现,其中包括如何使用pretrain的model,如何在VOC和COCO数据集上训练。在这篇文章中,我们来看一下如何能将作者提供的YOLO实现应用到自己的数据集上。

安装YOLO

首先,我们需要跟着主页上的教程安装YOLO。

1
2
3
git clone https://github.com/pjreddie/darknet
cd darknet
make

如果想使用GPU的话,要在make之前把Makefile中前两行置1。

1
2
GPU=1
CUDNN=1

调整我们的数据集

此链接中有我制作的数据集。我们来看一下数据集里都需要什么内容:

  • 训练集和验证集的图片images/
  • 训练集和验证集的label labels/:目录下的每个label文件名和图片对应,把图片文件名的格式.png/jpg/...改成.txt。label文件中,从第一行开始每行表示一个bounding box。每个bounding box包含5个参数,依次是类别(从0开始),x,y,w,h。其中(x,y)是bounding box的中心点位置,w表示bounding box的宽度,h表示bounding box的高度。注意,后四个参数都要除以图片的尺寸以获得[0,1]的数。

    1
    0 0.4140625 0.51953125 0.125 0.17578125
  • 一个文本文件,包含了训练集中所有图片的路径train.txt

    1
    2
    3
    4
    5
    /home/robot/Desktop/PVZ/data/images/1.png
    /home/robot/Desktop/PVZ/data/images/2.png
    /home/robot/Desktop/PVZ/data/images/3.png
    /home/robot/Desktop/PVZ/data/images/4.png
    ......
  • 一个文本文件,包含了验证集中所有图片的路径valid.txt:和train.txt类似。

  • 一个文本文件,包含了所有类别的名字zombie.name:从第一行开始每行写一个名字,第一行为第0类的名字,第二行为第1类的名字,以此类推。比如在我们的数据集中只有一类,这类名字叫Marty,因此只有一行,上面写着Marty。
    1
    Marty

到现在为止,我们的数据集就已经创建好了,接下来我们要做的就是把数据喂进YOLO中进行训练。

使用YOLO训练

接下来,我们来看一下如何训练。由于我在项目中使用的是YOLOv3-tiny这个网络,所以下面都以这个网络为例。

darknet/cfg/中有很多.cfg文件,这个文件描述了网络结构和一些训练参数,因为darknet实现了一套自己的底层框架,所以如果你先构建新的网络,也需要按照这个格式来。

打开darknet/cfg/yolov3-tiny.cfg文件,为了能使用我们自己的数据集,还要对其做一些调整。为了方便起见,我们创建两个.cfg文件,yolov3-tiny.cfg用于训练和yolov3-tina-testing.cfg用于测试。

  • yolov3-tiny.cfg文件中我们将[net]中的batch size设为64,subdivision设为2。在yolov3-tiny.cfg-testing文件中我们将[net]中的batch size设为1,subdivision设为1。
    yolov3-tiny.cfg

    1
    2
    3
    # Testing
    batch=1
    subdivisions=1

    yolov3-tiny-testing.cfg

    1
    2
    3
    # Testing
    batch=64
    subdivisions=2
  • 根据自己的数据集调整[net]中的widthheight,比如在我的数据集中图片是512*512的,所以有

    1
    2
    width=512
    height=512
  • 接下来,我们要把所有yolo层里的classes改成自己的数据集中的类别数。比如我的数据集中只有一类,因此classes=1

  • 最后也是最重要的一步,我们需要对每个yolo层前的卷积层的输出channel数做调整。根据上一篇文章,我们知道输出时每个格子对应的channel为\3*(4+1+num_classes\),所以要把filters改为这个数。比如在我的数据集中,只有一类,那输出channel数为18,因此改成filters=18

在修改完网络后,我们还需要建立一个新文件zombie.data,里面包含了训练时所用的一些目录的信息:

1
2
3
4
5
classes = 1
train = /home/robot/Desktop/PVZ/data/train.txt
valid = /home/robot/Desktop/PVZ/data/valid.txt
names = /home/robot/Desktop/PVZ/data/zombie.name
backup = backup

其中classes表示数据集中的类别数目,train指向上一节提到的train.txt所在位置,valid指向上一节提到的valid.txt所在位置,names指向类别名字的文件zombie.name所在位置,backup表示中间存储的目录,之后训练过程中存储下来的weights文件都可以在这个目录下找到。

现在我们就可以开始训练了,参照主页中的格式运行darknet:

1
./darknet detector train zombie.data yolov3-tiny.cfg darknet.conv.weights

最后一个命令行参数是pretrain的权重文件,关于YOLOv3-tiny的pretrain model可以在此链接处下载。

训练好后,我们在backup的目录下找到weights文件,然后进行测试:

1
./darknet detector test zombie.data yolov3-tiny-testing.cfg backup/[weights file]

然后就可以看到目录下出现了一个predictions.jpg文件,如图:


predictions.jpg
机器人小学期(5):使用PyTorch实现YOLOv3-tiny 机器人小学期(3):YOLOv3

评论

Your browser is out-of-date!

Update your browser to view this website correctly. Update my browser now

×