睿智的目标检测66——Pytorch搭建YoloV8目标检测平台
学习前言
又搞了个YoloV8，看起来似乎在抢这个名字。
源码下载
https://github.com/bubbliiiing/yolov8-pytorch
喜欢的可以点个star噢。
YoloV8改进的部分（不完全）
很多细节与YoloV7关系并不大，大概不是同一组人开发的原因。
1、主干部分：与此前的YoloV5系列差距不大，不过相比之前第一次卷积的卷积核缩小了，是3而不是6。另外CSP模块的预处理从三次卷积换成了两次卷积，具体的实现方式是第一次卷积的通道数扩充为原来的两倍，然后将卷积结果在通道上对半分割。另外借鉴了YoloV7的多堆叠结构。
2、加强特征提取部分：不再对主干网络获得的特征层进行卷积（目的估计是加快速度），另外CSP模块的预处理从三次卷积换成了两次卷积，实现方式与主干网络一样。
3、预测头：加入了DFL模块，DFL模块简单理解就是以概率的方式获得回归值，比如我们当前设置DFL的长度为8，那么某个回归值的计算方式为：
| 预测结果取softmax | 0.0 | 0.1 | 0.0 | 0.0 | 0.4 | 0.5 | 0.0 | 0.0 | 点乘 |
|---|---|---|---|---|---|---|---|---|---|
| 参考的固定值 | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 0.1 * 1 + 0.4 * 4 + 0.5 * 5 = 4.2 |
也就是预测的概率和range(0,8)进行点乘，获得回归值。在YoloV8中，DFL的长度为16，一共有四个需要回归的目标，所以回归头的通道为16 × \times × 4 = 64。
4、自适应多正样本匹配：在YoloV8中，参考YoloX使用了无anchors的实现，是一个无锚点算法，面对长宽不规则的目标比较有优势；在计算损失进行正样本匹配时，正样本需要满足两个条件：一是在真实框内、二是真实框topk最符合要求的正样本（预测框与真实框重合度高且种类预测准确）。
以上并非全部的改进部分，还存在一些其它的改进，这里只列出来了一些我比较感兴趣，而且非常有效的改进。
YoloV8实现思路
一、整体结构解析
在学习YoloV8之前，我们需要对YoloV8所作的工作有一定的了解，这有助于我们后面去了解网络的细节，YoloV8在预测方式上与之前的Yolo并没有多大的差别，依然分为三个部分。
分别是Backbone，FPN以及Yolo Head。
Backbone是YoloV8的主干特征提取网络，输入的图片首先会在主干网络里面进行特征提取，提取到的特征可以被称作特征层，是输入图片的特征集合。在主干部分，我们获取了三个特征层进行下一步网络的构建，这三个特征层我称它为有效特征层。
FPN是YoloV8的加强特征提取网络，在主干部分获得的三个有效特征层会在这一部分进行特征融合，特征融合的目的是结合不同尺度的特征信息。在FPN部分，已经获得的有效特征层被用于继续提取特征。在YoloV8里依然使用到了Panet的结构，我们不仅会对特征进行上采样实现特征融合，还会对特征再次进行下采样实现特征融合。
Yolo Head是YoloV8的分类器与回归器，通过Backbone和FPN，我们已经可以获得三个加强过的有效特征层。每一个特征层都有宽、高和通道数，此时我们可以将特征图看作一个又一个特征点的集合，每个特征点作为先验点，而不再存在先验框，每一个先验点都有通道数个特征。Yolo Head实际上所做的工作就是对特征点进行判断，判断特征点上的先验框是否有物体与其对应。YoloV8所用的解耦头是分开的，也就是分类和回归不在一个1X1卷积里实现。
因此，整个YoloV8网络所作的工作依然就是 特征提取-特征加强-预测先验框对应的物体情况。
二、网络结构解析
1、主干网络Backbone介绍
YoloV8所使用的主干特征提取网络主要为速度快做了一些优化：
1、颈部结构使用普通的步长为2的3x3卷积。
此前：
YoloV5最初使用了Focus结构来初步提取特征，在改进后使用了大卷积核的卷积来初步提取特征，速度都不快。
YoloV7则使用了三次卷积来初步提取特征，速度也不快。
YoloV8则使用普通的步长为2的3x3卷积核来初步提取特征（估计是感受野够了）。
self.stem = Conv(3, base_channels, 3, 2)
这样做会损失一些感受野，但是可以提高模型的速度。
2、CSP模块的预处理从三次卷积换成了两次卷积，并且借鉴了YoloV7的多堆叠结构。
具体的实现方式是第一次卷积的通道数扩充为原来的两倍，然后将卷积结果在通道上对半分割，这样可以减少一次卷积的次数，加快网络的速度。
实现代码如下：
class C2f(nn.Module):
# CSPNet结构结构，大残差结构
# c1为输入通道数，c2为输出通道数
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
super().__init__()
self.c = int(c2 * e)
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1)
self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
def forward(self, x):
# 进行一个卷积，然后划分成两份，每个通道都为c
y = list(self.cv1(x).split((self.c, self.c), 1))
# 每进行一次残差结构都保留，然后堆叠在一起，密集残差
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
整个主干实现代码为：
import torch
import torch.nn as nn
def autopad(k, p=None, d=1):
# kernel, padding, dilation
# 对输入的特征层进行自动padding，按照Same原则
if d > 1:
# actual kernel-size
k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]
if p is None:
# auto-pad
p = k // 2 if isinstance(k, int) else [x // 2 for x in k]
return p
class SiLU(nn.Module):
# SiLU激活函数
@staticmethod
def forward(x):
return x * torch.sigmoid(x)
class Conv(nn.Module):
# 标准卷积+标准化+激活函数
default_act = SiLU()
def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
super().__init__()
self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
self.bn = nn.BatchNorm2d(c2, eps=0.001, momentum=0.03, affine=True, track_running_stats=True)
self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()
def forward(self, x):
return self.act(self.bn(self.conv(x)))
def forward_fuse(self, x):
return self.act(self.conv(x))
class Bottleneck(nn.Module):
# 标准瓶颈结构，残差结构
# c1为输入通道数，c2为输出通道数
def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
super().__init__()
c_ = int(c2 * e) # hidden channels
self.cv1 = Conv(c1, c_, k[0], 1)
self.cv2 = Conv(c_, c2, k[1], 1, g=g)
self.add = shortcut and c1 == c2
def forward(self, x):
return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))
class C2f(nn.Module):
# CSPNet结构结构，大残差结构
# c1为输入通道数，c2为输出通道数
def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
super().__init__()
self.c = int(c2 * e)
self.cv1 = Conv(c1, 2 * self.c, 1, 1)
self.cv2 = Conv((2 + n) * self.c, c2, 1)
self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))
def forward(self, x):
# 进行一个卷积，然后划分成两份，每个通道都为c
y = list(self.cv1(x).split((self.c, self.c), 1))
# 每进行一次残差结构都保留，然后堆叠在一起，密集残差
y.extend(m(y[-1]) for m in self.m)
return self.cv2(torch.cat(y, 1))
class SPPF(nn.Module):
# SPP结构，5、9、13最大池化核的最大池化。
def __init__(self, c1, c2, k=5):
super().__init__()
c_ = c1 // 2
self.cv1 = Conv(c1, c_, 1, 1)
self.cv2 = Conv(c_ * 4, c2, 1, 1)
self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
def forward(self, x):
x = self.cv1(x)
y1 = self.m(x)
y2 = self.m(y1)
return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
class Backbone(nn.Module):
def __init__(self, base_channels, base_depth, deep_mul, phi, pretrained=False):
super().__init__()
#-----------------------------------------------#
# 输入图片是3, 640, 640
#-----------------------------------------------#
# 3, 640, 640 => 32, 640, 640 => 64, 320, 320
self.stem = Conv(3, base_channels, 3, 2)
# 64, 320, 320 => 128, 160, 160 => 128, 160, 160
self.dark2 = nn.Sequential(
Conv(base_channels, base_channels * 2, 3, 2),
C2f(base_channels * 2, base_channels * 2, base_depth, True),
)
# 128, 160, 160 => 256, 80, 80 => 256, 80, 80
self.dark3 = nn.Sequential(
Conv(base_channels * 2, base_channels * 4, 3, 2),
C2f(base_channels * 4, base_channels * 4, base_depth * 2, True),
)
# 256, 80, 80 => 512, 40, 40 => 512, 40, 40
self.dark4 = nn.Sequential(
Conv(base_channels * 4, base_channels * 8, 3, 2),
C2f(base_channels * 8, base_channels * 8, base_depth * 2, True),
)
# 512, 40, 40 => 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 20, 20
self.dark5 = nn.Sequential(
Conv(base_channels * 8, int(base_channels * 16 * deep_mul), 3, 2),
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)
)
if pretrained:
url = {
"n" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_n_backbone_weights.pth',
"s" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_s_backbone_weights.pth',
"m" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_m_backbone_weights.pth',
"l" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_l_backbone_weights.pth',
"x" : 'https://github.com/bubbliiiing/yolov8-pytorch/releases/download/v1.0/yolov8_x_backbone_weights.pth',
}[phi]
checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", model_dir="./model_data")
self.load_state_dict(checkpoint, strict=False)
print("Load weights from " + url.split('/')[-1])
def forward(self, x):
x = self.stem(x)
x = self.dark2(x)
#-----------------------------------------------#
# dark3的输出为256, 80, 80，是一个有效特征层
#-----------------------------------------------#
x = self.dark3(x)
feat1 = x
#-----------------------------------------------#
# dark4的输出为512, 40, 40，是一个有效特征层
#-----------------------------------------------#
x = self.dark4(x)
feat2 = x
#-----------------------------------------------#
# dark5的输出为1024 * deep_mul, 20, 20，是一个有效特征层
#-----------------------------------------------#
x = self.dark5(x)
feat3 = x
return feat1, feat2, feat3
2、构建FPN特征金字塔进行加强特征提取
在特征利用部分，YoloV8提取多特征层进行目标检测，一共提取三个特征层。
三个特征层位于主干部分的不同位置，分别位于中间层，中下层，底层，当输入为(640,640,3)的时候，三个特征层的shape分别为feat1=(80,80,256)、feat2=(40,40,512)、feat3=(20,20,1024 * deep_mul)。
deep_mul只是个系数，对深层的通道进行缩放，在YoloV8中应该是为了平衡计算量做的考虑。
在获得三个有效特征层后，我们利用这三个有效特征层进行FPN层的构建，构建方式为（在本博文中，将SPPCSPC结构归于FPN中）：
- feat3=(20,20,512)的特征层feat3进行上采样UmSampling2d后与feat2=(40,40,512)特征层进行结合，然后使用CSP模块进行特征提取获得P4，此时获得的特征层为(40,40,512)。
- P4=(40,40,512)进行上采样UmSampling2d后与feat1=(80,80,256)特征层进行结合，然后使用CSP模块进行特征提取获得P3，此时获得的特征层为(80,80,256)。
- P3=(80,80,256)的特征层进行一次3x3卷积进行下采样，下采样后与P4堆叠，然后使用CSP模块进行特征提取获得新P4，此时获得的特征层为(40,40,512)。
- P4=(40,40,512)的特征层进行一次3x3卷积进行下采样，下采样后与P5堆叠，然后使用CSP模块进行特征提取获得新P5，此时获得的特征层为(20,20,1024 * deep_mul)。
特征金字塔可以将不同shape的特征层进行特征融合，有利于提取出更好的特征。
#---------------------------------------------------#
# yolo_body
#---------------------------------------------------#
class YoloBody(nn.Module):
def __init__(self, input_shape, num_classes, phi, pretrained=False):
super(YoloBody, self).__init__()
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
#-----------------------------------------------#
# 输入图片是3, 640, 640
#-----------------------------------------------#
#---------------------------------------------------#
# 生成主干模型
# 获得三个有效特征层，他们的shape分别是：
# 256, 80, 80
# 512, 40, 40
# 1024 * deep_mul, 20, 20
#---------------------------------------------------#
self.backbone = Backbone(base_channels, base_depth, deep_mul, phi, pretrained=pretrained)
#------------------------加强特征提取网络------------------------#
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
self.conv3_for_upsample1 = C2f(int(base_channels * 16 * deep_mul) + base_channels * 8, base_channels * 8, base_depth, shortcut=False)
# 768, 80, 80 => 256, 80, 80
self.conv3_for_upsample2 = C2f(base_channels * 8 + base_channels * 4, base_channels * 4, base_depth, shortcut=False)
# 256, 80, 80 => 256, 40, 40
self.down_sample1 = Conv(base_channels * 4, base_channels * 4, 3, 2)
# 512 + 256, 40, 40 => 512, 40, 40
self.conv3_for_downsample1 = C2f(base_channels * 8 + base_channels * 4, base_channels * 8, base_depth, shortcut=False)
# 512, 40, 40 => 512, 20, 20
self.down_sample2 = Conv(base_channels * 8, base_channels * 8, 3, 2)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
self.conv3_for_downsample2 = C2f(int(base_channels * 16 * deep_mul) + base_channels * 8, int(base_channels * 16 * deep_mul), base_depth, shortcut=False)
#------------------------加强特征提取网络------------------------#
ch = [base_channels * 4, base_channels * 8, int(base_channels * 16 * deep_mul)]
self.shape = None
self.nl = len(ch)
# self.stride = torch.zeros(self.nl)
self.stride = torch.tensor([256 / x.shape[-2] for x in self.backbone.forward(torch.zeros(1, 3, 256, 256))]) # forward
self.reg_max = 16 # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
self.no = num_classes + self.reg_max * 4 # number of outputs per anchor
self.num_classes = num_classes
c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], num_classes) # channels
self.cv2 = nn.ModuleList(nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, num_classes, 1)) for x in ch)
self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
def fuse(self):
print('Fusing layers... ')
for m in self.modules():
if type(m) is Conv and hasattr(m, 'bn'):
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
delattr(m, 'bn') # remove batchnorm
m.forward = m.forward_fuse # update forward
return self
def forward(self, x):
# backbone
feat1, feat2, feat3 = self.backbone.forward(x)
#------------------------加强特征提取网络------------------------#
# 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 40, 40
P5_upsample = self.upsample(feat3)
# 1024 * deep_mul, 40, 40 cat 512, 40, 40 => 1024 * deep_mul + 512, 40, 40
P4 = torch.cat([P5_upsample, feat2], 1)
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
P4 = self.conv3_for_upsample1(P4)
# 512, 40, 40 => 512, 80, 80
P4_upsample = self.upsample(P4)
# 512, 80, 80 cat 256, 80, 80 => 768, 80, 80
P3 = torch.cat([P4_upsample, feat1], 1)
# 768, 80, 80 => 256, 80, 80
P3 = self.conv3_for_upsample2(P3)
# 256, 80, 80 => 256, 40, 40
P3_downsample = self.down_sample1(P3)
# 512, 40, 40 cat 256, 40, 40 => 768, 40, 40
P4 = torch.cat([P3_downsample, P4], 1)
# 768, 40, 40 => 512, 40, 40
P4 = self.conv3_for_downsample1(P4)
# 512, 40, 40 => 512, 20, 20
P4_downsample = self.down_sample2(P4)
# 512, 20, 20 cat 1024 * deep_mul, 20, 20 => 1024 * deep_mul + 512, 20, 20
P5 = torch.cat([P4_downsample, feat3], 1)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
P5 = self.conv3_for_downsample2(P5)
#------------------------加强特征提取网络------------------------#
# P3 256, 80, 80
# P4 512, 40, 40
# P5 1024 * deep_mul, 20, 20
shape = P3.shape # BCHW
# P3 256, 80, 80 => num_classes + self.reg_max * 4, 80, 80
# P4 512, 40, 40 => num_classes + self.reg_max * 4, 40, 40
# P5 1024 * deep_mul, 20, 20 => num_classes + self.reg_max * 4, 20, 20
x = [P3, P4, P5]
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
# num_classes + self.reg_max * 4 , 8400 => cls num_classes, 8400;
# box self.reg_max * 4, 8400
box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.num_classes), 1)
# origin_cls = [xi.split((self.reg_max * 4, self.num_classes), 1)[1] for xi in x]
dbox = self.dfl(box)
return dbox, cls, x, self.anchors.to(dbox.device), self.strides.to(dbox.device)
3、利用Yolo Head获得预测结果
利用FPN特征金字塔，我们可以获得三个加强特征，这三个加强特征的shape分别为(20,20,1024 * deep_mul)、(40,40,512)、(80,80,256)，然后我们利用这三个shape的特征层传入Yolo Head获得预测结果，YoloV8使用了解耦头，并且使用了DFL技术。
与之前Yolo系列不同的是，YoloV8在Yolo Head后使用了一个DFL结构来计算回归值，而不是直接获得回归值，DFL模块简单理解就是以概率的方式获得回归值，比如我们当前设置DFL的长度为8，那么某个回归值的计算方式为：
| 预测结果取softmax | 0.0 | 0.1 | 0.0 | 0.0 | 0.4 | 0.5 | 0.0 | 0.0 | 点乘 |
|---|---|---|---|---|---|---|---|---|---|
| 参考的固定值 | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 0.1 * 1 + 0.4 * 4 + 0.5 * 5 = 4.2 |
而对于每一个特征层，我们可以获得利用两个独立的卷积调整通道数，分别获得每个特征点目标对应预测框的种类和回归值，回归值相关的预测头的通道数与DFL的长度有关，在YoloV8中，DFL的长度均设为16，种类相关的预测头的通道数和需要区分的种类个数相关。
无论使用什么数据集，回归值相关的预测头的通道数均为
16
×
4
=
64
16\times4=64
16×4=64，三个特征层的shape为(20,20,64)，(40,40,64)，(80,80,64)。64可以分为四个16，用于计算四个回归系数。
计算完回归系数后。三个特征层的特征层的shape为(20,20,4)，(40,40,4)，(80,80,4)
如果使用的是voc训练集，类则为20种，种类相关的预测头的通道数为20，三个特征层的shape为(20,20,20)，(40,40,20)，(80,80,20)。用于判断每一个特征点所包含的物体种类。
如果使用的是coco训练集，类则为80种，种类相关的预测头的通道数为80，三个特征层的shape为(20,20,80)，(40,40,80)，(80,80,80)。用于判断每一个特征点所包含的物体种类。
实现代码如下：
import numpy as np
import torch
import torch.nn as nn
from nets.backbone import Backbone, C2f, Conv, SiLU, autopad
from utils.utils_bbox import make_anchors
def fuse_conv_and_bn(conv, bn):
# 混合Conv2d + BatchNorm2d 减少计算量
# Fuse Conv2d() and BatchNorm2d() layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/
fusedconv = nn.Conv2d(conv.in_channels,
conv.out_channels,
kernel_size=conv.kernel_size,
stride=conv.stride,
padding=conv.padding,
dilation=conv.dilation,
groups=conv.groups,
bias=True).requires_grad_(False).to(conv.weight.device)
# 准备kernel
w_conv = conv.weight.clone().view(conv.out_channels, -1)
w_bn = torch.diag(bn.weight.div(torch.sqrt(bn.eps + bn.running_var)))
fusedconv.weight.copy_(torch.mm(w_bn, w_conv).view(fusedconv.weight.shape))
# 准备bias
b_conv = torch.zeros(conv.weight.size(0), device=conv.weight.device) if conv.bias is None else conv.bias
b_bn = bn.bias - bn.weight.mul(bn.running_mean).div(torch.sqrt(bn.running_var + bn.eps))
fusedconv.bias.copy_(torch.mm(w_bn, b_conv.reshape(-1, 1)).reshape(-1) + b_bn)
return fusedconv
class DFL(nn.Module):
# DFL模块
# Distribution Focal Loss (DFL) proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
def __init__(self, c1=16):
super().__init__()
self.conv = nn.Conv2d(c1, 1, 1, bias=False).requires_grad_(False)
x = torch.arange(c1, dtype=torch.float)
self.conv.weight.data[:] = nn.Parameter(x.view(1, c1, 1, 1))
self.c1 = c1
def forward(self, x):
# bs, self.reg_max * 4, 8400
b, c, a = x.shape
# bs, 4, self.reg_max, 8400 => bs, self.reg_max, 4, 8400 => b, 4, 8400
# 以softmax的方式，对0~16的数字计算百分比，获得最终数字。
return self.conv(x.view(b, 4, self.c1, a).transpose(2, 1).softmax(1)).view(b, 4, a)
# return self.conv(x.view(b, self.c1, 4, a).softmax(1)).view(b, 4, a)
#---------------------------------------------------#
# yolo_body
#---------------------------------------------------#
class YoloBody(nn.Module):
def __init__(self, input_shape, num_classes, phi, pretrained=False):
super(YoloBody, self).__init__()
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
#-----------------------------------------------#
# 输入图片是3, 640, 640
#-----------------------------------------------#
#---------------------------------------------------#
# 生成主干模型
# 获得三个有效特征层，他们的shape分别是：
# 256, 80, 80
# 512, 40, 40
# 1024 * deep_mul, 20, 20
#---------------------------------------------------#
self.backbone = Backbone(base_channels, base_depth, deep_mul, phi, pretrained=pretrained)
#------------------------加强特征提取网络------------------------#
self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
self.conv3_for_upsample1 = C2f(int(base_channels * 16 * deep_mul) + base_channels * 8, base_channels * 8, base_depth, shortcut=False)
# 768, 80, 80 => 256, 80, 80
self.conv3_for_upsample2 = C2f(base_channels * 8 + base_channels * 4, base_channels * 4, base_depth, shortcut=False)
# 256, 80, 80 => 256, 40, 40
self.down_sample1 = Conv(base_channels * 4, base_channels * 4, 3, 2)
# 512 + 256, 40, 40 => 512, 40, 40
self.conv3_for_downsample1 = C2f(base_channels * 8 + base_channels * 4, base_channels * 8, base_depth, shortcut=False)
# 512, 40, 40 => 512, 20, 20
self.down_sample2 = Conv(base_channels * 8, base_channels * 8, 3, 2)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
self.conv3_for_downsample2 = C2f(int(base_channels * 16 * deep_mul) + base_channels * 8, int(base_channels * 16 * deep_mul), base_depth, shortcut=False)
#------------------------加强特征提取网络------------------------#
ch = [base_channels * 4, base_channels * 8, int(base_channels * 16 * deep_mul)]
self.shape = None
self.nl = len(ch)
# self.stride = torch.zeros(self.nl)
self.stride = torch.tensor([256 / x.shape[-2] for x in self.backbone.forward(torch.zeros(1, 3, 256, 256))]) # forward
self.reg_max = 16 # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x)
self.no = num_classes + self.reg_max * 4 # number of outputs per anchor
self.num_classes = num_classes
c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], num_classes) # channels
self.cv2 = nn.ModuleList(nn.Sequential(Conv(x, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch)
self.cv3 = nn.ModuleList(nn.Sequential(Conv(x, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, num_classes, 1)) for x in ch)
self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity()
def fuse(self):
print('Fusing layers... ')
for m in self.modules():
if type(m) is Conv and hasattr(m, 'bn'):
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
delattr(m, 'bn') # remove batchnorm
m.forward = m.forward_fuse # update forward
return self
def forward(self, x):
# backbone
feat1, feat2, feat3 = self.backbone.forward(x)
#------------------------加强特征提取网络------------------------#
# 1024 * deep_mul, 20, 20 => 1024 * deep_mul, 40, 40
P5_upsample = self.upsample(feat3)
# 1024 * deep_mul, 40, 40 cat 512, 40, 40 => 1024 * deep_mul + 512, 40, 40
P4 = torch.cat([P5_upsample, feat2], 1)
# 1024 * deep_mul + 512, 40, 40 => 512, 40, 40
P4 = self.conv3_for_upsample1(P4)
# 512, 40, 40 => 512, 80, 80
P4_upsample = self.upsample(P4)
# 512, 80, 80 cat 256, 80, 80 => 768, 80, 80
P3 = torch.cat([P4_upsample, feat1], 1)
# 768, 80, 80 => 256, 80, 80
P3 = self.conv3_for_upsample2(P3)
# 256, 80, 80 => 256, 40, 40
P3_downsample = self.down_sample1(P3)
# 512, 40, 40 cat 256, 40, 40 => 768, 40, 40
P4 = torch.cat([P3_downsample, P4], 1)
# 768, 40, 40 => 512, 40, 40
P4 = self.conv3_for_downsample1(P4)
# 512, 40, 40 => 512, 20, 20
P4_downsample = self.down_sample2(P4)
# 512, 20, 20 cat 1024 * deep_mul, 20, 20 => 1024 * deep_mul + 512, 20, 20
P5 = torch.cat([P4_downsample, feat3], 1)
# 1024 * deep_mul + 512, 20, 20 => 1024 * deep_mul, 20, 20
P5 = self.conv3_for_downsample2(P5)
#------------------------加强特征提取网络------------------------#
# P3 256, 80, 80
# P4 512, 40, 40
# P5 1024 * deep_mul, 20, 20
shape = P3.shape # BCHW
# P3 256, 80, 80 => num_classes + self.reg_max * 4, 80, 80
# P4 512, 40, 40 => num_classes + self.reg_max * 4, 40, 40
# P5 1024 * deep_mul, 20, 20 => num_classes + self.reg_max * 4, 20, 20
x = [P3, P4, P5]
for i in range(self.nl):
x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1)
if self.shape != shape:
self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
self.shape = shape
# num_classes + self.reg_max * 4 , 8400 => cls num_classes, 8400;
# box self.reg_max * 4, 8400
box, cls = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2).split((self.reg_max * 4, self.num_classes), 1)
# origin_cls = [xi.split((self.reg_max * 4, self.num_classes), 1)[1] for xi in x]
dbox = self.dfl(box)
return dbox, cls, x, self.anchors.to(dbox.device), self.strides.to(dbox.device)
三、预测结果的解码
1、获得预测框与得分
由第二步我们可以获得三个特征层的预测结果（假设为COCO数据集）：
80x80的特征层对应两个输出：回归输出(20,20,4)；种类输出(20,20,80)。
40x40的特征层对应两个输出：回归输出(40,40,4)；种类输出(40,40,80)。
20x20的特征层对应两个输出：回归输出(80,80,4)；种类输出(80,80,80)。
将所有特征层在高和宽上平铺后堆叠，可以得到总的回归输出为(8400, 4)；种类输出为(8400, 80)。
但是这个预测结果并不对应着最终的预测框在图片上的位置，还需要解码才可以完成。在YoloV8里，每一个特征层上每一个特征点对应一个预测框。
回归输出4个参数用于判断每一个特征点的回归参数，回归参数调整后可以获得预测框，前两个序号的内容代表预测框左上角的距离，后两个序号的内容代表预测框右下角的距离；
种类输出80个参数用于判断每一个特征点所包含的物体种类。
以(20,20)这个特征层为例，该特征层相当于将图像划分成20x20个特征点，如果某个特征点落在物体的对应框内，就用于预测该物体。
如图所示，蓝色的点为20x20的特征点，此时我们对左图黑色点进行解码操作演示：
1、进行预测框左上角进行计算，利用特征点坐标减去Regression预测结果前两个序号的内容获得预测框的左上角
2、进行预测框右下角进行计算，利用特征点坐标加上Regression预测结果后两个序号的内容获得预测框的右下角
3、此时获得的预测框就可以绘制在图片上了。
除去这样的解码操作，还有非极大抑制的操作需要进行，防止同一种类的框的堆积。
def dist2bbox(distance, anchor_points, xywh=True, dim=-1):
"""Transform distance(ltrb) to box(xywh or xyxy)."""
# 左上右下
lt, rb = torch.split(distance, 2, dim)
x1y1 = anchor_points - lt
x2y2 = anchor_points + rb
if xywh:
c_xy = (x1y1 + x2y2) / 2
wh = x2y2 - x1y1
return torch.cat((c_xy, wh), dim) # xywh bbox
return torch.cat((x1y1, x2y2), dim) # xyxy bbox
def decode_box(self, inputs):
# dbox batch_size, 4, 8400
# cls batch_size, 20, 8400
dbox, cls, origin_cls, anchors, strides = inputs
# 获得中心宽高坐标
dbox = dist2bbox(dbox, anchors.unsqueeze(0), xywh=True, dim=1) * strides
y = torch.cat((dbox, cls.sigmoid()), 1).permute(0, 2, 1)
# 进行归一化，到0~1之间
y[:, :, :4] = y[:, :, :4] / torch.Tensor([self.input_shape[1], self.input_shape[0], self.input_shape[1], self.input_shape[0]]).to(y.device)
return y
2、得分筛选与非极大抑制
得到最终的预测结果后还要进行得分排序与非极大抑制筛选。
得分筛选就是筛选出得分满足confidence置信度的预测框。
非极大抑制就是筛选出一定区域内属于同一种类得分最大的框。
得分筛选与非极大抑制的过程可以概括如下：
1、找出该图片中得分大于门限函数的框。在进行重合框筛选前就进行得分的筛选可以大幅度减少框的数量。
2、对种类进行循环，非极大抑制的作用是筛选出一定区域内属于同一种类得分最大的框，对种类进行循环可以帮助我们对每一个类分别进行非极大抑制。
3、根据得分对该种类进行从大到小排序。
4、每次取出得分最大的框，计算其与其它所有预测框的重合程度，重合程度过大的则剔除。
得分筛选与非极大抑制后的结果就可以用于绘制预测框了。
下图是经过非极大抑制的。
下图是未经过非极大抑制的。
实现代码为：
def non_max_suppression(self, prediction, num_classes, input_shape, image_shape, letterbox_image, conf_thres=0.5, nms_thres=0.4):
#----------------------------------------------------------#
# 将预测结果的格式转换成左上角右下角的格式。
# prediction [batch_size, num_anchors, 85]
#----------------------------------------------------------#
box_corner = prediction.new(prediction.shape)
box_corner[:, :, 0] = prediction[:, :, 0] - prediction[:, :, 2] / 2
box_corner[:, :, 1] = prediction[:, :, 1] - prediction[:, :, 3] / 2
box_corner[:, :, 2] = prediction[:, :, 0] + prediction[:, :, 2] / 2
box_corner[:, :, 3] = prediction[:, :, 1] + prediction[:, :, 3] / 2
prediction[:, :, :4] = box_corner[:, :, :4]
output = [None for _ in range(len(prediction))]
for i, image_pred in enumerate(prediction):
#----------------------------------------------------------#
# 对种类预测部分取max。
# class_conf [num_anchors, 1] 种类置信度
# class_pred [num_anchors, 1] 种类
#----------------------------------------------------------#
class_conf, class_pred = torch.max(image_pred[:, 4:4 + num_classes], 1, keepdim=True)
#----------------------------------------------------------#
# 利用置信度进行第一轮筛选
#----------------------------------------------------------#
conf_mask = (class_conf[:, 0] >= conf_thres).squeeze()
#----------------------------------------------------------#
# 根据置信度进行预测结果的筛选
#----------------------------------------------------------#
image_pred = image_pred[conf_mask]
class_conf = class_conf[conf_mask]
class_pred = class_pred[conf_mask]
if not image_pred.size(0):
continue
#-------------------------------------------------------------------------#
# detections [num_anchors, 6]
# 6的内容为：x1, y1, x2, y2, class_conf, class_pred
#-------------------------------------------------------------------------#
detections = torch.cat((image_pred[:, :4], class_conf.float(), class_pred.float()), 1)
#------------------------------------------#
# 获得预测结果中包含的所有种类
#------------------------------------------#
unique_labels = detections[:, -1].cpu().unique()
if prediction.is_cuda:
unique_labels = unique_labels.cuda()
detections = detections.cuda()
for c in unique_labels:
#------------------------------------------#
# 获得某一类得分筛选后全部的预测结果
#------------------------------------------#
detections_class = detections[detections[:, -1] == c]
#------------------------------------------#
# 使用官方自带的非极大抑制会速度更快一些！
# 筛选出一定区域内，属于同一种类得分最大的框
#------------------------------------------#
keep = nms(
detections_class[:, :4],
detections_class[:, 4],
nms_thres
)
max_detections = detections_class[keep]
# # 按照存在物体的置信度排序
# _, conf_sort_index = torch.sort(detections_class[:, 4]*detections_class[:, 5], descending=True)
# detections_class = detections_class[conf_sort_index]
# # 进行非极大抑制
# max_detections = []
# while detections_class.size(0):
# # 取出这一类置信度最高的，一步一步往下判断，判断重合程度是否大于nms_thres，如果是则去除掉
# max_detections.append(detections_class[0].unsqueeze(0))
# if len(detections_class) == 1:
# break
# ious = bbox_iou(max_detections[-1], detections_class[1:])
# detections_class = detections_class[1:][ious < nms_thres]
# # 堆叠
# max_detections = torch.cat(max_detections).data
# Add max detections to outputs
output[i] = max_detections if output[i] is None else torch.cat((output[i], max_detections))
if output[i] is not None:
output[i] = output[i].cpu().numpy()
box_xy, box_wh = (output[i][:, 0:2] + output[i][:, 2:4])/2, output[i][:, 2:4] - output[i][:, 0:2]
output[i][:, :4] = self.yolo_correct_boxes(box_xy, box_wh, input_shape, image_shape, letterbox_image)
return output
四、训练部分
1、计算loss所需内容
计算loss实际上是网络的预测结果和网络的真实结果的对比。
和网络的预测结果一样，网络的损失也由两个部分组成，分别是回归部分、种类部分。回归部分是特征点的回归参数判断、种类部分是特征点包含的物体的种类。
2、正样本的匹配过程
在YoloV8中，训练时正样本的匹配过程可以分为三部分。
- 根据空间距离判断特征点是否在真实框中。
- 根据代价函数判断特征点是否在真实框内的topk中。
- 去重等后处理。
所谓正样本匹配，就是寻找哪些特征点被认为有对应的真实框，并且负责这个真实框的预测。
a、判断特征点是否在预测框中
在步骤中，首先根据空间距离判断特征点是否在预测框中。YoloV8会对每个真实框进行粗匹配。找到哪些特征点上的哪些先验框可以负责该真实框的预测。
代码根据真实框与特征点的坐标情况，利用特征点坐标减去真实框左上角，利用真实框右下角减去特征点坐标，如果这几个值全都大于0则特征点在真实框内部。
def select_candidates_in_gts(xy_centers, gt_bboxes, eps=1e-9, roll_out=False):
"""select the positive anchor center in gt
Args:
xy_centers (Tensor): shape(h*w, 4)
gt_bboxes (Tensor): shape(b, n_boxes, 4)
Return:
(Tensor): shape(b, n_boxes, h*w)
"""
n_anchors = xy_centers.shape[0]
bs, n_boxes, _ = gt_bboxes.shape
# 计算每个真实框距离每个anchors锚点的左上右下的距离，然后求min
# 保证真实框在锚点附近，包围锚点
if roll_out:
bbox_deltas = torch.empty((bs, n_boxes, n_anchors), device=gt_bboxes.device)
for b in range(bs):
lt, rb = gt_bboxes[b].view(-1, 1, 4).chunk(2, 2) # left-top, right-bottom
bbox_deltas[b] = torch.cat((xy_centers[None] - lt, rb - xy_centers[None]),
dim=2).view(n_boxes, n_anchors, -1).amin(2).gt_(eps)
return bbox_deltas
else:
# 真实框的坐上右下left-top, right-bottom
lt, rb = gt_bboxes.view(-1, 1, 4).chunk(2, 2)
# 真实框距离每个anchors锚点的左上右下的距离
bbox_deltas = torch.cat((xy_centers[None] - lt, rb - xy_centers[None]), dim=2).view(bs, n_boxes, n_anchors, -1)
# return (bbox_deltas.min(3)[0] > eps).to(gt_bboxes.dtype)
return bbox_deltas.amin(3).gt_(eps)
b、判断特征点是否在真实框内的topk中
在YoloV8中，我们会计算一个Cost代价矩阵，代表每个真实框和每个特征点之间的代价关系，Cost代价矩阵由两个部分组成：
1、每个真实框和当前特征点预测框的重合程度；
2、每个真实框和当前特征点预测框的种类预测准确度；
之前的Yolo在做代价函数时，都是以小为好，YoloV8是以大为好，这两者并没有本质区别，差别就是有没有用1去减。
每个真实框和当前特征点预测框的重合程度越高，代表这个特征点已经尝试去拟合该真实框了，因此它的Cost代价就会越大。
每个真实框和当前特征点预测框的种类预测准确度越高，也代表这个特征点已经尝试去拟合该真实框了，因此它的Cost代价就会越大。
Cost代价矩阵的目的是自适应的找到当前特征点应该去拟合的真实框，重合度越高越需要拟合，分类越准越需要拟合，在一定半径内越需要拟合。
YoloV8没有使用OTA的思想，每个真实框最大匹配13个特征点。
因此，判断特征点是否在真实框内的topk中的过程总结如下：
1、计算每个真实框和每个特征点预测框的重合程度。取alpha指数。
2、计算每个真实框和每个特征点预测框的种类预测准确度。取beta指数。
3、相加得到Cost代价矩阵。
4、将Cost最大的k个点作为该真实框的正样本。
def get_pos_mask(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes, anc_points, mask_gt):
# pd_scores bs, num_total_anchors, num_classes
# pd_bboxes bs, num_total_anchors, 4
# gt_labels bs, n_max_boxes, 1
# gt_bboxes bs, n_max_boxes, 4
#
# align_metric是一个算出来的代价值，某个先验点属于某个真实框的类的概率乘上某个先验点与真实框的重合程度
# overlaps是某个先验点与真实框的重合程度
# align_metric, overlaps bs, max_num_obj, 8400
align_metric, overlaps = self.get_box_metrics(pd_scores, pd_bboxes, gt_labels, gt_bboxes)
# 正样本锚点需要同时满足：
# 1、在真实框内
# 2、是真实框topk最重合的正样本
# 3、满足mask_gt
# get in_gts mask b, max_num_obj, 8400
# 判断先验点是否在真实框内
mask_in_gts = select_candidates_in_gts(anc_points, gt_bboxes, roll_out=self.roll_out)
# get topk_metric mask b, max_num_obj, 8400
# 判断锚点是否在真实框的topk中
mask_topk = self.select_topk_candidates(align_metric * mask_in_gts, topk_mask=mask_gt.repeat([1, 1, self.topk]).bool())
# merge all mask to a final mask, b, max_num_obj, h*w
# 真实框存在，非padding
mask_pos = mask_topk * mask_in_gts * mask_gt
return mask_pos, align_metric, overlaps
def get_box_metrics(self, pd_scores, pd_bboxes, gt_labels, gt_bboxes):
if self.roll_out:
align_metric = torch.empty((self.bs, self.n_max_boxes, pd_scores.shape[1]), device=pd_scores.device)
overlaps = torch.empty((self.bs, self.n_max_boxes, pd_scores.shape[1]), device=pd_scores.device)
ind_0 = torch.empty(self.n_max_boxes, dtype=torch.long)
for b in range(self.bs):
ind_0[:], ind_2 = b, gt_labels[b].squeeze(-1).long()
# 获得属于这个类别的得分
# bs, max_num_obj, 8400
bbox_scores = pd_scores[ind_0, :, ind_2]
# 计算真实框和预测框的ciou
# bs, max_num_obj, 8400
overlaps[b] = bbox_iou(gt_bboxes[b].unsqueeze(1), pd_bboxes[b].unsqueeze(0), xywh=False, CIoU=True).squeeze(2).clamp(0)
align_metric[b] = bbox_scores.pow(self.alpha) * overlaps[b].pow(self.beta)
else:
# 2, b, max_num_obj
ind = torch.zeros([2, self.bs, self.n_max_boxes], dtype=torch.long)
# b, max_num_obj
# [0]代表第几个图片的
ind[0] = torch.arange(end=self.bs).view(-1, 1).repeat(1, self.n_max_boxes)
# [1]真是标签是什么
ind[1] = gt_labels.long().squeeze(-1)
# 获得属于这个类别的得分
# 取出某个先验点属于某个类的概率
# b, max_num_obj, 8400
bbox_scores = pd_scores[ind[0], :, ind[1]]
# 计算真实框和预测框的ciou
# bs, max_num_obj, 8400
overlaps = bbox_iou(gt_bboxes.unsqueeze(2), pd_bboxes.unsqueeze(1), xywh=False, CIoU=True).squeeze(3).clamp(0)
align_metric = bbox_scores.pow(self.alpha) * overlaps.pow(self.beta)
return align_metric, overlaps
def select_topk_candidates(self, metrics, largest=True, topk_mask=None):
"""
Args:
metrics : (b, max_num_obj, h*w).
topk_mask : (b, max_num_obj, topk) or None
"""
# 8400
num_anchors = metrics.shape[-1]
# b, max_num_obj, topk
topk_metrics, topk_idxs = torch.topk(metrics, self.topk, dim=-1, largest=largest)
if topk_mask is None:
topk_mask = (topk_metrics.max(-1, keepdim=True) > self.eps).tile([1, 1, self.topk])
# b, max_num_obj, topk
topk_idxs[~topk_mask] = 0
# b, max_num_obj, topk, 8400 -> b, max_num_obj, 8400
# 这一步得到的is_in_topk为b, max_num_obj, 8400
# 代表每个真实框对应的top k个先验点
if self.roll_out:
is_in_topk = torch.empty(metrics.shape, dtype=torch.long, device=metrics.device)
for b in range(len(topk_idxs)):
is_in_topk[b] = F.one_hot(topk_idxs[b], num_anchors).sum(-2)
else:
is_in_topk = F.one_hot(topk_idxs, num_anchors).sum(-2)
# 判断锚点是否在真实框的topk中
is_in_topk = torch.where(is_in_topk > 1, 0, is_in_topk)
return is_in_topk.to(metrics.dtype)
c、去重等后处理
在上述的处理过程中，会存在多个真实框的topk个特征点是同一个特征点的情况。此时我们需要反过来操作，寻找哪个真实框是最适合当前这个特征点的。
在这里的判断哪个真实框是最适合当前这个特征点的方式比较简单，通过判断每个真实框与每个预测框的重合程度即可，重合程度越大，越适合。
def select_highest_overlaps(mask_pos, overlaps, n_max_boxes):
"""if an anchor box is assigned to multiple gts,
the one with the highest iou will be selected.
Args:
mask_pos (Tensor): shape(b, n_max_boxes, h*w)
overlaps (Tensor): shape(b, n_max_boxes, h*w)
Return:
target_gt_idx (Tensor): shape(b, h*w)
fg_mask (Tensor): shape(b, h*w)
mask_pos (Tensor): shape(b, n_max_boxes, h*w)
"""
# b, n_max_boxes, 8400 -> b, 8400
fg_mask = mask_pos.sum(-2)
# 如果有一个anchor被指派去预测多个真实框
if fg_mask.max() > 1:
# b, n_max_boxes, 8400
mask_multi_gts = (fg_mask.unsqueeze(1) > 1).repeat([1, n_max_boxes, 1])
# 如果有一个anchor被指派去预测多个真实框，首先计算这个anchor最重合的真实框
# 然后做一个onehot
# b, 8400
max_overlaps_idx = overlaps.argmax(1)
# b, 8400, n_max_boxes
is_max_overlaps = F.one_hot(max_overlaps_idx, n_max_boxes)
# b, n_max_boxes, 8400
is_max_overlaps = is_max_overlaps.permute(0, 2, 1).to(overlaps.dtype)
# b, n_max_boxes, 8400
mask_pos = torch.where(mask_multi_gts, is_max_overlaps, mask_pos)
fg_mask = mask_pos.sum(-2)
# 找到每个anchor符合哪个gt
target_gt_idx = mask_pos.argmax(-2) # (b, h*w)
return target_gt_idx, fg_mask, mask_pos
3、计算Loss
由第一部分可知，YoloV8的损失由两个部分组成：
1、回归部分，由上一部分可知道每个真实框与每个预测框对应关系，可通过预测框找到对应的特征点。因为YoloV8使用了DFL来进行最后的回归预测，所以在回归部分还需要增加上DFL损失。YoloV8的回归损失由iou损失与DFL损失组成。
- iou损失部分比较简单，直接通过计算预测框与真实框的重合程度并使用1-重合程度即可。
- DFL损失则是以概率的方式去计算回归损失，因此要使用到交叉熵
DFL把回归目标定成了分类目标，以真实框的左上角点的x坐标为例，它一般不会位于具体的网格点上，此时它的坐标便不是整数（计算损失不是相对于真实图的，而是相对于每个特征层的网格图的）。
假设真实框的x坐标为7.9，那么它距离8更近，距离7更远。我们便可以这样使用两个交叉熵，预测结果与7求交叉熵，给低一些的权重，预测结果与8求交叉熵，给高一点的权重。
(8-7.9) * cross_entropy(pred_dist, 7)
+
(7.9-7) * cross_entropy(pred_dist, 8)
2、种类部分，由上一部分可知道每个真实框与每个预测框对应关系，可通过预测框找到对应的特征点。取出该先验框的种类预测结果，根据真实框的种类和先验框的种类预测结果计算交叉熵损失，作为种类部分的Loss组成。
不过种类损失的计算部分，标签此时不为1，而是通过重合程度计算出来的，通过代价函数乘上预测框与真实框的重合程度再除上这个真实框对应的最大代价值。
class BboxLoss(nn.Module):
def __init__(self, reg_max=16, use_dfl=False):
super().__init__()
self.reg_max = reg_max
self.use_dfl = use_dfl
def forward(self, pred_dist, pred_bboxes, anchor_points, target_bboxes, target_scores, target_scores_sum, fg_mask):
# 计算IOU损失
# weight代表损失中标签应该有的置信度，0最小，1最大
weight = torch.masked_select(target_scores.sum(-1), fg_mask).unsqueeze(-1)
# 计算预测框和真实框的重合程度
iou = bbox_iou(pred_bboxes[fg_mask], target_bboxes[fg_mask], xywh=False, CIoU=True)
# 然后1-重合程度，乘上应该有的置信度，求和后求平均。
loss_iou = ((1.0 - iou) * weight).sum() / target_scores_sum
# 计算DFL损失
if self.use_dfl:
target_ltrb = bbox2dist(anchor_points, target_bboxes, self.reg_max)
loss_dfl = self._df_loss(pred_dist[fg_mask].view(-1, self.reg_max + 1), target_ltrb[fg_mask]) * weight
loss_dfl = loss_dfl.sum() / target_scores_sum
else:
loss_dfl = torch.tensor(0.0).to(pred_dist.device)
return loss_iou, loss_dfl
@staticmethod
def _df_loss(pred_dist, target):
# Return sum of left and right DFL losses
# Distribution Focal Loss (DFL) proposed in Generalized Focal Loss https://ieeexplore.ieee.org/document/9792391
tl = target.long() # target left
tr = tl + 1 # target right
wl = tr - target # weight left
wr = 1 - wl # weight right
# 一个点一般不会处于anchor点上，一般是xx.xx。如果要用DFL的话，不可能直接一个cross_entropy就能拟合
# 所以把它认为是相对于xx.xx左上角锚点与右下角锚点的距离 如果距离右下角锚点距离小，wl就小，左上角损失就小
# 如果距离左上角锚点距离小，wr就小，右下角损失就小
return (F.cross_entropy(pred_dist, tl.view(-1), reduction="none").view(tl.shape) * wl +
F.cross_entropy(pred_dist, tr.view(-1), reduction="none").view(tl.shape) * wr).mean(-1, keepdim=True)
def xywh2xyxy(x):
"""
Convert bounding box coordinates from (x, y, width, height) format to (x1, y1, x2, y2) format where (x1, y1) is the
top-left corner and (x2, y2) is the bottom-right corner.
Args:
x (np.ndarray) or (torch.Tensor): The input bounding box coordinates in (x, y, width, height) format.
Returns:
y (np.ndarray) or (torch.Tensor): The bounding box coordinates in (x1, y1, x2, y2) format.
"""
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
y[..., 0] = x[..., 0] - x[..., 2] / 2 # top left x
y[..., 1] = x[..., 1] - x[..., 3] / 2 # top left y
y[..., 2] = x[..., 0] + x[..., 2] / 2 # bottom right x
y[..., 3] = x[..., 1] + x[..., 3] / 2 # bottom right y
return y
# Criterion class for computing training losses
class Loss:
def __init__(self, model):
self.bce = nn.BCEWithLogitsLoss(reduction='none')
self.stride = model.stride # model strides
self.nc = model.num_classes # number of classes
self.no = model.no
self.reg_max = model.reg_max
self.use_dfl = model.reg_max > 1
roll_out_thr = 64
self.assigner = TaskAlignedAssigner(topk=10,
num_classes=self.nc,
alpha=0.5,
beta=6.0,
roll_out_thr=roll_out_thr)
self.bbox_loss = BboxLoss(model.reg_max - 1, use_dfl=self.use_dfl)
self.proj = torch.arange(model.reg_max, dtype=torch.float)
def preprocess(self, targets, batch_size, scale_tensor):
if targets.shape[0] == 0:
out = torch.zeros(batch_size, 0, 5, device=targets.device)
else:
# 获得图像索引
i = targets[:, 0]
_, counts = i.unique(return_counts=True)
out = torch.zeros(batch_size, counts.max(), 5, device=targets.device)
# 对batch进行循环，然后赋值
for j in range(batch_size):
matches = i == j
n = matches.sum()
if n:
out[j, :n] = targets[matches, 1:]
# 缩放到原图大小。
out[..., 1:5] = xywh2xyxy(out[..., 1:5].mul_(scale_tensor))
return out
def bbox_decode(self, anchor_points, pred_dist):
if self.use_dfl:
# batch, anchors, channels
b, a, c = pred_dist.shape
# DFL的解码
pred_dist = pred_dist.view(b, a, 4, c // 4).softmax(3).matmul(self.proj.to(pred_dist.device).type(pred_dist.dtype))
# pred_dist = pred_dist.view(b, a, c // 4, 4).transpose(2,3).softmax(3).matmul(self.proj.type(pred_dist.dtype))
# pred_dist = (pred_dist.view(b, a, c // 4, 4).softmax(2) * self.proj.type(pred_dist.dtype).view(1, 1, -1, 1)).sum(2)
# 然后解码获得预测框
return dist2bbox(pred_dist, anchor_points, xywh=False)
def __call__(self, preds, batch):
# 获得使用的device
device = preds[1].device
# box, cls, dfl三部分的损失
loss = torch.zeros(3, device=device)
# 获得特征，并进行划分
feats = preds[2] if isinstance(preds, tuple) else preds
pred_distri, pred_scores = torch.cat([xi.view(feats[0].shape[0], self.no, -1) for xi in feats], 2).split((self.reg_max * 4, self.nc), 1)
# bs, num_classes + self.reg_max * 4 , 8400 => cls bs, num_classes, 8400;
# box bs, self.reg_max * 4, 8400
pred_scores = pred_scores.permute(0, 2, 1).contiguous()
pred_distri = pred_distri.permute(0, 2, 1).contiguous()
# 获得batch size与dtype
dtype = pred_scores.dtype
batch_size = pred_scores.shape[0]
# 获得输入图片大小
imgsz = torch.tensor(feats[0].shape[2:], device=device, dtype=dtype) * self.stride[0]
# 获得anchors点和步长对应的tensor
anchor_points, stride_tensor = make_anchors(feats, self.stride, 0.5)
# 把一个batch中的东西弄一个矩阵
# 0为属于第几个图片
# 1为种类
# 2:为框的坐标
targets = torch.cat((batch[:, 0].view(-1, 1), batch[:, 1].view(-1, 1), batch[:, 2:]), 1)
# 先进行初步的处理，对输入进来的gt进行padding，到最大数量，并把框的坐标进行缩放
# bs, max_boxes_num, 5
targets = self.preprocess(targets.to(device), batch_size, scale_tensor=imgsz[[1, 0, 1, 0]])
# bs, max_boxes_num, 5 => bs, max_boxes_num, 1 ; bs, max_boxes_num, 4
gt_labels, gt_bboxes = targets.split((1, 4), 2) # cls, xyxy
# 求哪些框是有目标的，哪些是填充的
# bs, max_boxes_num
mask_gt = gt_bboxes.sum(2, keepdim=True).gt_(0)
# pboxes
# 对预测结果进行解码，获得预测框
# bs, 8400, 4
pred_bboxes = self.bbox_decode(anchor_points, pred_distri) # xyxy, (b, h*w, 4)
# 对预测框与真实框进行分配
# target_bboxes bs, 8400, 4
# target_scores bs, 8400, 80
# fg_mask bs, 8400
_, target_bboxes, target_scores, fg_mask, _ = self.assigner(
pred_scores.detach().sigmoid(), (pred_bboxes.detach() * stride_tensor).type(gt_bboxes.dtype),
anchor_points * stride_tensor, gt_labels, gt_bboxes, mask_gt
)
target_bboxes /= stride_tensor
target_scores_sum = max(target_scores.sum(), 1)
# 计算分类的损失
# loss[1] = self.varifocal_loss(pred_scores, target_scores, target_labels) / target_scores_sum # VFL way
loss[1] = self.bce(pred_scores, target_scores.to(dtype)).sum() / target_scores_sum # BCE
# 计算bbox的损失
if fg_mask.sum():
loss[0], loss[2] = self.bbox_loss(pred_distri, pred_bboxes, anchor_points, target_bboxes, target_scores,
target_scores_sum, fg_mask)
loss[0] *= 7.5 # box gain
loss[1] *= 0.5 # cls gain
loss[2] *= 1.5 # dfl gain
return loss.sum() # loss(box, cls, dfl) # * batch_size
训练自己的YoloV8模型
首先前往Github下载对应的仓库，下载完后利用解压软件解压，之后用编程软件打开文件夹。
注意打开的根目录必须正确，否则相对目录不正确的情况下，代码将无法运行。
一定要注意打开后的根目录是文件存放的目录。
一、数据集的准备
本文使用VOC格式进行训练，训练前需要自己制作好数据集，如果没有自己的数据集，可以通过Github连接下载VOC12+07的数据集尝试下。
训练前将标签文件放在VOCdevkit文件夹下的VOC2007文件夹下的Annotation中。
训练前将图片文件放在VOCdevkit文件夹下的VOC2007文件夹下的JPEGImages中。
此时数据集的摆放已经结束。
二、数据集的处理
在完成数据集的摆放之后，我们需要对数据集进行下一步的处理，目的是获得训练用的2007_train.txt以及2007_val.txt，需要用到根目录下的voc_annotation.py。
voc_annotation.py里面有一些参数需要设置。
分别是annotation_mode、classes_path、trainval_percent、train_percent、VOCdevkit_path，第一次训练可以仅修改classes_path
'''
annotation_mode用于指定该文件运行时计算的内容
annotation_mode为0代表整个标签处理过程，包括获得VOCdevkit/VOC2007/ImageSets里面的txt以及训练用的2007_train.txt、2007_val.txt
annotation_mode为1代表获得VOCdevkit/VOC2007/ImageSets里面的txt
annotation_mode为2代表获得训练用的2007_train.txt、2007_val.txt
'''
annotation_mode = 0
'''
必须要修改，用于生成2007_train.txt、2007_val.txt的目标信息
与训练和预测所用的classes_path一致即可
如果生成的2007_train.txt里面没有目标信息
那么就是因为classes没有设定正确
仅在annotation_mode为0和2的时候有效
'''
classes_path = 'model_data/voc_classes.txt'
'''
trainval_percent用于指定(训练集+验证集)与测试集的比例，默认情况下 (训练集+验证集):测试集 = 9:1
train_percent用于指定(训练集+验证集)中训练集与验证集的比例，默认情况下 训练集:验证集 = 9:1
仅在annotation_mode为0和1的时候有效
'''
trainval_percent = 0.9
train_percent = 0.9
'''
指向VOC数据集所在的文件夹
默认指向根目录下的VOC数据集
'''
VOCdevkit_path = 'VOCdevkit'
classes_path用于指向检测类别所对应的txt，以voc数据集为例，我们用的txt为：
训练自己的数据集时，可以自己建立一个cls_classes.txt，里面写自己所需要区分的类别。
三、开始网络训练
通过voc_annotation.py我们已经生成了2007_train.txt以及2007_val.txt，此时我们可以开始训练了。
训练的参数较多，大家可以在下载库后仔细看注释，其中最重要的部分依然是train.py里的classes_path。
classes_path用于指向检测类别所对应的txt，这个txt和voc_annotation.py里面的txt一样！训练自己的数据集必须要修改！
修改完classes_path后就可以运行train.py开始训练了，在训练多个epoch后，权值会生成在logs文件夹中。
其它参数的作用如下：
#---------------------------------#
# Cuda 是否使用Cuda
# 没有GPU可以设置成False
#---------------------------------#
Cuda = True
#---------------------------------------------------------------------#
# distributed 用于指定是否使用单机多卡分布式运行
# 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
# Windows系统下默认使用DP模式调用所有显卡，不支持DDP。
# DP模式：
# 设置 distributed = False
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py
# DDP模式：
# 设置 distributed = True
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
#---------------------------------------------------------------------#
distributed = False
#---------------------------------------------------------------------#
# sync_bn 是否使用sync_bn，DDP模式多卡可用
#---------------------------------------------------------------------#
sync_bn = False
#---------------------------------------------------------------------#
# fp16 是否使用混合精度训练
# 可减少约一半的显存、需要pytorch1.7.1以上
#---------------------------------------------------------------------#
fp16 = False
#---------------------------------------------------------------------#
# classes_path 指向model_data下的txt，与自己训练的数据集相关
# 训练前一定要修改classes_path，使其对应自己的数据集
#---------------------------------------------------------------------#
classes_path = 'model_data/voc_classes.txt'
#----------------------------------------------------------------------------------------------------------------------------#
# 权值文件的下载请看README，可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的，因为特征是通用的。
# 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分，用于进行特征提取。
# 预训练权重对于99%的情况都必须要用，不用的话主干部分的权值太过随机，特征提取效果不明显，网络训练的结果也不会好
#
# 如果训练过程中存在中断训练的操作，可以将model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。
# 同时修改下方的 冻结阶段 或者 解冻阶段 的参数，来保证模型epoch的连续性。
#
# 当model_path = ''的时候不加载整个模型的权值。
#
# 此处使用的是整个模型的权重，因此是在train.py进行加载的。
# 如果想要让模型从0开始训练，则设置model_path = ''，下面的Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#
# 一般来讲，网络从0开始的训练效果会很差，因为权值太过随机，特征提取效果不明显，因此非常、非常、非常不建议大家从0开始训练！
# 从0开始训练有两个方案：
# 1、得益于Mosaic数据增强方法强大的数据增强能力，将UnFreeze_Epoch设置的较大（300及以上）、batch较大（16及以上）、数据较多（万以上）的情况下，
# 可以设置mosaic=True，直接随机初始化参数开始训练，但得到的效果仍然不如有预训练的情况。（像COCO这样的大数据集可以这样做）
# 2、了解imagenet数据集，首先训练分类模型，获得网络的主干部分权值，分类模型的 主干部分 和该模型通用，基于此进行训练。
#----------------------------------------------------------------------------------------------------------------------------#
model_path = 'model_data/yolov8_s.pth'
#------------------------------------------------------#
# input_shape 输入的shape大小，一定要是32的倍数
#------------------------------------------------------#
input_shape = [640, 640]
#------------------------------------------------------#
# phi 所使用到的yolov8的版本
# n : 对应yolov8_n
# s : 对应yolov8_s
# m : 对应yolov8_m
# l : 对应yolov8_l
# x : 对应yolov8_x
#------------------------------------------------------#
phi = 's'
#----------------------------------------------------------------------------------------------------------------------------#
# pretrained 是否使用主干网络的预训练权重，此处使用的是主干的权重，因此是在模型构建的时候进行加载的。
# 如果设置了model_path，则主干的权值无需加载，pretrained的值无意义。
# 如果不设置model_path，pretrained = True，此时仅加载主干开始训练。
# 如果不设置model_path，pretrained = False，Freeze_Train = Fasle，此时从0开始训练，且没有冻结主干的过程。
#----------------------------------------------------------------------------------------------------------------------------#
pretrained = False
#------------------------------------------------------------------#
# mosaic 马赛克数据增强。
# mosaic_prob 每个step有多少概率使用mosaic数据增强，默认50%。
#
# mixup 是否使用mixup数据增强，仅在mosaic=True时有效。
# 只会对mosaic增强后的图片进行mixup的处理。
# mixup_prob 有多少概率在mosaic后使用mixup数据增强，默认50%。
# 总的mixup概率为mosaic_prob * mixup_prob。
#
# special_aug_ratio 参考YoloX，由于Mosaic生成的训练图片，远远脱离自然图片的真实分布。
# 当mosaic=True时，本代码会在special_aug_ratio范围内开启mosaic。
# 默认为前70%个epoch，100个世代会开启70个世代。
#------------------------------------------------------------------#
mosaic = True
mosaic_prob = 0.5
mixup = True
mixup_prob = 0.5
special_aug_ratio = 0.7
#------------------------------------------------------------------#
# label_smoothing 标签平滑。一般0.01以下。如0.01、0.005。
#------------------------------------------------------------------#
label_smoothing = 0
#----------------------------------------------------------------------------------------------------------------------------#
# 训练分为两个阶段，分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
# 冻结训练需要的显存较小，显卡非常差的情况下，可设置Freeze_Epoch等于UnFreeze_Epoch，Freeze_Train = True，此时仅仅进行冻结训练。
#
# 在此提供若干参数设置建议，各位训练者根据自己的需求进行灵活调整：
# （一）从整个模型的预训练权重开始训练：
# Adam：
# Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 100，Freeze_Train = True，optimizer_type = 'adam'，Init_lr = 1e-3，weight_decay = 0。（冻结）
# Init_Epoch = 0，UnFreeze_Epoch = 100，Freeze_Train = False，optimizer_type = 'adam'，Init_lr = 1e-3，weight_decay = 0。（不冻结）
# SGD：
# Init_Epoch = 0，Freeze_Epoch = 50，UnFreeze_Epoch = 300，Freeze_Train = True，optimizer_type = 'sgd'，Init_lr = 1e-2，weight_decay = 5e-4。（冻结）
# Init_Epoch = 0，UnFreeze_Epoch = 300，Freeze_Train = False，optimizer_type = 'sgd'，Init_lr = 1e-2，weight_decay = 5e-4。（不冻结）
# 其中：UnFreeze_Epoch可以在100-300之间调整。
# （二）从0开始训练：
# Init_Epoch = 0，UnFreeze_Epoch >= 300，Unfreeze_batch_size >= 16，Freeze_Train = False（不冻结训练）
# 其中：UnFreeze_Epoch尽量不小于300。optimizer_type = 'sgd'，Init_lr = 1e-2，mosaic = True。
# （三）batch_size的设置：
# 在显卡能够接受的范围内，以大为好。显存不足与数据集大小无关，提示显存不足（OOM或者CUDA out of memory）请调小batch_size。
# 受到BatchNorm层影响，batch_size最小为2，不能为1。
# 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大，因为关系到学习率的自动调整。
#----------------------------------------------------------------------------------------------------------------------------#
#------------------------------------------------------------------#
# 冻结阶段训练参数
# 此时模型的主干被冻结了，特征提取网络不发生改变
# 占用的显存较小，仅对网络进行微调
# Init_Epoch 模型当前开始的训练世代，其值可以大于Freeze_Epoch，如设置：
# Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100
# 会跳过冻结阶段，直接从60代开始，并调整对应的学习率。
# （断点续练时使用）
# Freeze_Epoch 模型冻结训练的Freeze_Epoch
# (当Freeze_Train=False时失效)
# Freeze_batch_size 模型冻结训练的batch_size
# (当Freeze_Train=False时失效)
#------------------------------------------------------------------#
Init_Epoch = 0
Freeze_Epoch = 50
Freeze_batch_size = 32
#------------------------------------------------------------------#
# 解冻阶段训练参数
# 此时模型的主干不被冻结了，特征提取网络会发生改变
# 占用的显存较大，网络所有的参数都会发生改变
# UnFreeze_Epoch 模型总共训练的epoch
# SGD需要更长的时间收敛，因此设置较大的UnFreeze_Epoch
# Adam可以使用相对较小的UnFreeze_Epoch
# Unfreeze_batch_size 模型在解冻后的batch_size
#------------------------------------------------------------------#
UnFreeze_Epoch = 300
Unfreeze_batch_size = 16
#------------------------------------------------------------------#
# Freeze_Train 是否进行冻结训练
# 默认先冻结主干训练后解冻训练。
#------------------------------------------------------------------#
Freeze_Train = True
#------------------------------------------------------------------#
# 其它训练参数：学习率、优化器、学习率下降有关
#------------------------------------------------------------------#
#------------------------------------------------------------------#
# Init_lr 模型的最大学习率
# Min_lr 模型的最小学习率，默认为最大学习率的0.01
#------------------------------------------------------------------#
Init_lr = 1e-2
Min_lr = Init_lr * 0.01
#------------------------------------------------------------------#
# optimizer_type 使用到的优化器种类，可选的有adam、sgd
# 当使用Adam优化器时建议设置 Init_lr=1e-3
# 当使用SGD优化器时建议设置 Init_lr=1e-2
# momentum 优化器内部使用到的momentum参数
# weight_decay 权值衰减，可防止过拟合
# adam会导致weight_decay错误，使用adam时建议设置为0。
#------------------------------------------------------------------#
optimizer_type = "sgd"
momentum = 0.937
weight_decay = 5e-4
#------------------------------------------------------------------#
# lr_decay_type 使用到的学习率下降方式，可选的有step、cos
#------------------------------------------------------------------#
lr_decay_type = "cos"
#------------------------------------------------------------------#
# save_period 多少个epoch保存一次权值
#------------------------------------------------------------------#
save_period = 10
#------------------------------------------------------------------#
# save_dir 权值与日志文件保存的文件夹
#------------------------------------------------------------------#
save_dir = 'logs'
#------------------------------------------------------------------#
# eval_flag 是否在训练时进行评估，评估对象为验证集
# 安装pycocotools库后，评估体验更佳。
# eval_period 代表多少个epoch评估一次，不建议频繁的评估
# 评估需要消耗较多的时间，频繁评估会导致训练非常慢
# 此处获得的mAP会与get_map.py获得的会有所不同，原因有二：
# （一）此处获得的mAP为验证集的mAP。
# （二）此处设置评估参数较为保守，目的是加快评估速度。
#------------------------------------------------------------------#
eval_flag = True
eval_period = 10
#------------------------------------------------------------------#
# num_workers 用于设置是否使用多线程读取数据
# 开启后会加快数据读取速度，但是会占用更多内存
# 内存较小的电脑可以设置为2或者0
#------------------------------------------------------------------#
num_workers = 4
四、训练结果预测
训练结果预测需要用到两个文件，分别是yolo.py和predict.py。
我们首先需要去yolo.py里面修改model_path以及classes_path，这两个参数必须要修改。
model_path指向训练好的权值文件，在logs文件夹里。
classes_path指向检测类别所对应的txt。
完成修改后就可以运行predict.py进行检测了。运行后输入图片路径即可检测。
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
