YOLO11 改进 – 即插即用 _ RepViTBlock重参数化视觉Transformer块：结构重参数化技术破解训练推理效率瓶颈，实现精度与速度兼得
前言
本文介绍了轻量级卷积神经网络RepViT及其核心模块RepViTBlock在YOLOv11中的结合应用。RepViT通过融合轻量级ViTs的高效架构设计，对轻量级CNN进行“现代化”改造，采用结构重参数化技术，其核心的RepViT Block分离Token Mixer和Channel Mixer，提高了空间和通道信息利用效率。我们将RepViTBlock集成进YOLOv11，实验表明，RepViT在视觉任务中表现优于现有轻量级ViTs，与SAM结合后推理速度大幅提升，改进后的YOLOv11也取得了良好实验结果。
文章目录： YOLOv11改进大全：卷积层、轻量化、注意力机制、损失函数、Backbone、SPPF、Neck、检测头全方位优化汇总
专栏链接: YOLOv11改进专栏
文章目录
[TOC]
介绍
摘要
近期，轻量级视觉Transformer（Vision Transformers, ViTs）在资源受限的移动计算平台上展现出相较于传统轻量级卷积神经网络（CNNs）更优异的性能表现与更低的推理延迟。尽管已有研究揭示了轻量级ViTs与轻量级CNNs之间存在诸多结构性关联，但两者在模块级架构、宏观设计理念以及微观实现细节方面的显著差异性尚未得到系统性的深入探讨。本研究从视觉Transformer的技术视角重新审视轻量级CNN的高效架构设计范式，并着重强调其在移动设备应用场景中的巨大潜力。具体而言，通过系统融合轻量级ViTs的高效架构设计要素，逐步优化标准轻量级CNN（如MobileNetV3）的移动端适配特性，最终构建了一个全新的轻量级CNN架构家族，命名为 RepViT 。大量实验验证结果表明，RepViT在多种视觉任务中均显著优于现有的轻量级ViTs模型，并展现出理想的延迟性能指标。特别地，在ImageNet数据集上的测试中，RepViT首次实现了在iPhone 12移动设备上以1.0毫秒超低延迟达到超过80%的Top-1分类精度，成功突破了轻量级模型在精度与速度权衡方面的性能瓶颈。此外，将RepViT与SAM（Segment Anything Model）技术结合后，RepViT-SAM的推理速度相比先进的MobileSAM实现近10倍的显著提升。相关代码与预训练模型已全面开源，详细信息参见： https://github.com/THU-MIG/RepViT 。
文章链接
论文地址： 论文地址
代码地址： 代码地址
基本原理：
RepViT（Revisiting Mobile CNN From ViT Perspective）是一种新型的轻量级卷积神经网络（CNN），旨在通过引入轻量级视觉变换器（ViT）的高效架构设计，提升移动设备上的计算机视觉任务性能。RepViT的设计灵感来源于对现有轻量级CNN（如MobileNetV3）的“现代化”改造，结合了ViT的结构优势，形成了一种新的轻量级CNN家族。
RepViT的技术原理
-
架构设计 ： RepViT通过将轻量级ViT的高效设计理念融入到MobileNetV3的架构中，逐步增强其移动友好性。具体来说，RepViT采用了结构重参数化技术，使得模型在训练期间能够更好地学习，同时在推理时减少计算和内存开销。
-
RepViT Block ： RepViT Block是RepViT的核心构建模块，其设计包括以下几个关键要素：
-
分离的Token Mixer和Channel Mixer ：RepViT Block将token混合器和通道混合器分开处理，这样可以更有效地利用空间信息和通道信息。具体来说，深度可分离卷积（DW卷积）被放置在前面，以便更好地处理空间信息，而可选的挤压与激励（SE）层则在DW卷积之后，以增强模型的表达能力。
-
结构重参数化 ：通过结构重参数化技术，RepViT Block在推理时消除了与跳跃连接相关的计算和内存成本，这对于资源受限的移动设备尤为重要。
-
高效的卷积操作 ：RepViT Block使用了高效的卷积操作，确保在保持较低延迟的同时，能够实现较高的准确率。
-
-
性能表现 ： RepViT在多个计算机视觉任务中表现出色。例如，在ImageNet数据集上，RepViT能够在1.0毫秒的延迟下达到超过80%的Top-1准确率。此外，RepViT的最大模型RepViT-M2.3在仅2.3毫秒的延迟下获得了83.7%的准确率，显示出其在移动设备上的高效性。
-
与SAM的结合 ： RepViT还与Segment Anything Model（SAM）结合，形成RepViT-SAM模型，显著提高了在移动设备上的推理速度，几乎比现有的MobileSAM快10倍，同时保持了良好的零-shot迁移能力。
核心代码
class RepViTBlock(nn.Module):
def __init__(self, inp, hidden_dim, oup, kernel_size, stride, use_se, use_hs):
super(RepViTBlock, self).__init__()
assert stride in [1, 2]
self.identity = stride == 1 and inp == oup
assert(hidden_dim == 2 * inp)
if stride == 2:
self.token_mixer = nn.Sequential(
Conv2d_BN(inp, inp, kernel_size, stride, (kernel_size - 1) // 2, groups=inp),
SqueezeExcite(inp, 0.25) if use_se else nn.Identity(),
Conv2d_BN(inp, oup, ks=1, stride=1, pad=0)
)
self.channel_mixer = Residual(nn.Sequential(
# pw
Conv2d_BN(oup, 2 * oup, 1, 1, 0),
nn.GELU() if use_hs else nn.GELU(),
# pw-linear
Conv2d_BN(2 * oup, oup, 1, 1, 0, bn_weight_init=0),
))
else:
assert(self.identity)
self.token_mixer = nn.Sequential(
RepVGGDW(inp),
SqueezeExcite(inp, 0.25) if use_se else nn.Identity(),
)
self.channel_mixer = Residual(nn.Sequential(
# pw
Conv2d_BN(inp, hidden_dim, 1, 1, 0),
nn.GELU() if use_hs else nn.GELU(),
# pw-linear
Conv2d_BN(hidden_dim, oup, 1, 1, 0, bn_weight_init=0),
))
def forward(self, x):
return self.channel_mixer(self.token_mixer(x))
实验
脚本
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO
if __name__ == '__main__':
# 修改为自己的配置文件地址
model = YOLO('/root/ultralytics-main/ultralytics/cfg/models/11/yolov11-RepViTBlock.yaml')
# 修改为自己的数据集地址
model.train(data='/root/ultralytics-main/ultralytics/cfg/datasets/coco8.yaml',
cache=False,
imgsz=640,
epochs=10,
single_cls=False, # 是否是单类别检测
batch=8,
close_mosaic=10,
workers=0,
optimizer='SGD',
# amp=True,
project='runs/train',
name='RepViTBlock',
)
