YOLOV5(7)——Swintransformer(无痛版）

以下内容均以 YOLOV5的7.0版本为主

以下为配置文件 yolov5s_Swin.yaml 复制下面的内容到models

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32


# 当前结构数量:max(round(number * depth_multiple), 1)(Conv数量   C3内部的残差数量)
# args :作为参数传入类中(内部元素可能会做调整)
# 1.Conv/C3/SPPF：输入维度：ch[-1] (ch 会储存上一个结构的输出维度)    输出维度：args[0] * width_multiple 
         # eg:args [64,6,2,2]  => [3,32,6,2,2]    Conv2d(3,32,6,2,2)  in,out,k,s,p   (若没有p,则为k//2)

# 2.Concat：输入维度 : c2 = sum(ch[x] for x in from)   args:[1] 表明 torch.cat(x, 1)在第一维上拼接

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]  # (640,640)
  [[-1, 1, Swin_Transformer_block, [192,1]],    # 0    [2, 80, 80, 192]  /8    
   [-1, 1, Swin_Transformer_block, [384,2]],    # 1    [2, 40, 40, 384]  /16
   [-1, 3, Swin_Transformer_block, [768,3]],    # 2    [2, 20, 20, 768]  /32              
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, New_Conv, [512, 1, 1]],                     # 3 (20,20)
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],     # 4 (40,40)
   [[-1, 1], 1, New_Concat, [1]],               #    5 (40,40)
   [-1, 3, C3, [512, False]],                      # 6 (40,40)                 csp2_x 无残差

   [-1, 1, Conv, [256, 1, 1]],                     # 7 (40,40)
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],     # 8 (80,80)
   [[-1, 0], 1, New_Concat, [1]],  #                 9 (80,80)
   [-1, 3, C3, [256, False]],  #    (P3/8-small)     10 (80,80)

   [-1, 1, Conv, [256, 3, 2]],                     # 11 (40,40)
   [[-1, 7], 1, Concat, [1]],  # cat head P4        12 (40,40)
   [-1, 3, C3, [512, False]],  #  (P4/16-medium)     13 (40,40)
   
   [-1, 1, Conv, [512, 3, 2]],                     # 14 (20,20) 
   [[-1, 3], 1, Concat, [1]],  # cat head P5       # 15 (20,20)
   [-1, 3, C3, [1024, False]],  #    (P5/32-large) # 16 (20,20)

   [[10, 13, 16], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

将以下内容写入 common.py

import torchvision.models as models
swin_t=models.swin_transformer.swin_t(pretrained=True,progress=True)

#  利用 torchvision.models 构建 swin_t 

class Swin_Transformer_block(nn.Module):
    def __init__(self,slice):
        super().__init__()
        if(slice==1):
            self.model=nn.Sequential(
                swin_t.features[0:4] #  [2, 80, 80, 192]  /8
            )
        elif(slice==2):
            self.model=nn.Sequential(
                swin_t.features[4:6] # [2, 40, 40, 384]  /16
            )
        elif(slice==3):
            self.model=nn.Sequential(
                swin_t.features[6:8] # [2, 20, 20, 768]  /32
            )
    def forward(self,x):
        return self.model(x)

将以下内容复制到yolo.py中

        elif m is Swin_Transformer_block:
            c2=args[0]
            args=[*args[1:]]

此时运行yolo.py，查看是否配置完成（别忘记改cfg路径），成功则配置完成，此时的yolov5 backbone已经为 swintransformer