Gold-YOLO(NeurIPS 2023)论文与代码解析

本文提出Gold-YOLO,一种基于Gather-and-Distribute机制的改进版目标检测模型,通过增强信息融合能力和引入MAE-style预训练,优化了YOLO系列的性能,特别是在Nano和Small版本上实现了延迟与精度的良好平衡。
摘要由CSDN通过智能技术生成

paper:Gold-YOLO: Efficient Object Detector via Gather-and-Distribute Mechanism

official implementation:https://github.com/huawei-noah/Efficient-Computing/tree/master/Detection/Gold-YOLO

存在的问题

在过去几年里,YOLO系列已经成为了实时目标检测领域最先进以及最常用的方法。许多研究通过修改模型架构、数据增强、设计新的损失函数将baseline提升到了一个更高的水平。但现有的模型仍然存在信息融合的问题,尽管FPN和PANet在一定程度上缓解了该问题。

传统的neck如FPN以及相关变体的结构如图3(a)所示,但是这种信息融合的方法存在一个明显的缺陷:当需要跨层融合信息时(如level-1和level-3),FPN式的结构无法无损的传输信息,这阻碍了YOLO系列更好的进行信息融合。

本文的创新点

针对FPN式结构存在的问题,本文在TopFormer理论的基础上,提出了一种新的聚合-分发(GD)机制,它通过融合多层特征并将全局信息注入到更高层,在YOLO中实现高效的信息交换。这显著增加了neck的信息融合能力,同时没有显著增加延迟。

基于此提出了一个新的模型Gold-YOLO,它提高了多尺度特征融合的能力,并在所有尺度上实现了延迟和精度之间的理想平衡。

此外,本文首次在YOLO系列中实现了MAE-style的预训练,使得YOLO系列可以从无监督预训练中受益。

方法介绍

如图3所示,在FPN的结构中,只能完全融合相邻层的信息,对于其它层的信息,只能间接的“递归”获得。这种传输模式可能导致计算过程中信息的丢失。为了避免这种情况,本文放弃了递归方法,构建了一种新的聚合-分发机制。通过使用一个统一的模块从各层收集和融合信息,然后将其分发到不同的层。

具体实现中,聚合与分发的过程对应三个模块:特征对齐模块(Feature Alignment Module, FAM)、信息融合模块(Information Fusion Module, IFM)、信息注入模块(Information Injection Module, Inject)。完整结构如图2所示

  • gather过程包括两步。首先,FAM从不同层收集和对齐特征。然后,IFM通过融合对齐的特征得到全局信息。
  • 在获得全局信息后,inject模块将这些信息distribute到每个level中,并使用简单的注意力操作进行注入,从而提高分支的检测能力。

为了增加模型检测不同大小对象的能力,提出了两个分支,low-stage GDhigh-stage GD。如图2所示,neck的输入包括backbone提取的特征图B2,B3,B4,B5,其中 \(B_{i}\in \mathbb{R}^{N\times C_{Bi}\times R_{Bi}}\),N是batch size,C是通道数,R=HxW。

Low GD

结构如图4(a)所示

Low-FAM

在Low-FAM中,用average pooling下采样得到一个统一大小的 \(F_{align}\)。这里选择 \(R_{B4}=\frac{1}{4}R\) 作为目标大小。

Low-IFM

Low-IFM包括多层重参数化卷积Block (RepBlock) 和一个split操作。具体来说,RepBlock取 \(F_{align}\)(\(channel=sum(C_{B2},C_{B3},C_{B4},C_{B5})\))作为输入得到 \(F_{fuse}\)(\(channel=C_{B4}+C_{B5}\)),然后沿通道维度split成 \(F_{inj\_P3}\) 和 \(F_{inj\_P4}\)。如下

Information injection module

为了更有效的将全局信息注入到不同的层,作者采用注意力机制来融合信息,如图5所示。具体来说,同时输入局部信息(当前层)和全局信息(IFM生成的)并分别记为 \(F_{local}\) 和 \(F_{inj}\),\(F_{inj}\) 通过两个不同的卷积层分别得到 \(F_{global\_embed}\) 和 \(F_{act}\)。\(F_{local}\) 通过卷积得到 \(F_{local\_embed}\)。然后通过注意力计算得到融合特征 \(F_{out}\)。其中 \(F_{local}\) 等于 \(Bi\),具体如下

High GD

High-GD融合Low-GD得到的特征 {P3, P4, P5},如图4(b)所示

High-FAM

High-FAM和Low-FAM的操作一样,通过全局平均池化下采样来对齐大小,目标大小为 \(R_{P5}=\frac{1}{8}R\)。

Hign-IFM

High-IFM包括多个transformer block和一个split操作。具体包括三步

  1. High-FAM的输出 \(F_{align}\) 通过transformer block融合得到 \(F_{fuse}\)
  2. \(F_{fuse}\) 通过1x1卷积通道降维到 \(sum(C_{P4},C_{P5})\)
  3. 沿通道进行split操作得到 \(F_{inj\_N4}\) 和 \(F_{inj\_N5}\)

具体如下

式(8)中的transformer融合模块包括多个堆叠的transformer block,每个block包含一个multi-head attention block、一个ffn、一个residual connection。具体配置和LeViT一样,K,Q的维度设为D(例如16),V的维度为2D(例如32)。考虑到推理速度,替换掉了一些速度不友好的操作,每个卷积的LN换成了BN,所有的GELU激活换成了ReLU。为了增强transformer block中的局部连接,两个1x1卷积中间增加了一层深度卷积。FFN的expansion factor设为2。

Information injection module

这里和Low-GD中的结构一样,其中 \(F_{local}\) 等于 \(Pi\),具体如下

Enhanced cross-layer information flow

为了进一步提升性能,作者借鉴YOLOv6里的PAFPN提出了一个Inject-LAF模块。这个模块是注入模块的增加,其中在注入模块的输入位置新加了一个轻量的相邻层融合模块(lightweight adjacent layer fusion, LAF)。为了实现速度和精度的平衡,设计了两种LAF:low-level LAF和high-level LAF,分别用于低层注入(合并相邻两层的特征)和高层注入(合并相邻一层的特征),具体结构如图5(b)所示。

代码解析

官方的实现是基于YOLOv6的实现,其中n,s的neck是"RepGDNeck",m的neck是"GDNeck",l的neck是"GDNeck2",因为从实验结果看,提升比较明显的事nano和small版本,因此这里只解析一下RepGDNeck的实现。具体实现代码在Efficient-Computing/Detection/Gold-YOLO/gold_yolo/reppan.py中,forward实现如下

def forward(self, input):
    (c2, c3, c4, c5) = input  # [(16,32,160,160),(16,64,80,80),(16,128,40,40),(16,256,20,20)]
    
    # Low-GD
    ## use conv fusion global info
    low_align_feat = self.low_FAM(input)  # (16,480,40,40)
    low_fuse_feat = self.low_IFM(low_align_feat)  # (16,96,40,40)
    low_global_info = low_fuse_feat.split(self.trans_channels[0:2], dim=1)  # [(16,64,40,40),(16,32,40,40)]
    
    ## inject low-level global info to p4
    c5_half = self.reduce_layer_c5(c5)  # (16,64,20,20)
    p4_adjacent_info = self.LAF_p4([c3, c4, c5_half])  # (16,64,40,40)
    p4 = self.Inject_p4(p4_adjacent_info, low_global_info[0])  # (16,64,40,40)
    p4 = self.Rep_p4(p4)  # (16,64,40,40), 式(7)

    ## inject low-level global info to p3
    p4_half = self.reduce_layer_p4(p4)  # (16,32,40,40)
    p3_adjacent_info = self.LAF_p3([c2, c3, p4_half])  # (16,32,80,80)
    p3 = self.Inject_p3(p3_adjacent_info, low_global_info[1])  # (16,32,80,80)
    p3 = self.Rep_p3(p3)  # (16,32,80,80)

    # High-GD
    ## use transformer fusion global info
    high_align_feat = self.high_FAM([p3, p4, c5])  # (16,352,10,10)
    high_fuse_feat = self.high_IFM(high_align_feat)  # (16,352,10,10)
    high_fuse_feat = self.conv_1x1_n(high_fuse_feat)  # (16,192,10,10)
    high_global_info = high_fuse_feat.split(self.trans_channels[2:4], dim=1)  # [(16,64,10,10),(16,128,10,10)]

    ## inject low-level global info to n4
    n4_adjacent_info = self.LAF_n4(p3, p4_half)  # (16,64,40,40)
    n4 = self.Inject_n4(n4_adjacent_info, high_global_info[0])  # (16,64,40,40)
    n4 = self.Rep_n4(n4)  # (16,64,40,40)

    ## inject low-level global info to n5
    n5_adjacent_info = self.LAF_n5(n4, c5_half)  # (16,128,20,20)
    n5 = self.Inject_n5(n5_adjacent_info, high_global_info[1])  # (16,128,20,20)
    n5 = self.Rep_n5(n5)  # (16,128,20,20)

    outputs = [p3, n4, n5]  # [(16,32,80,80),(16,64,40,40),(16,128,20,20)]
    
    return outputs

首先是Low-GD,self.low_FAM的实现如下 

def forward(self, x):
    x_l, x_m, x_s, x_n = x
    # [(16,32,160,160),(16,64,80,80),(16,128,40,40),(16,256,20,20)]
    B, C, H, W = x_s.shape
    output_size = np.array([H, W])

    if torch.onnx.is_in_onnx_export():
        self.avg_pool = onnx_AdaptiveAvgPool2d

    x_l = self.avg_pool(x_l, output_size)
    x_m = self.avg_pool(x_m, output_size)
    x_n = F.interpolate(x_n, size=(H, W), mode='bilinear', align_corners=False)

    out = torch.cat([x_l, x_m, x_s, x_n], 1)  # (16,480,40,40)
    return out

self.low_IFM的实现如下,其中的block是RepVGGBlock

self.low_IFM = nn.Sequential(
        Conv(extra_cfg.fusion_in, extra_cfg.embed_dim_p, kernel_size=1, stride=1, padding=0),  # 480,96
        *[block(extra_cfg.embed_dim_p, extra_cfg.embed_dim_p) for _ in range(extra_cfg.fuse_block_num)],  # 3
        Conv(extra_cfg.embed_dim_p, sum(extra_cfg.trans_channels[0:2]), kernel_size=1, stride=1, padding=0),
)

接着通过split操作得到low_global_info,即inject模块的输入 \(F_{inj\_P3}\) 和 \(F_{inj\_P4}\)。

以上分别对应式(1)~式(3)

接下来是inject模块,self.LAF_p4的实现如下

def forward(self, x):
    N, C, H, W = x[1].shape
    output_size = (H, W)
    
    if torch.onnx.is_in_onnx_export():
        self.downsample = onnx_AdaptiveAvgPool2d
        output_size = np.array([H, W])
    
    x0 = self.downsample(x[0], output_size)
    x1 = self.cv1(x[1])
    x2 = F.interpolate(x[2], size=(H, W), mode='bilinear', align_corners=False)
    return self.cv_fuse(torch.cat((x0, x1, x2), dim=1))

self.Inject_p4的实现如下

def forward(self, x_l, x_g):
    '''
    x_g: global features
    x_l: local features
    '''
    B, C, H, W = x_l.shape
    g_B, g_C, g_H, g_W = x_g.shape
    use_pool = H < g_H
    
    local_feat = self.local_embedding(x_l)
    
    global_act = self.global_act(x_g)  # 式(4)
    global_feat = self.global_embedding(x_g)  # 式(5)
    
    if use_pool:
        avg_pool = get_avg_pool()
        output_size = np.array([H, W])
        
        sig_act = avg_pool(global_act, output_size)
        global_feat = avg_pool(global_feat, output_size)
    
    else:
        sig_act = F.interpolate(self.act(global_act), size=(H, W), mode='bilinear', align_corners=False)
        global_feat = F.interpolate(global_feat, size=(H, W), mode='bilinear', align_corners=False)
    
    out = local_feat * sig_act + global_feat  # 式(6)
    return out

然后self.Rep_p4对应式(7)。接下来inject to p3和p4的操作是一致的。

接下来是Hign-GD,输入是C5以及Low-GD的输出P3、P4。

self.high_FAM的实现如下

def forward(self, inputs):
    B, C, H, W = get_shape(inputs[-1])
    H = (H - 1) // self.stride + 1
    W = (W - 1) // self.stride + 1
    
    output_size = np.array([H, W])
    
    if not hasattr(self, 'pool'):
        self.pool = nn.functional.adaptive_avg_pool2d
    
    if torch.onnx.is_in_onnx_export():
        self.pool = onnx_AdaptiveAvgPool2d
    
    out = [self.pool(inp, output_size) for inp in inputs]
    
    return torch.cat(out, dim=1)

self.high_IFM采用transformer block,这里具体实现就不贴了。然后是1x1卷积接split操作。上面对应式(8)~式(10)

然后是inject模块,首先self.LAF_n4只融合相邻一层的特征

def forward(self, x1, x2):
    if torch.onnx.is_in_onnx_export():
        self.pool = onnx_AdaptiveAvgPool2d
    else:
        self.pool = nn.functional.adaptive_avg_pool2d
    
    N, C, H, W = x2.shape
    output_size = np.array([H, W])
    x1 = self.pool(x1, output_size)
    
    return torch.cat([x1, x2], 1)

接下来的self.Inject_n4和self.Rep_n4与low-GD中的self.Inject_p4和self.Rep_p4是一样的。

实验结果

GOLD-YOLO和其他YOLO的效果对比如下,可以看出主要提升在nano和small版本上。

  • 24
    点赞
  • 29
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

00000cj

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值