以下内容均以 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