计算机系统应用教程网站

网站首页 > 技术文章 正文

目标检测算法丨YOLOv3-Darknet53构建特征金字塔

btikc 2024-09-01 15:30:11 技术文章 15 ℃ 0 评论

网络构建

前面我们已经完成了Darknet53的骨干网络构建。

目标检测算法丨YOLOv3介绍及Darknet53代码实现

接下来我们来了解一下如何利用前面提取的特征构建特征金字塔和进行分类预测以及回归预测。

从特征提取到预测结果,总共分为两个部分:

1. 构建FPN特征金字塔,进行加强特征提取

特征图金字塔网络FPN(Feature Pyramid Networks)是2017年提出的一种网络,FPN主要解决的是物体检测中的多尺度问题,通过简单的网络连接改变,在基本不增加原有模型计算量的情况下,大幅度提升了小物体检测的性能。

2. 利用YOLO Head对三组有效特征层进行预测

YOLO Head 本质上是一个3 x 3 卷积加上1 x 1卷积,一个起到特征整合,一个用于调整通道数


构建FPN特征金字塔进行加强特征提取

我们可以通过YOLO的网络结构可以看到,YOLOv3提取多特征层进行检测的,一共提取了三层特征层,位于骨干网络中的不同位置,这个三个特征层的shape分别为(52,52,256)(26,26,512)(13,13,1024)

所以在获得了这个三个特征层之后,我们就可以进行FPN层的构建,构建方式如下:

1. 13 x 13 x 1024的特征层会进行5次的卷积处理,然后利用YOLO Head获得预测结果。一部分经过上采样之后和26 x 26 x512特征层进行结合,

2. 经过结合的特征再进行5次的卷积处理,同样的会利用YOLO Head获得预测结果,然后进行上采样与52 x 52 x256特征层进行结合

3. 然后52 x 52 x256也同样的会进行5次的卷积处理,然后利用YOLO Head获得预测结果。

通过特征金字塔,我们就可以将不同尺寸的特征层进行融合,能够提取出更好的特征。


利用YOLO Head获取预测结果

在经过特征金字塔加强了特征提取之后,我们就可以利用这三种不同尺寸的特征输入到YOLO Head中得到预测结果。

对三个特征层分别进行处理,假设我们现在要对VOC数据进行预测,那么我们的输出层的shape分别为(13,13,75)(26,26,75)(52,52,75),最后一个75是通道数,是基于VOC数据的,该数据集的类别为20种,而YOLOv3针对每一个特征层的每一个特征点存在3个先验框,所以最后预测的通道数为 3 x 25(20个类别 + 4个坐标 + 1个置信度)

使用coco数据集的时候,类别数量为80,所以最后的维度是255


网络代码实现

from collections import OrderedDict
import torch
import torch.nn as nn

from nets.darknet import darknet53
from utils.utils import get_classes, get_anchors

def conv2d(filter_in, filter_out, kernel_size):
    pad = (kernel_size - 1) // 2 if kernel_size else 0
    return nn.Sequential(OrderedDict([
        ("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=1, padding=pad, bias=False)),
        ("bn", nn.BatchNorm2d(filter_out)),
        ("relu", nn.LeakyReLU(0.1)),
    ]))

#------------------------------------------------------------------------#
#   make_last_layers里面一共有七个卷积,前五个用于提取特征。
#   后两个用于获得yolo网络的预测结果
#------------------------------------------------------------------------#
def make_last_layers(filters_list, in_filters, out_filter):
    m = nn.Sequential(
        conv2d(in_filters, filters_list[0], 1),
        conv2d(filters_list[0], filters_list[1], 3),
        conv2d(filters_list[1], filters_list[0], 1),
        conv2d(filters_list[0], filters_list[1], 3),
        conv2d(filters_list[1], filters_list[0], 1),
        conv2d(filters_list[0], filters_list[1], 3),
        nn.Conv2d(filters_list[1], out_filter, kernel_size=1, stride=1, padding=0, bias=True)
    )
    return m

class YoloBody(nn.Module):
    def __init__(self, anchors_mask, num_classes, pretrained = False):
        super(YoloBody, self).__init__()
        #---------------------------------------------------#
        #   生成darknet53的主干模型
        #   获得三个有效特征层,他们的shape分别是:
        #   52,52,256
        #   26,26,512
        #   13,13,1024
        #---------------------------------------------------#
        self.backbone = darknet53()
        if pretrained:
            self.backbone.load_state_dict(torch.load("model_data/darknet53_backbone_weights.pth"))

        #---------------------------------------------------#
        #   out_filters : [64, 128, 256, 512, 1024]
        #---------------------------------------------------#
        out_filters = self.backbone.layers_out_filters

        #------------------------------------------------------------------------#
        #   计算yolo_head的输出通道数,对于voc数据集而言
        #   final_out_filter0 = final_out_filter1 = final_out_filter2 = 75
        #------------------------------------------------------------------------#
        self.last_layer0            = make_last_layers([512, 1024], out_filters[-1], len(anchors_mask[0]) * (num_classes + 5))

        self.last_layer1_conv       = conv2d(512, 256, 1)
        self.last_layer1_upsample   = nn.Upsample(scale_factor=2, mode='nearest')
        self.last_layer1            = make_last_layers([256, 512], out_filters[-2] + 256, len(anchors_mask[1]) * (num_classes + 5))

        self.last_layer2_conv       = conv2d(256, 128, 1)
        self.last_layer2_upsample   = nn.Upsample(scale_factor=2, mode='nearest')
        self.last_layer2            = make_last_layers([128, 256], out_filters[-3] + 128, len(anchors_mask[2]) * (num_classes + 5))

    def forward(self, x):
        #---------------------------------------------------#
        #   获得三个有效特征层,他们的shape分别是:
        #   52,52,256;26,26,512;13,13,1024
        #---------------------------------------------------#
        x2, x1, x0 = self.backbone(x)

        #---------------------------------------------------#
        #   第一个特征层
        #   out0 = (batch_size,255,13,13)
        #---------------------------------------------------#
        # 13,13,1024 -> 13,13,512 -> 13,13,1024 -> 13,13,512 -> 13,13,1024 -> 13,13,512
        out0_branch = self.last_layer0[:5](x0)
        out0        = self.last_layer0[5:](out0_branch)

        # 13,13,512 -> 13,13,256 -> 26,26,256
        x1_in = self.last_layer1_conv(out0_branch)
        x1_in = self.last_layer1_upsample(x1_in)

        # 26,26,256 + 26,26,512 -> 26,26,768
        x1_in = torch.cat([x1_in, x1], 1)
        #---------------------------------------------------#
        #   第二个特征层
        #   out1 = (batch_size,255,26,26)
        #---------------------------------------------------#
        # 26,26,768 -> 26,26,256 -> 26,26,512 -> 26,26,256 -> 26,26,512 -> 26,26,256
        out1_branch = self.last_layer1[:5](x1_in)
        out1        = self.last_layer1[5:](out1_branch)

        # 26,26,256 -> 26,26,128 -> 52,52,128
        x2_in = self.last_layer2_conv(out1_branch)
        x2_in = self.last_layer2_upsample(x2_in)

        # 52,52,128 + 52,52,256 -> 52,52,384
        x2_in = torch.cat([x2_in, x2], 1)
        #---------------------------------------------------#
        #   第一个特征层
        #   out3 = (batch_size,255,52,52)
        #---------------------------------------------------#
        # 52,52,384 -> 52,52,128 -> 52,52,256 -> 52,52,128 -> 52,52,256 -> 52,52,128
        out2 = self.last_layer2(x2_in)
        return out0, out1, out2

from config import *

# 获classes和anchor
class_names,num_classes = get_classes('../model_data/voc_classes.txt')
anchors,num_anchors = get_anchors('../model_data/yolo_anchors.txt')

model = YoloBody(anchors_mask,num_classes,pretrained=pretrained)
print(model)

Tags:

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表