YOLOv10优化：backbone改进 | 轻量化之王MobileNetV4 开源 | Top-1 精度 87%，手机推理速度 3.8ms，原地起飞！

原创

AI小怪兽

发布于 2024-06-14 13:01:11

3310

发布于 2024-06-14 13:01:11

💡💡💡创新点：轻量化之王MobileNetV4 开源 | Top-1 精度 87%，手机推理速度 3.8ms，原地起飞！

MobileNetV4（MNv4），其特点是针对移动设备设计的通用高效架构。创新1）：引入了通用倒瓶颈（UIB）搜索块，这是一个统一且灵活的结构，它融合了倒瓶颈（IB）、ConvNext、前馈网络（FFN）以及一种新颖的额外深度可分（ExtraDW）变体；创新2）：一种优化的神经结构搜索(NAS)配方，提高了MNv4的搜索效率；创新3）：为了进一步提升准确度，引入了一种新颖的蒸馏技术。

1.YOLOv10介绍

添加描述

论文： https://arxiv.org/pdf/2405.14458

代码： GitHub - THU-MIG/yolov10: YOLOv10: Real-Time End-to-End Object Detection

摘要：在过去的几年里，由于其在计算成本和检测性能之间的有效平衡，YOLOS已经成为实时目标检测领域的主导范例。研究人员已经探索了YOLOS的架构设计、优化目标、数据增强策略等，并取得了显著进展。然而，对用于后处理的非最大抑制（NMS）的依赖妨碍了YOLOS的端到端部署，并且影响了推理延迟。此外，YOLOS中各部件的设计缺乏全面和彻底的检查，导致明显的计算冗余，限制了模型的性能。这导致次优的效率，以及相当大的性能改进潜力。在这项工作中，我们的目标是从后处理和模型架构两个方面进一步推进YOLOS的性能-效率边界。为此，我们首先提出了用于YOLOs无NMS训练的持续双重分配，该方法带来了有竞争力的性能和低推理延迟。此外，我们还介绍了YOLOS的整体效率-精度驱动模型设计策略。我们从效率和精度两个角度对YOLOS的各个组件进行了全面优化，大大降低了计算开销，增强了性能。我们努力的成果是用于实时端到端对象检测的新一代YOLO系列，称为YOLOV10。广泛的实验表明，YOLOV10在各种模型规模上实现了最先进的性能和效率。例如，在COCO上的类似AP下，我们的YOLOV10-S比RT-DETR-R18快1.8倍，同时具有2.8倍更少的参数和FLOPS。与YOLOV9-C相比，YOLOV10-B在性能相同的情况下，延迟减少了46%，参数减少了25%。

1.YOLOv10介绍

论文： https://arxiv.org/pdf/2405.14458

代码： GitHub - THU-MIG/yolov10: YOLOv10: Real-Time End-to-End Object Detection

1.1 C2fUIB介绍

为了解决这个问题，我们提出了一种基于秩的块设计方案，旨在通过紧凑的架构设计降低被证明是冗余的阶段复杂度。我们首先提出了一个紧凑的倒置块（CIB）结构，它采用廉价的深度可分离卷积进行空间混合，以及成本效益高的点对点卷积进行通道混合

C2fUIB只是用CIB结构替换了YOLOv8中 C2f的Bottleneck结构

实现代码ultralytics/nn/modules/block.py

class CIB(nn.Module):
    """Standard bottleneck."""

    def __init__(self, c1, c2, shortcut=True, e=0.5, lk=False):
        """Initializes a bottleneck module with given input/output channels, shortcut option, group, kernels, and
        expansion.
        """
        super().__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = nn.Sequential(
            Conv(c1, c1, 3, g=c1),
            Conv(c1, 2 * c_, 1),
            Conv(2 * c_, 2 * c_, 3, g=2 * c_) if not lk else RepVGGDW(2 * c_),
            Conv(2 * c_, c2, 1),
            Conv(c2, c2, 3, g=c2),
        )

        self.add = shortcut and c1 == c2

    def forward(self, x):
        """'forward()' applies the YOLO FPN to input data."""
        return x + self.cv1(x) if self.add else self.cv1(x)

class C2fCIB(C2f):
    """Faster Implementation of CSP Bottleneck with 2 convolutions."""

    def __init__(self, c1, c2, n=1, shortcut=False, lk=False, g=1, e=0.5):
        """Initialize CSP bottleneck layer with two convolutions with arguments ch_in, ch_out, number, shortcut, groups,
        expansion.
        """
        super().__init__(c1, c2, n, shortcut, g, e)
        self.m = nn.ModuleList(CIB(self.c, self.c, shortcut, e=1.0, lk=lk) for _ in range(n))

1.2 PSA介绍

具体来说，我们在1×1卷积后将特征均匀地分为两部分。我们只将一部分输入到由多头自注意力模块（MHSA）和前馈网络（FFN）组成的NPSA块中。然后，两部分通过1×1卷积连接并融合。此外，遵循将查询和键的维度分配为值的一半，并用BatchNorm替换LayerNorm以实现快速推理。

实现代码ultralytics/nn/modules/block.py

class Attention(nn.Module):
    def __init__(self, dim, num_heads=8,
                 attn_ratio=0.5):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.key_dim = int(self.head_dim * attn_ratio)
        self.scale = self.key_dim ** -0.5
        nh_kd = nh_kd = self.key_dim * num_heads
        h = dim + nh_kd * 2
        self.qkv = Conv(dim, h, 1, act=False)
        self.proj = Conv(dim, dim, 1, act=False)
        self.pe = Conv(dim, dim, 3, 1, g=dim, act=False)

    def forward(self, x):
        B, _, H, W = x.shape
        N = H * W
        qkv = self.qkv(x)
        q, k, v = qkv.view(B, self.num_heads, -1, N).split([self.key_dim, self.key_dim, self.head_dim], dim=2)

        attn = (
            (q.transpose(-2, -1) @ k) * self.scale
        )
        attn = attn.softmax(dim=-1)
        x = (v @ attn.transpose(-2, -1)).view(B, -1, H, W) + self.pe(v.reshape(B, -1, H, W))
        x = self.proj(x)
        return x

class PSA(nn.Module):

    def __init__(self, c1, c2, e=0.5):
        super().__init__()
        assert(c1 == c2)
        self.c = int(c1 * e)
        self.cv1 = Conv(c1, 2 * self.c, 1, 1)
        self.cv2 = Conv(2 * self.c, c1, 1)
        
        self.attn = Attention(self.c, attn_ratio=0.5, num_heads=self.c // 64)
        self.ffn = nn.Sequential(
            Conv(self.c, self.c*2, 1),
            Conv(self.c*2, self.c, 1, act=False)
        )
        
    def forward(self, x):
        a, b = self.cv1(x).split((self.c, self.c), dim=1)
        b = b + self.attn(b)
        b = b + self.ffn(b)
        return self.cv2(torch.cat((a, b), 1))

1.3 SCDown

OLOs通常利用常规的3×3标准卷积，步长为2，同时实现空间下采样（从H×W到H/2×W/2）和通道变换（从C到2C）。这引入了不可忽视的计算成本O(9HWC^2)和参数数量O(18C^2)。相反，我们提议将空间缩减和通道增加操作解耦，以实现更高效的下采样。具体来说，我们首先利用点对点卷积来调整通道维度，然后利用深度可分离卷积进行空间下采样。这将计算成本降低到O(2HWC^2 + 9HWC)，并将参数数量减少到O(2C^2 + 18C)。同时，它最大限度地保留了下采样过程中的信息，从而在减少延迟的同时保持了有竞争力的性能。

实现代码ultralytics/nn/modules/block.py

class SCDown(nn.Module):
    def __init__(self, c1, c2, k, s):
        super().__init__()
        self.cv1 = Conv(c1, c2, 1, 1)
        self.cv2 = Conv(c2, c2, k=k, s=s, g=c2, act=False)

    def forward(self, x):
        return self.cv2(self.cv1(x))

2.原理介绍

论文： https://arxiv.org/pdf/2404.10518

摘要：我们介绍了最新一代的mobilenet，被称为MobileNetV4 (MNv4)，具有普遍有效的移动设备架构设计。在其核心，我们引入了通用倒瓶颈(UIB)搜索块，这是一个统一而灵活的结构，它融合了倒瓶颈(IB)， ConvNext，前馈网络(FFN)和一个新的Extra depth(引渡)变体。除了UIB之外，我们还推出了Mobile MQA，这是一款专为移动加速器量身定制的注意力块，可显著提高39%的速度。介绍了一种优化的神经结构搜索(NAS)配方，提高了MNv4的搜索效率。UIB, Mobile MQA和精致的NAS配方的集成产生了一套新的MNv4模型，这些模型在移动cpu, dsp, gpu以及专用加速器(如Apple Neural Engine和Google Pixel EdgeTPU)上大多是最优的，这是任何其他模型测试中没有发现的特征。最后，为了进一步提高精度，我们介绍了一种新的蒸馏技术。通过这种技术的增强，我们的MNv4-Hybrid-Large模型提供了87%的ImageNet-1K精度，Pixel 8 EdgeTPU运行时间仅为3.8ms。

3. mobilenetv4加入YOLOv10

3.1 新建ultralytics/nn/backbone/mobilenetv4.py

核心代码：

class UniversalInvertedBottleneckBlock(nn.Module):
    def __init__(self, 
            inp, 
            oup, 
            start_dw_kernel_size, 
            middle_dw_kernel_size, 
            middle_dw_downsample,
            stride,
            expand_ratio
        ):
        super().__init__()
        # Starting depthwise conv.
        self.start_dw_kernel_size = start_dw_kernel_size
        if self.start_dw_kernel_size:            
            stride_ = stride if not middle_dw_downsample else 1
            self._start_dw_ = conv_2d(inp, inp, kernel_size=start_dw_kernel_size, stride=stride_, groups=inp, act=False)
        # Expansion with 1x1 convs.
        expand_filters = make_divisible(inp * expand_ratio, 8)
        self._expand_conv = conv_2d(inp, expand_filters, kernel_size=1)
        # Middle depthwise conv.
        self.middle_dw_kernel_size = middle_dw_kernel_size
        if self.middle_dw_kernel_size:
            stride_ = stride if middle_dw_downsample else 1
            self._middle_dw = conv_2d(expand_filters, expand_filters, kernel_size=middle_dw_kernel_size, stride=stride_, groups=expand_filters)
        # Projection with 1x1 convs.
        self._proj_conv = conv_2d(expand_filters, oup, kernel_size=1, stride=1, act=False)
        
        # Ending depthwise conv.
        # this not used
        # _end_dw_kernel_size = 0
        # self._end_dw = conv_2d(oup, oup, kernel_size=_end_dw_kernel_size, stride=stride, groups=inp, act=False)
        
    def forward(self, x):
        if self.start_dw_kernel_size:
            x = self._start_dw_(x)
            # print("_start_dw_", x.shape)
        x = self._expand_conv(x)
        # print("_expand_conv", x.shape)
        if self.middle_dw_kernel_size:
            x = self._middle_dw(x)
            # print("_middle_dw", x.shape)
        x = self._proj_conv(x)
        # print("_proj_conv", x.shape)
        return x

def build_blocks(layer_spec):
    if not layer_spec.get('block_name'):
        return nn.Sequential()
    block_names = layer_spec['block_name']
    layers = nn.Sequential()
    if block_names == "convbn":
        schema_ = ['inp', 'oup', 'kernel_size', 'stride']
        args = {}
        for i in range(layer_spec['num_blocks']):
            args = dict(zip(schema_, layer_spec['block_specs'][i]))
            layers.add_module(f"convbn_{i}", conv_2d(**args))
    elif block_names == "uib":
        schema_ =  ['inp', 'oup', 'start_dw_kernel_size', 'middle_dw_kernel_size', 'middle_dw_downsample', 'stride', 'expand_ratio']
        args = {}
        for i in range(layer_spec['num_blocks']):
            args = dict(zip(schema_, layer_spec['block_specs'][i]))
            layers.add_module(f"uib_{i}", UniversalInvertedBottleneckBlock(**args))
    elif block_names == "fused_ib":
        schema_ = ['inp', 'oup', 'stride', 'expand_ratio', 'act']
        args = {}
        for i in range(layer_spec['num_blocks']):
            args = dict(zip(schema_, layer_spec['block_specs'][i]))
            layers.add_module(f"fused_ib_{i}", InvertedResidual(**args))
    else:
        raise NotImplementedError
    return layers


class MobileNetV4(nn.Module):
    def __init__(self, model):
        # MobileNetV4ConvSmall  MobileNetV4ConvMedium  MobileNetV4ConvLarge
        # MobileNetV4HybridMedium  MobileNetV4HybridLarge
        """Params to initiate MobilenNetV4
        Args:
            model : support 5 types of models as indicated in 
            "https://github.com/tensorflow/models/blob/master/official/vision/modeling/backbones/mobilenet.py"        
        """
        super().__init__()
        assert model in MODEL_SPECS.keys()
        self.model = model
        self.spec = MODEL_SPECS[self.model]
       
        # conv0
        self.conv0 = build_blocks(self.spec['conv0'])
        # layer1
        self.layer1 = build_blocks(self.spec['layer1'])
        # layer2
        self.layer2 = build_blocks(self.spec['layer2'])
        # layer3
        self.layer3 = build_blocks(self.spec['layer3'])
        # layer4
        self.layer4 = build_blocks(self.spec['layer4'])
        # layer5   
        self.layer5 = build_blocks(self.spec['layer5'])
        self.features = nn.ModuleList([self.conv0, self.layer1, self.layer2, self.layer3, self.layer4, self.layer5])     
        self.channel = [i.size(1) for i in self.forward(torch.randn(1, 3, 640, 640))]
        
    def forward(self, x):
        input_size = x.size(2)
        scale = [4, 8, 16, 32]
        features = [None, None, None, None]
        for f in self.features:
            x = f(x)
            if input_size // x.size(2) in scale:
                features[scale.index(input_size // x.size(2))] = x
        return features

3.2 yolov10n-mobilenetv4.yaml

# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  n: [0.33, 0.25, 1024] 

# YOLOv8.0n backbone
backbone:
  # [from, repeats, module, args]
  - [-1, 1, MobileNetV4ConvSmall , []]  # 4
  - [-1, 1, SPPF, [1024, 5]] # 5
  - [-1, 1, PSA, [1024]] # 6

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  - [-1, 3, C2f, [512]] # 9

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  - [-1, 3, C2f, [256]] # 12 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 9], 1, Concat, [1]] # cat head P4
  - [-1, 3, C2f, [512]] # 15 (P4/16-medium)

  - [-1, 1, SCDown, [512, 3, 2]]
  - [[-1, 6], 1, Concat, [1]] # cat head P5
  - [-1, 3, C2fCIB, [1024, True, True]] # 18 (P5/32-large)

  - [[12, 15, 18], 1, v10Detect, [nc]] # Detect(P3, P4, P5)

3.3 yolov10s-mobilenetv4.yaml

# Parameters
nc: 80 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
  # [depth, width, max_channels]
  s: [0.33, 0.50, 1024]

backbone:
  # [from, repeats, module, args]
  - [-1, 1, MobileNetV4ConvSmall , []]  # 4
  - [-1, 1, SPPF, [1024, 5]] # 5
  - [-1, 1, PSA, [1024]] # 6

# YOLOv8.0n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 3], 1, Concat, [1]] # cat backbone P4
  - [-1, 3, C2f, [512]] # 9

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 2], 1, Concat, [1]] # cat backbone P3
  - [-1, 3, C2f, [256]] # 12 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 9], 1, Concat, [1]] # cat head P4
  - [-1, 3, C2f, [512]] # 15 (P4/16-medium)

  - [-1, 1, SCDown, [512, 3, 2]]
  - [[-1, 6], 1, Concat, [1]] # cat head P5
  - [-1, 3, C2fCIB, [1024, True, True]] # 18 (P5/32-large)

  - [[12, 15, 18], 1, v10Detect, [nc]] # Detect(P3, P4, P5)