关键点检测(6)——yolov8-neck的搭建
话接上文。之前学习了backbone,这里就学习neck了。我这里还是会以先以yolov8-pose的网络结构为例进行展示,然后再学习其neck层如何搭建。而Neck(颈部,连接部)是一个中间层,用于对来自backbone的特征进行融合,以提升模型地性能。yolov8并没有使用Neck这个概念,但其中架构图中Head中类似PANet功能地部分也可以归为Neck。所以我们这里会主要学习一下FPN和PANet,而下一节主要是说目标检测模型的决策部分,负责产生最终的检测结果。
在学习YOLOv8的neck之前,我这里对其backbone再啰嗦一下。
Backbone是yolov8的主干特征提取网络,输入图像首先会在主干网络里进行特征提取,提取到的特征可以被称为特征层,是输入图像的特征集合。在主干部分,我们获得了三个特征层进行下一步网络的构建,这三个特征层,我们称为有效特征层。之前也搭建了(在上个博客)。
而这节课的重点就是YOLOv8的Neck层。
1,yolov8的yaml配置文件
首先,我们仍然展示一下yolov8-pose.yaml文件。看看其网络的构造:
# Ultralytics YOLO 🚀, AGPL-3.0 license
# YOLOv8-pose keypoints/pose estimation model. For Usage examples see https://docs.ultralytics.com/tasks/pose
# Parameters
nc: 1 # number of classes
kpt_shape: [17, 3] # number of keypoints, number of dims (2 for x,y or 3 for x,y,visible)
scales: # model compound scaling constants, i.e. 'model=yolov8n-pose.yaml' will call yolov8-pose.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024]
s: [0.33, 0.50, 1024]
m: [0.67, 0.75, 768]
l: [1.00, 1.00, 512]
x: [1.00, 1.25, 512]
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [[15, 18, 21], 1, Pose, [nc, kpt_shape]] # Pose(P3, P4, P5)
确实如我们所说,yolov8并没有neck之说,特征融合都集中在head里面。因为我们上一节已经展示了yolov8如何通过yaml文件加载backbone,所以我们这里直接通过yaml里面的head内容进行打印即可。
为了回顾,我先打印整个yolo-pose的网络结构(因为通过yaml内容只打印head会报错。。要么修改一下代码,要么就全打印了。这里为了方便,就全打印了):
我们通过yaml解析的内容,可以看到上面的neck的内容。neck的代码的解析
第10层:[-1, 1, nn.Upsample, [None, 2, 'nearest']] :本层是上采样层。-1代表将上层的输出作为本层的输入。None代表上采样的size(输出尺寸)不指定。2代表scale_factor=2,表示输出的尺寸是输入尺寸的2倍。nearest代表使用的上采样算法为最近邻插值算法。经过这层之后,特征图的长和宽变成原来的两倍,通道数不变,所以最终尺寸为40*40*1024。
第11层:[[-1, 6], 1, Concat, [1]] # cat backbone P4 :本层是concat层,[-1, 6]代表将上层和第6层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是40*40*1024,第6层的输出是40*40*512,最终本层的输出尺寸为40*40*1536。
第12层:[-1, 3, C2f, [512]] :本层是C2f模块。3代表本层重复3次。512代表输出通道数。与Backbone中C2f不同的是,此处的C2f的bottleneck模块的shortcut=False。
第13层:[-1, 1, nn.Upsample, [None, 2, 'nearest']] :本层也是上采样层(参考第10层)。经过这层之后,特征图的长和宽变成原来的两倍,通道数不变,所以最终尺寸为80*80*512。
第14层:[[-1, 4], 1, Concat, [1]] # cat backbone P3 :本层是concat层,[-1, 4]代表将上层和第4层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是80*80*512,第6层的输出是80*80*256,最终本层的输出尺寸为80*80*768。 -
第15层:[-1, 3, C2f, [256]] # 15 (P3/8-small) :本层是C2f模块。3代表本层重复3次。256代表输出通道数。经过这层之后,特征图尺寸变为80*80*256,特征图的长宽已经变成输入图像的1/8。
第16层:[-1, 1, Conv, [256, 3, 2]] :进行卷积操作(256代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为40*40*256(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样)。
第17层:[[-1, 12], 1, Concat, [1]] # cat head P4 :本层是concat层,[-1, 12]代表将上层和第12层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是40*40*256,第12层的输出是40*40*512,最终本层的输出尺寸为40*40*768。
第18层:[-1, 3, C2f, [512]] # 18 (P4/16-medium) :本层是C2f模块。3代表本层重复3次。512代表输出通道数。经过这层之后,特征图尺寸变为40*40*512,特征图的长宽已经变成输入图像的1/16。
第19层:[-1, 1, Conv, [512, 3, 2]] :进行卷积操作(512代表输出通道数,3代表卷积核大小k,2代表stride步长),输出特征图尺寸为20*20*512(卷积的参数都没变,所以都是长宽变成原来的1/2,和之前一样)。
第20层:[[-1, 9], 1, Concat, [1]] # cat head P5 :本层是concat层,[-1, 9]代表将上层和第9层的输出作为本层的输入。[1]代表concat拼接的维度是1。从上面的分析可知,上层的输出尺寸是20*20*512,第9层的输出是20*20*1024,最终本层的输出尺寸为20*20*1536。
第21层:[-1, 3, C2f, [1024]] # 21 (P5/32-large) :本层是C2f模块。3代表本层重复3次。1024代表输出通道数。经过这层之后,特征图尺寸变为20*20*1024,特征图的长宽已经变成输入图像的1/32。
第22层:[[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5) :本层是Detect层,[15, 18, 21]代表将第15、18、21层的输出(分别是80*80*256、40*40*512、20*20*1024)作为本层的输入。nc是数据集的类别数。
2,yolov8的neck的架构图
Neck部分就是负责多尺度特征融合,通过将来自backbone不同阶段的特征图进行融合,增强特征表示能力。具体来说,YOLOv8的Neck部分包括以下组件。
- PAA模块(Probabilistic Anchor Assignment):用于智能的分批锚框,以优化正负样本的选择,提高模型的训练效果。
- PAN模块(Path Aggregation Network):包括两个PAN模块,用于不同层次特征的路线聚合,通过自底向上和自顶向下的路线增强特征图的表达能力。
照着图来说,SPPF是属于backbone的。所以我这里就不放在neck里面说了。
YOLOv8中引入了PAN-FPN(Path Aggregation Network -- Feature Pyramid Network)作为其特征金字塔网络,进一步增强了多尺度特征的表示能力。对比YOLOv5,YOLOv8将YOLOv5中PAN-FPN上采样阶段中的卷积层结构删除了,同时也将C3模块替换为了C2f模块。
虽然前面提到了yolov8并没有明显的区分neck和head,但是mmyolo里面仍然画出了两种构造图,一张是区分了head和neck,一种是没有。我截图如下:
但是无所谓哈,因为本质上都是一样的。叫法不一,但是内容不变。我们仍然按照之前的叫法来做。最后,YOLO Head 是YOLOv8的分类器与回归器。通过Backbone和FPN,我们已经可以获得三个加强过的有效特征层。每一个特征层都有宽,高和通道数,此时我们可以将特征图看作一个又一个特征点的集合,每个特征点作为先验点,而不再存在先验框,每一个先验点都有通道数个特征。YOLO Head 实际上所作的就是对特征点进行判断,判断特征点上的先验框是否有物体与其对应。YOLOv8所用的解耦头是分开的,也就是分类和回归不在一个1*1 卷积里面实现。
而YOLOv8 HEAD整体参考了 PAN 网络骨架(基于FPN发展而来)。参考架构图 FPN的思想主要体现在 4~9层 和 10~15层之间地关联关系。 PANet思想主要体现在 16~21层地信息传递。我们这里分别介绍
3,PAN-FPN结构
YOLOv8中引入了 PAN-FPN(Path Aggregation Network-Feature Pyramid Network)作为其特征金字塔网络,进一步增强了多尺度特征的表示能力。
3.1 FPN 网络骨架
实验表明,浅层的网络更关注于细节信息,高层的网络更关注语义信息,而高层的语义信息能够帮助我们准确的检测出目标,因此我们可以利用最后一个卷积层的feature map进行预测。这种方法呢存在大多数深度网络中,比如VGG,ResNet,DenseNet等,他们都是利用深度网络的最后一层特征来进行分类。这种方法的优点是速度快,需要内存小。缺点就是我们仅仅关注深层网络中最后一层的特征,却忽略了其他层的特征,而且检测小物体的性能急剧下降,但是细节信息可以在一定程度上提升检测的精度。
为了解决这些问题。研究者将处理过的低层特征和处理过的高层特征进行累加,这样做的目的是因为低层特征可以提供更加准确的位置信息,而多次的降采样和上采样操作使得深层网络的定位信息存在误差,因此我们将其结合起来实验,就构建了一个更深的特征金字塔,融合了多层特征信息,并在不同的特征进行输出。这就是FPN。
FPN的架构图如下:
上图可以看到,自顶向下的过程通过上采用(Up-Sampling)的方式将顶层的小特征图放大到上一个stage的特征图一样的大小。上采样的方法是最近邻插值法:实验最近邻插值法可以在上采样的过程中最大程度地保留特征图地语义信息(有利于分类)。从而与backbone地bottom-up过程中相应地具有丰富地空间信息地特征图进行融合,从而得到既有良好地空间信息又有较强烈地语义信息地特征图。
顺带提一句bottom-up地过程:自底向上就是卷积网络地前向过程,我们可以选择不同地骨干网络,例如ResNet-50或ResNet-18等。前向网络地返回值依次是C2,C3,C4, C5,是每次池化之后得到地Feature Map。在残差网络中,C2,C3,C4,C5经过地降采样次数分别是2,3,4,5.即分别对应原图中地步长分别是4,8,16,32.这里之所以没有使用C1,是考虑到由于C1地尺度过大,训练过程中会消耗很多地显存。
在YOLOv8中,FPN主要负责构建从低层到高层的多尺度特征图。其主要过程如下:
1,自顶向下路径
- 上采样操作:对卷积后地特征图进行2倍上采样,恢复到上一层地特征图大小。比如20*20*512地特征图,经过上采样后变为40*40*512
- 逐元素相加:每一层的上采样特征与相应的低层特征进行逐元素相加,以补充空间信息和增强语义信息。即上采样后地40*40*512地特征和自底向上地特征40*40*512进行逐个元素相加,得到地特征是40*40*1532.
2,横向链接
- 通道拼接:利用1*1卷积调整通道数,使得上采样特征与底层特征的通道数一致,然后在通道维度上拼接在一起
3.2 PANet 网络骨架
PANet(Path Aggregation Network)主要贡献点就是提出了一个自顶向下和自底向上地双向融合骨干网络,同时在最底层和最高层之间添加了一条“short-cut”,用于缩短层之间地路径,将网络浅层的较准确的位置信息传递,融合到深层的特征中。
PANet还提出了自适应特征池化和全连接融合两个模块。其中自适应特征池化可以用于聚合不同层之间地特征,保证特征地完整性和多样性,而通过全连接融合可以得到更加准确地预测mask。
PANet地整体结构如下,我们可以之间看这个架构图:
如上图所示,PANet地网络结构由五个核心模块组成。其中:
- A是一个backbone提取特征然后结合FPN,通过这个自底向上模块,可以得到C1,C2,C3,C4,C5五组不同尺寸的 Feature Map,然后FPN通过自顶向下的路径得到四组Feature Map:P5,P4,P3,P2;
- B是PANet在FPN的自底向上路径之后又增加地自底向上地特征融合层,通过这个路径PANet得到了N2, N3, N4, N5 总共四个 Feature Map;
- C是自适应池化特征层,在FPN中,每个Feature Map都会输出一个预测结果,FPN这么做的原因就是处于感受野和网络深度成正比的关系,即网络越深,网络上的像素点的感受野越大,但是实际上不同层次的Feature Map的不同特性并不仅仅只有感受野,还有他们不同的侧重点,基于这个思想,PANet提出了融合所有层的 Feature Map的池化操作--自适应特征池化层。自适应特征池化层先将通过RPN提取的ROI压缩为一维的特征向量,然后通过取max或者取和的方式进行不同Feature Map的融合;
- D是PANet的bounding box预测头,最后会在自适应特征池化层的基础上做bounding box和类别的预测;
- E是用于预测掩码的全连接融合层,全连接和FCN都被广泛的应用到分割图的预测。
4,手写yolov8的neck代码
4.1 根据模型梳理neck结构
FPN是YOLOv8的加强特征提取网络,在主干部分获得的三个有效特征层会在这一部分进行特征融合,特征融合的目的是结合不同尺度的特征信息。在FPN部分,已经获得的有效特征层被用于继续提取特征。我们再扣一个清晰的图来结合代码分析:
在特征利用部分,YOLOv8提取多特征层进行目标检测,一共提取三个特征层。三个特种层分辨位于主干部分的不同位置,分别位置中间层,中下层,底层。当输入为(640, 640, 3)的时候,三个特种层的shape分别是feat1=(80,80,256)、feat2=(40,40,512)、feat3=(20,20,1024 * deep_mul)。
在获得三个有效特征层后,我们利用这三个有效特征层进行FPN层的构建,首先就是feat3=(20,20,512)的特征层:feat3进行上采样UmSampling2d后与feat2=(40,40,512)特征层进行结合,然后使用CSP模块进行特征提取获得P4,此时获得的特征层为(40,40,512)。第二个 feat2=(40,40,512)的特征层:进行上采样UmSampling2d后与feat1=(80,80,256)特征层进行结合,然后使用CSP模块进行特征提取获得P3,此时获得的特征层为(80,80,256)。第三个feat1=(80,80,256)的特征层:进行一次3x3卷积进行下采样,下采样后与P4堆叠,然后使用CSP模块进行特征提取获得新P4,此时获得的特征层为(40,40,512)。
P4=(40,40,512)的特征层进行一次3x3卷积进行下采样,下采样后与P5堆叠,然后使用CSP模块进行特征提取获得新P5,此时获得的特征层为(20,20,1024 * deep_mul)。感觉是图上画错了。最后输出。
所以说neck总共是图中的10——21层。而上一节已经实现了0~9层。对于所用到Conv, C2f,SPPF网络也讲过了。所以这里没有难点,直接进行代码整合。
4.2 concat代码
其实concat并不是一个模块,只是将两个输入拼接而已。如下图所示:
我圈出来的,就是将40*40*512 的特征和 40*40*256的特征进行拼接,最后得到 40*40*(512+256)=40*40*768。就是这么朴实无华。
代码如下,只是包装了以下torch.cat函数:
class Concat(nn.Module):
"""Concatenate a list of tensors along dimension."""
def __init__(self, dimension=1):
"""Concatenates a list of tensors along a specified dimension."""
super().__init__()
self.d = dimension
def forward(self, x):
"""Forward pass for the YOLOv8 mask Proto module."""
return torch.cat(x, self.d)
4.3 neck的代码组合
上面我们已经梳理清楚了。下面直接写即可。代码如下(包括了backbone的代码):
class PoseModel(nn.Module):
"""YOLOv8 pose model."""
def __init__(self, base_channels, base_depth, deep_mul):
super().__init__()
self.model = nn.Sequential(
# backbone
# stem layer 640*640*3 -> 320*320*64
Conv(c1=3, c2=base_channels, k=3, s=2, p=1),
# stage layer1(dark2) 320*320*64 -> 160*160*128 -> 160*160*128
Conv(c1=base_channels, c2=base_channels * 2, k=3, s=2, p=1),
C2f(base_channels * 2, base_channels * 2, base_depth, True),
# stage layer2(dark3) 160*160*128 -> 80*80*256 -> 80*80*256
Conv(c1=base_channels * 2, c2=base_channels * 4, k=3, s=2, p=1),
C2f(base_channels * 4, base_channels * 4, base_depth * 2, True),
# stage layer3(dark4) 80*80*256 -> 40*40*512 -> 40*40*512
Conv(c1=base_channels * 4, c2=base_channels * 8, k=3, s=2, p=1),
C2f(base_channels * 8, base_channels * 8, base_depth * 2, True),
# stage layer4(dark5) 40*40*512 -> 20*20*512 -> 20*20*512
Conv(c1=base_channels * 8, c2=int(base_channels * 16 * deep_mul), k=3, s=2, p=1),
C2f(int(base_channels * 16 * deep_mul), int(base_channels * 16 * deep_mul), base_depth, True),
SPPF(int(base_channels * 16 * deep_mul), int(base_channels * 16 * deep_mul), k=5),
# neck 加强特征提取
# 1024*deep_mul + 512, 40, 40 --> 512, 40, 40
nn.Upsample(scale_factor=2, mode="nearest"),
Concat(), # cat backbone P4
C2f(int(base_channels * 16*deep_mul) + base_channels * 8, base_channels * 8, base_depth, False),
# 768, 80, 80 -> 256, 80, 80
nn.Upsample(scale_factor=2, mode="nearest"),
Concat(), # cat backbone P3
# 15 (P3/8 - small)
C2f(base_channels * 8 + base_channels * 4, base_channels * 4, base_depth, False),
# down_sample 256, 80, 80 -> 256, 40, 40
Conv(c1=base_channels * 4, c2=base_channels * 4, k=3, s=2, p=1),
Concat(), # cat head P4
# 18 (P4/16 - medium)
# 512 + 256, 40, 40 ==> 512, 40, 40
C2f(base_channels * 8 + base_channels * 4, base_channels * 8, base_depth, False),
# down_sample 512, 40, 40 --> 512, 20, 20
Conv(c1=base_channels * 8, c2=base_channels * 8, k=3, s=2, p=1),
Concat(), # cat head P5
# 21 (P5/32-large)
# 1024*deep_mul + 512, 20, 20 --> 1024*deep_mul, 20, 20
C2f(base_channels * 8 + int(base_channels * 16 * deep_mul), int(base_channels * 16 * deep_mul), base_depth, False),
)
def forward(self, x):
# backbone 基础卷积提取特征
feat1 = x
feat2 = x
feat3 = x
for layer in self.model[:5]:
feat1 = layer(feat1)
for layer in self.model[:7]:
feat2 = layer(feat2)
feat3 = self.model[:10](feat3)
print("feat1, feat2, feat3 shape is ", feat1.shape, feat2.shape, feat3.shape)
#------------------------加强特征提取网络------------------------#
# 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 40, 40
P5_upsample = self.model[10:11](feat3)
# 1024 * deep_mul, 40, 40 cat 512, 40, 40 => 1024 * deep_mul + 512, 40, 40
P4 = self.model[11:12]([P5_upsample, feat2])
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
P4 = self.model[12:13](P4)
# 512, 40, 40 => 512, 80, 80
P4_upsample = self.model[13:14](P4)
# 512, 80, 80 cat 256, 80, 80 => 768, 80, 80
P3 = self.model[14:15]([P4_upsample, feat1])
# 768, 80, 80 => 256, 80, 80
P3 = self.model[15:16](P3)
# 256, 80, 80 => 256, 40, 40
P3_downsample = self.model[16:17](P3)
# 512, 40, 40 cat 256, 40, 40 => 768, 40, 40
P4 = self.model[17:18]([P3_downsample, P4])
# 768, 40, 40 => 512, 40, 40
P4 = self.model[18:19](P4)
# 512, 40, 40 => 512, 20, 20
P4_downsample = self.model[19:20](P4)
# 512, 20, 20 cat 1024 * deep_mul, 20, 20 => 1024 * deep_mul + 512, 20, 20
P5 = self.model[20:21]([P4_downsample, feat3])
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
P5 = self.model[21:22](P5)
return P3, P4, P5
4.4 预训练权重的backbone和neck的提取
提取代码如下:
# 加载整个模型的权重
weight_path = r"D:\work\workdata\public\keypoint_code\ultralytics\weights\yolov8n-pose.pt"
weights = torch.load(weight_path, map_location=torch.device('cpu'))
weights_model = weights["model"] if isinstance(weights, dict) else weights
full_model_state_dict = torch.nn.Module.state_dict(weights_model)
# 筛选出backbone部分的权重
prefixes = [f'model.{i}.' for i in range(0, 22)]
backbone_state_dict = {k: v for k, v in full_model_state_dict.items() if any(k.startswith(prefix) for prefix in prefixes)}
# 保存筛选后的权重到新的.pth文件
output_path = r"D:\work\workdata\public\keypoint_code\ultralytics\weights\yolov8n-pose-backbone-neck.pt"
torch.save(backbone_state_dict, output_path)
4.5 使用自己写的neck加载预训练模型
当我们准备好了从预训练权重中抽取的模型的时候,我们就可以使用自己写的模型来加载预训练权重了。
代码如下:
phi = "n"
depth_dict = {'n' : 0.33, 's' : 0.33, 'm' : 0.67, 'l' : 1.00, 'x' : 1.00,}
width_dict = {'n' : 0.25, 's' : 0.50, 'm' : 0.75, 'l' : 1.00, 'x' : 1.25,}
deep_width_dict = {'n' : 1.00, 's' : 1.00, 'm' : 0.75, 'l' : 0.50, 'x' : 0.50,}
dep_mul, wid_mul, deep_mul = depth_dict[phi], width_dict[phi], deep_width_dict[phi]
base_channels = int(wid_mul * 64) # 64
base_depth = max(round(dep_mul * 3), 1) # 3
#-----------------------------------------------#
# 输入图片是640, 640, 3
#-----------------------------------------------#
#---------------------------------------------------#
# 生成主干模型和特征加强网络
# 获得三个有效特征层,他们的shape分别是:
# 256 * deep_mul, 80, 80
# 512 * deep_mul, 40, 40
# 1024 * deep_mul, 20, 20
#---------------------------------------------------#
print(base_channels, base_depth, deep_mul)
model = PoseModel(base_channels, base_depth, deep_mul)
print(model)
ckpt_path = r"D:\work\workdata\public\keypoint_code\ultralytics\weights\yolov8n-pose-backbone-neck.pt"
saved_weights = torch.load(ckpt_path, map_location=torch.device('cpu'))
# 尝试加载权重到模型
model.load_state_dict(saved_weights)
random_input = torch.randn(8, 3, 640, 640)
res = model.forward(random_input)
print(len(res))
print(res[0].shape, res[1].shape, res[2].shape)
print("successful")
打印输出结果:
这样就完成了对yolov8 neck的所有理解。