简单总结一下 YOLO 的核心代码

想用 YOLO 做点事情,但是不敢碰纯 C 的原版 YOLO 代码。网上 Tensoeflow 实现的 YOLO 倒是有几份,但要么是不能用,要么是代码写得太魔幻看不懂改不动。于是作死自己造了个轮子,深入理解一下 YOLO。

YOLO 大致可以分为特征提取网络和检测网络。

特征提取网络

这部分不用多说,唯一一点是它加入了一个 reorg 的层,在缩小特征图尺寸的同时增加通道,“几乎不丢信息”。

其实这部分可以使用其他网络来代替的。通用的网络后面再接几层卷积,进行特征变换。

检测网络

拿到的特征图尺寸会比较小(例如 13x13),每一个 “像素点” 都是一个 “格子”(所以,13x13x1024 的特征图就是 13x13 个格子,每个格子维度是 1024 维)。

reshape

我们希望每个格子产生多个(例如,5 个)预测框,以供我们进行挑选。每个预测框都有 x y w h confidence [class] 这 6 个属性(注意顺序),其中 xywh 是一个预测框的坐标,confidence 是 “我们有多大的把握它准确框住了一个物体”。最后的 [class] 是说,这个被框住的物体属于各个类(比如,总共有 7 个类)的概率。所以一个预测框的数据长度是 4+1+len(class),由于产生了多个预测框,所以需要的滤波器的长度为 box_per_grid * (4 + 1 + len(class))。为了方便使用,我们作一下 reshape:

获取检测结果

YOLO 使用了 Anchor,可以理解为 “默认大小”,或者干脆理解为 “长宽比” 都行。在获取结果的时候,直接把这个乘上去就好了。

但是,这里的 xy 是 相对于格子的偏移,而不是 在整张图片中的位置。这个好办,弄一个 meshgrid 就好了:

然后根据 confidence 和 class 的预测过滤一下就好了。

Loss

做 Loss 相对比较复杂一点,但也复杂不到哪儿去。

我们可以知道每张图的 bbox 和对应的 labelbbox 是 YOLO 自己定义的格式,“一个检测框的中心相对于整张图的位置,边长相对于整张图的比例”,但现在网络输出需要的是 相对于格子。所以,在特征图那个层面直接做这个转换就好了:

我们拿到的标注是一个 “平铺直叙” 的标注,这里需要将它们转回到和特征图一样的维度的标注。

这里还能顺手做一个 “这个格子里面有没有物体” 的 mask

好了,准备工作基本完成了,下面可以开始算 loss 了:

找到 IoU 最大的 box,只用它来做计算:

下面就是做减法、加法,计算总 Loss 了:

大致就这样了

发表评论

电子邮件地址不会被公开。 必填项已用*标注