学习资源站

25-替换主干网络之MobileViTv3（移动端轻量化网络的进一步升级）_yolov5 vit

YOLOv5改进系列（24）——替换主干网络之MobileViTv3（移动端轻量化网络的进一步升级）

🚀一、MobileViT v3介绍    

        在之前的研究中，CNN模型足够轻量化但是精准度有待提高，ViT模型具有较好的识别能力但是模型参数量大，计算复杂，都不能满足移动端实时高效检测的需求。

       MobileViT模型是2021年苹果公司提出的基于轻量化的ViT模型，该模型既具备ViT模型准确检测的优越性能，也具备CNN模型的轻量化优点，能极大程度上减少模型参数，对移动端友好，具备部署于移动设备的可能性。 

      MobileViT v3是该公司2022年9月推出的第3个版本，该模型相较于初始版本有以下四个改进：

  • 首先，将3×3卷积层替换为1×1卷积层；
  • 第二，将局部表示块和全局表示块的特征融合在一起；
  • 第三，在生成MobileViT Block输出之前，在融合块中添加输入特征作为最后一步；
  • 最后，在局部表示块中，将普通的3×3卷积层替换为深度3×3卷积层。

MobileViTv3网络结构图如下所示：

模型预测处理过程如下：

（1）将输入图像连接3×3标准卷积并做2倍下采样；之后通过5个MV2模块（如图（a）所示），其中步长为1的MV2模块进行特征提取，步长为2的MV2模块做2倍下采样；

（2）将得到的特征图间隔传入MobileViTV3 Block（如图（b）所示）和步长为2的MV2模块；

（3）接着使用3×3标准卷积进行通道压缩；

（4）最后进行全局平均池化来获取预测结果 。

MobileViT v3 Block模块是MobileViT v3核心部分，由局部表征模块、全局表征模块、融合模块三部分组成，具体介绍如下：

（1）局部表征模块 

对于输入的特征图像X\in R^{H\times W\times C}（H、W为图像的高、宽；C为输入特征图像的通道） ，首先采用一个3×3的深度卷积层和1×1的卷积层得到输出X_{L}\in R^{H\times W\times d}（d为输出特征图像的通道）。经过局部表征模块可以将特征图像X的局部空间信息映射到特定的维度d中。

（2）全局表征模块

将局部表征模块中的输出X_{L}作为全局表征模块的输入，然后将X_{L}展平为无重叠的图片块X_{U}\in R^{P\times N\times d}（P为展平图片块的大小，P=w\times h，w、h分别为展平图片块的宽、高，N为图片块的个数，N=W\times H/P）。随后，将每个图片块上相同位置上的像素p\in \left \{ 1,...,P \right \}送入Transformer模块中进行编码，得到X_{G}\in R^{P\times N\times d}。 为了避免图片块之间的信息丢失，在全局表征模块的最后会将编码后的所有图片块重组还原，得到X_{F}\in R^{H\times W\times d}，然后输出到下一个模块。 

（3）融合模块

将X_{F}送入特征融合模块中，使其映射到一个低维空间，得到X^{'}\in R^{H\times ×W\times ×C}。 然后通过拼接操作，将局部表征模块的输出X_{L}与X^{'}拼接起来，得到X^{''}\in R^{H\times ×W\times ×2C}。接着采用1×1的卷积层来融合局部特征X^{'}与全局特征X^{''}，得到X^{*}\in R^{H\times ×W\times ×2C}。最后把刚获取的X^{*}和原始X进行相加操作，得到输出Y\in R^{H\times W\times C}。

该过程需要伪代码可以私聊我~

🚀二、具体添加方法 

第①步：在common.py中添加MobileViT v3模块

首先，定义卷积层。

分为1×1卷积层和n×n（n=3）卷积层

  1. def conv_1x1_bn(inp, oup):
  2.     return nn.Sequential(
  3.         nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
  4.         nn.BatchNorm2d(oup),
  5.         nn.SiLU()
  6.     )
  8. def conv_nxn_bn(inp, oup, kernal_size=3, stride=1):
  9.     return nn.Sequential(
  10.         nn.Conv2d(inp, oup, kernal_size, stride, 1, bias=False),
  11.         nn.BatchNorm2d(oup),
  12.         nn.SiLU()
  13.     )

 接着，构造ViT模块。

Transformer Encoder模块中编码

  1. class PreNorm(nn.Module):
  2.     def __init__(self, dim, fn):
  3.         super().__init__()
  4.         self.norm = nn.LayerNorm(dim)
  5.         self.fn = fn # mg
  7.     def forward(self, x, **kwargs):
  8.         return self.fn(self.norm(x), **kwargs)
  10. class Attention(nn.Module):
  11.     def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
  12.         super().__init__()
  13.         inner_dim = dim_head *  heads
  14.         project_out = not (heads == 1 and dim_head == dim)
  16.         self.heads = heads
  17.         self.scale = dim_head ** -0.5
  19.         self.attend = nn.Softmax(dim = -1)
  20.         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
  22.         self.to_out = nn.Sequential(
  23.             nn.Linear(inner_dim, dim),
  24.             nn.Dropout(dropout)# mg
  25.         ) if project_out else nn.Identity()
  27.     def forward(self, x):
  28.         qkv = self.to_qkv(x).chunk(3, dim=-1)
  29.         q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h = self.heads), qkv)
  31.         dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
  32.         attn = self.attend(dots)
  33.         out = torch.matmul(attn, v)
  34.         out = rearrange(out, 'b p h n d -> b p n (h d)')
  35.         return self.to_out(out)
  37. class FeedForward(nn.Module):
  38.     def __init__(self, dim, hidden_dim, dropout=0.):
  39.         super().__init__()
  40.         self.net = nn.Sequential(
  41.             nn.Linear(dim, hidden_dim),
  42.             nn.SiLU(),
  43.             nn.Dropout(dropout),
  44.             nn.Linear(hidden_dim, dim),
  45.             nn.Dropout(dropout)
  46.         )
  48.     def forward(self, x):
  49.         return self.net(x)
  51. class MBTransformer(nn.Module):
  52.     def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
  53.         super().__init__()
  54.         self.layers = nn.ModuleList([])
  55.         for _ in range(depth):
  56.             self.layers.append(nn.ModuleList([
  57.                 PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
  58.                 PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
  59.             ]))
  61.     def forward(self, x):
  62.         for attn, ff in self.layers:
  63.             x = attn(x) + x
  64.             x = ff(x) + x
  65.         return x

 然后，MV2模块。

分为stride=1和stride=2两种。

  1. class MV2Block(nn.Module):
  2.     def __init__(self, inp, oup, stride=1, expansion=4):
  3.         super().__init__()
  4.         self.stride = stride
  5.         assert stride in [1, 2]
  7.         hidden_dim = int(inp * expansion)
  8.         self.use_res_connect = self.stride == 1 and inp == oup
  10.         if expansion == 1:
  11.             self.conv = nn.Sequential(
  12.                 # dw
  13.                 nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
  14.                 nn.BatchNorm2d(hidden_dim),
  15.                 nn.SiLU(),
  16.                 # pw-linear
  17.                 nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
  18.                 nn.BatchNorm2d(oup),
  19.             )
  20.         else:
  21.             self.conv = nn.Sequential(
  22.                 # pw
  23.                 nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
  24.                 nn.BatchNorm2d(hidden_dim),
  25.                 nn.SiLU(),
  26.                 # dw
  27.                 nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
  28.                 nn.BatchNorm2d(hidden_dim),
  29.                 nn.SiLU(),
  30.                 # pw-linear
  31.                 nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
  32.                 nn.BatchNorm2d(oup),
  33.             )
  35.     def forward(self, x):
  36.         if self.use_res_connect:
  37.             return x + self.conv(x)
  38.         else:
  39.             return self.conv(x)

 最后，核心模块 MobileViTv3_block。

介绍部分看上面就行~

  1. class MobileViTv3_block(nn.Module):
  2.     def __init__(self, channel, dim, depth=2, kernel_size=3, patch_size=(2, 2), mlp_dim=int(64*2), dropout=0.):
  3.         super().__init__()
  4.         self.ph, self.pw = patch_size
  5.         self.mv01 = MV2Block(channel, channel) 
  6.         self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
  7.         self.conv3 = conv_1x1_bn(dim, channel)
  8.         self.conv2 = conv_1x1_bn(channel, dim)
  9.         self.transformer = MBTransformer(dim, depth, 4, 8, mlp_dim, dropout)
  10.         self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)
  12.     def forward(self, x):
  13.         y = x.clone()
  14.         x = self.conv1(x)
  15.         x = self.conv2(x)
  16.         z = x.clone()
  17.         _, _, h, w = x.shape
  18.         x = rearrange(x, 'b d (h ph) (w pw) -> b (ph pw) (h w) d', ph=self.ph, pw=self.pw)
  19.         x = self.transformer(x)
  20.         x = rearrange(x, 'b (ph pw) (h w) d -> b d (h ph) (w pw)', h=h//self.ph, w=w//self.pw, ph=self.ph, pw=self.pw)
  21.         x = self.conv3(x)
  22.         x = torch.cat((x, z), 1)
  23.         x = self.conv4(x)
  24.         x = x + y
  25.         x = self.mv01(x)
  26.         return x

以下是完整代码：

将以下代码复制粘贴到common.py文件的末尾

  1. from einops import rearrange
  4. def conv_1x1_bn(inp, oup):
  5.     return nn.Sequential(
  6.         nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
  7.         nn.BatchNorm2d(oup),
  8.         nn.SiLU()
  9.     )
  11. def conv_nxn_bn(inp, oup, kernal_size=3, stride=1):
  12.     return nn.Sequential(
  13.         nn.Conv2d(inp, oup, kernal_size, stride, 1, bias=False),
  14.         nn.BatchNorm2d(oup),
  15.         nn.SiLU()
  16.     )
  18. class PreNorm(nn.Module):
  19.     def __init__(self, dim, fn):
  20.         super().__init__()
  21.         self.norm = nn.LayerNorm(dim)
  22.         self.fn = fn # mg
  24.     def forward(self, x, **kwargs):
  25.         return self.fn(self.norm(x), **kwargs)
  27. class Attention(nn.Module):
  28.     def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
  29.         super().__init__()
  30.         inner_dim = dim_head *  heads
  31.         project_out = not (heads == 1 and dim_head == dim)
  33.         self.heads = heads
  34.         self.scale = dim_head ** -0.5
  36.         self.attend = nn.Softmax(dim = -1)
  37.         self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
  39.         self.to_out = nn.Sequential(
  40.             nn.Linear(inner_dim, dim),
  41.             nn.Dropout(dropout)# mg
  42.         ) if project_out else nn.Identity()
  44.     def forward(self, x):
  45.         qkv = self.to_qkv(x).chunk(3, dim=-1)
  46.         q, k, v = map(lambda t: rearrange(t, 'b p n (h d) -> b p h n d', h = self.heads), qkv)
  48.         dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
  49.         attn = self.attend(dots)
  50.         out = torch.matmul(attn, v)
  51.         out = rearrange(out, 'b p h n d -> b p n (h d)')
  52.         return self.to_out(out)
  54. class FeedForward(nn.Module):
  55.     def __init__(self, dim, hidden_dim, dropout=0.):
  56.         super().__init__()
  57.         self.net = nn.Sequential(
  58.             nn.Linear(dim, hidden_dim),
  59.             nn.SiLU(),
  60.             nn.Dropout(dropout),
  61.             nn.Linear(hidden_dim, dim),
  62.             nn.Dropout(dropout)
  63.         )
  65.     def forward(self, x):
  66.         return self.net(x)
  68. class MBTransformer(nn.Module):
  69.     def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
  70.         super().__init__()
  71.         self.layers = nn.ModuleList([])
  72.         for _ in range(depth):
  73.             self.layers.append(nn.ModuleList([
  74.                 PreNorm(dim, Attention(dim, heads, dim_head, dropout)),
  75.                 PreNorm(dim, FeedForward(dim, mlp_dim, dropout))
  76.             ]))
  78.     def forward(self, x):
  79.         for attn, ff in self.layers:
  80.             x = attn(x) + x
  81.             x = ff(x) + x
  82.         return x
  84. class MV2Block(nn.Module):
  85.     def __init__(self, inp, oup, stride=1, expansion=4):
  86.         super().__init__()
  87.         self.stride = stride
  88.         assert stride in [1, 2]
  90.         hidden_dim = int(inp * expansion)
  91.         self.use_res_connect = self.stride == 1 and inp == oup
  93.         if expansion == 1:
  94.             self.conv = nn.Sequential(
  95.                 # dw
  96.                 nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
  97.                 nn.BatchNorm2d(hidden_dim),
  98.                 nn.SiLU(),
  99.                 # pw-linear
  100.                 nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
  101.                 nn.BatchNorm2d(oup),
  102.             )
  103.         else:
  104.             self.conv = nn.Sequential(
  105.                 # pw
  106.                 nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
  107.                 nn.BatchNorm2d(hidden_dim),
  108.                 nn.SiLU(),
  109.                 # dw
  110.                 nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
  111.                 nn.BatchNorm2d(hidden_dim),
  112.                 nn.SiLU(),
  113.                 # pw-linear
  114.                 nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
  115.                 nn.BatchNorm2d(oup),
  116.             )
  118.     def forward(self, x):
  119.         if self.use_res_connect:
  120.             return x + self.conv(x)
  121.         else:
  122.             return self.conv(x)
  124. class MobileViTv3_block(nn.Module):
  125.     def __init__(self, channel, dim, depth=2, kernel_size=3, patch_size=(2, 2), mlp_dim=int(64*2), dropout=0.):
  126.         super().__init__()
  127.         self.ph, self.pw = patch_size
  128.         self.mv01 = MV2Block(channel, channel) 
  129.         self.conv1 = conv_nxn_bn(channel, channel, kernel_size)
  130.         self.conv3 = conv_1x1_bn(dim, channel)
  131.         self.conv2 = conv_1x1_bn(channel, dim)
  132.         self.transformer = MBTransformer(dim, depth, 4, 8, mlp_dim, dropout)
  133.         self.conv4 = conv_nxn_bn(2 * channel, channel, kernel_size)
  135.     def forward(self, x):
  136.         y = x.clone()
  137.         x = self.conv1(x)
  138.         x = self.conv2(x)
  139.         z = x.clone()
  140.         _, _, h, w = x.shape
  141.         x = rearrange(x, 'b d (h ph) (w pw) -> b (ph pw) (h w) d', ph=self.ph, pw=self.pw)
  142.         x = self.transformer(x)
  143.         x = rearrange(x, 'b (ph pw) (h w) d -> b d (h ph) (w pw)', h=h//self.ph, w=w//self.pw, ph=self.ph, pw=self.pw)
  144.         x = self.conv3(x)
  145.         x = torch.cat((x, z), 1)
  146.         x = self.conv4(x)
  147.         x = x + y
  148.         x = self.mv01(x)
  149.         return x

第②步：修改yolo.py文件

再来修改yolo.py，在parse_model函数中找到 elif m is Concat: 语句，在其后面加上下面代码：

  1.         elif m in [MobileViTv3_block]:
  2.             c1, c2 = ch[f], args[0]
  3.             if c2 != no:  
  4.                 c2 = make_divisible(c2 * gw, 8)
  5.             args = [c1, c2]
  6.             if m in [MobileViTv3_block]:
  7.                 args.insert(2, n)  
  8.                 n = 1

 第③步：创建自定义的yaml文件  

yaml文件配置完整代码如下：

  1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
  3. # Parameters
  4. nc: 80  # number of classes
  5. depth_multiple: 0.33  # model depth multiple
  6. width_multiple: 0.50  # layer channel iscyy multiple
  7. anchors:
  8.   - [10,13, 16,30, 33,23]  # P3/8
  9.   - [30,61, 62,45, 59,119]  # P4/16
  10.   - [116,90, 156,198, 373,326]  # P5/32
  13. # YOLOv5 v6.0 backbone
  14. backbone:
  15.   # [from, number, module, args]
  16.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
  17.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
  18.    [-1, 3, C3, [128]],
  19.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
  20.    [-1, 6, MobileViTv3_block, [256]],
  21.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
  22.    [-1, 9, C3, [512]],
  23.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
  24.    [-1, 3, C3, [1024]],
  25.    [-1, 1, SPPF, [1024, 5]],  # 9
  26.   ]
  27. # YOLOv5 v6.0 head
  28. head:
  29.   [[-1, 1, Conv, [512, 1, 1]],
  30.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  31.    [[-1, 6], 1, Concat, [1]],  # cat backbone P4
  32.    [-1, 3, C3, [512, False]],  # 13
  34.    [-1, 1, Conv, [256, 1, 1]],
  35.    [-1, 1, nn.Upsample, [None, 2, 'nearest']],
  36.    [[-1, 4], 1, Concat, [1]],  # cat backbone P3
  37.    [-1, 3, C3, [256, False]],  # 17 (P3/8-small)
  39.    [-1, 1, Conv, [256, 3, 2]],
  40.    [[-1, 14], 1, Concat, [1]],  # cat head P4
  41.    [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)
  43.    [-1, 1, Conv, [512, 3, 2]],
  44.    [[-1, 10], 1, Concat, [1]],  # cat head P5
  45.    [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
  47.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  48.   ]

第④步 验证是否加入成功

运行yolo.py

 这样就OK啦~