devbox
Monodyee
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

改进YOLO：YOLOv8结合swin transformer

作者：Monodyee | 2024-04-06 10:46:29

赞

踩

改进YOLO：YOLOv8结合swin transformer

目录

1、修改yaml文件

2、添加 SwinTransformer.py

3、修改 tasks.py

4、根目录增加文件

1、修改yaml文件

修改 ultralytics/cfg/models/v8/yolov8.yaml

  1. backbone:
  2.   # [from, repeats, module, args]
  3.   - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
  4.   - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
  5.   - [-1, 3, C2f, [128, True]]
  6.   - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
  7.   # - [-1, 6, C2f, [256, True]]
  8.   - [-1, 6, SwinTransformer, [256, True]]
  9.   - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
  10.   - [-1, 6, C2f, [512, True]]
  11.   - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
  12.   - [-1, 3, C2f, [1024, True]]
  13.   - [-1, 1, SPPF, [1024, 5]]  # 9

C2f 那一行，替换为 SwinTransformer

2、添加 SwinTransformer.py

在 ultralytics/nn 下新增该文件

  1. import torch
  2. import torch.nn as nn
  3. # from .conv import Conv
  4. from ultralytics.nn.modules.conv import Conv
  5. import torch.nn.functional as F
  6.  
  7. from timm.models.layers import DropPath as TimmDropPath
  8. from timm.models.layers import trunc_normal_
  9. class DropPath(TimmDropPath):
  10.  
  11.     def __init__(self, drop_prob=None):
  12.         super().__init__(drop_prob=drop_prob)
  13.         self.drop_prob = drop_prob
  14.  
  15.     def __repr__(self):
  16.         msg = super().__repr__()
  17.         msg += f'(drop_prob={self.drop_prob})'
  18.         return msg
  19.  
  20. class WindowAttention(nn.Module):
  21.  
  22.     def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
  23.  
  24.         super().__init__()
  25.         self.dim = dim
  26.         self.window_size = window_size  # Wh, Ww
  27.         self.num_heads = num_heads
  28.         head_dim = dim // num_heads
  29.         self.scale = qk_scale or head_dim ** -0.5
  30.  
  31.         # define a parameter table of relative position bias
  32.         self.relative_position_bias_table = nn.Parameter(
  33.             torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
  34.  
  35.         # get pair-wise relative position index for each token inside the window
  36.         coords_h = torch.arange(self.window_size[0])
  37.         coords_w = torch.arange(self.window_size[1])
  38.         coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
  39.         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
  40.         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
  41.         relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
  42.         relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
  43.         relative_coords[:, :, 1] += self.window_size[1] - 1
  44.         relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
  45.         relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
  46.         self.register_buffer("relative_position_index", relative_position_index)
  47.  
  48.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
  49.         self.attn_drop = nn.Dropout(attn_drop)
  50.         self.proj = nn.Linear(dim, dim)
  51.         self.proj_drop = nn.Dropout(proj_drop)
  52.  
  53.         nn.init.normal_(self.relative_position_bias_table, std=.02)
  54.         self.softmax = nn.Softmax(dim=-1)
  55.  
  56.     def forward(self, x, mask=None):
  57.  
  58.         B_, N, C = x.shape
  59.         qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
  60.         q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
  61.  
  62.         q = q * self.scale
  63.         attn = (q @ k.transpose(-2, -1))
  64.  
  65.         relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
  66.             self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
  67.         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
  68.         attn = attn + relative_position_bias.unsqueeze(0)
  69.  
  70.         if mask is not None:
  71.             nW = mask.shape[0]
  72.             attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
  73.             attn = attn.view(-1, self.num_heads, N, N)
  74.             attn = self.softmax(attn)
  75.         else:
  76.             attn = self.softmax(attn)
  77.  
  78.         attn = self.attn_drop(attn)
  79.  
  80.         # print(attn.dtype, v.dtype)
  81.         try:
  82.             x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
  83.         except:
  84.             # print(attn.dtype, v.dtype)
  85.             x = (attn.half() @ v).transpose(1, 2).reshape(B_, N, C)
  86.         x = self.proj(x)
  87.         x = self.proj_drop(x)
  88.         return x
  89.  
  90.  
  91. class SwinTransformer(nn.Module):
  92.     # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
  93.     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
  94.         super(SwinTransformer, self).__init__()
  95.         c_ = int(c2 * e)  # hidden channels
  96.         self.cv1 = Conv(c1, c_, 1, 1)
  97.         self.cv2 = Conv(c1, c_, 1, 1)
  98.         self.cv3 = Conv(2 * c_, c2, 1, 1)
  99.         num_heads = c_ // 32
  100.         self.m = SwinTransformerBlock(c_, c_, num_heads, n)
  101.         # self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  102.  
  103.     def forward(self, x):
  104.         y1 = self.m(self.cv1(x))
  105.         y2 = self.cv2(x)
  106.         return self.cv3(torch.cat((y1, y2), dim=1))
  107.  
  108.  
  109. class SwinTransformerB(nn.Module):
  110.     # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
  111.     def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
  112.         super(SwinTransformerB, self).__init__()
  113.         c_ = int(c2)  # hidden channels
  114.         self.cv1 = Conv(c1, c_, 1, 1)
  115.         self.cv2 = Conv(c_, c_, 1, 1)
  116.         self.cv3 = Conv(2 * c_, c2, 1, 1)
  117.         num_heads = c_ // 32
  118.         self.m = SwinTransformerBlock(c_, c_, num_heads, n)
  119.         # self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  120.  
  121.     def forward(self, x):
  122.         x1 = self.cv1(x)
  123.         y1 = self.m(x1)
  124.         y2 = self.cv2(x1)
  125.         return self.cv3(torch.cat((y1, y2), dim=1))
  126.  
  127.  
  128. class SwinTransformerC(nn.Module):
  129.     # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
  130.     def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
  131.         super(SwinTransformerC, self).__init__()
  132.         c_ = int(c2 * e)  # hidden channels
  133.         self.cv1 = Conv(c1, c_, 1, 1)
  134.         self.cv2 = Conv(c1, c_, 1, 1)
  135.         self.cv3 = Conv(c_, c_, 1, 1)
  136.         self.cv4 = Conv(2 * c_, c2, 1, 1)
  137.         num_heads = c_ // 32
  138.         self.m = SwinTransformerBlock(c_, c_, num_heads, n)
  139.         # self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
  140.  
  141.     def forward(self, x):
  142.         y1 = self.cv3(self.m(self.cv1(x)))
  143.         y2 = self.cv2(x)
  144.         return self.cv4(torch.cat((y1, y2), dim=1))
  145.  
  146.  
  147. class Mlp(nn.Module):
  148.  
  149.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.):
  150.         super().__init__()
  151.         out_features = out_features or in_features
  152.         hidden_features = hidden_features or in_features
  153.         self.fc1 = nn.Linear(in_features, hidden_features)
  154.         self.act = act_layer()
  155.         self.fc2 = nn.Linear(hidden_features, out_features)
  156.         self.drop = nn.Dropout(drop)
  157.  
  158.     def forward(self, x):
  159.         x = self.fc1(x)
  160.         x = self.act(x)
  161.         x = self.drop(x)
  162.         x = self.fc2(x)
  163.         x = self.drop(x)
  164.         return x
  165.  
  166.  
  167. def window_partition(x, window_size):
  168.     B, H, W, C = x.shape
  169.     assert H % window_size == 0, 'feature map h and w can not divide by window size'
  170.     x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
  171.     windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
  172.     return windows
  173.  
  174.  
  175. def window_reverse(windows, window_size, H, W):
  176.     B = int(windows.shape[0] / (H * W / window_size / window_size))
  177.     x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
  178.     x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
  179.     return x
  180.  
  181.  
  182. class SwinTransformerLayer(nn.Module):
  183.  
  184.     def __init__(self, dim, num_heads, window_size=8, shift_size=0,
  185.                  mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
  186.                  act_layer=nn.SiLU, norm_layer=nn.LayerNorm):
  187.         super().__init__()
  188.         self.dim = dim
  189.         self.num_heads = num_heads
  190.         self.window_size = window_size
  191.         self.shift_size = shift_size
  192.         self.mlp_ratio = mlp_ratio
  193.         # if min(self.input_resolution) <= self.window_size:
  194.         #     # if window size is larger than input resolution, we don't partition windows
  195.         #     self.shift_size = 0
  196.         #     self.window_size = min(self.input_resolution)
  197.         assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
  198.         self.norm1 = norm_layer(dim)
  199.         self.attn = WindowAttention(
  200.             dim, window_size=(self.window_size, self.window_size), num_heads=num_heads,
  201.             qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
  202.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
  203.         self.norm2 = norm_layer(dim)
  204.         mlp_hidden_dim = int(dim * mlp_ratio)
  205.         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
  206.     def create_mask(self, H, W):
  207.         # calculate attention mask for SW-MSA
  208.         img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
  209.         h_slices = (slice(0, -self.window_size),
  210.                     slice(-self.window_size, -self.shift_size),
  211.                     slice(-self.shift_size, None))
  212.         w_slices = (slice(0, -self.window_size),
  213.                     slice(-self.window_size, -self.shift_size),
  214.                     slice(-self.shift_size, None))
  215.         cnt = 0
  216.         for h in h_slices:
  217.             for w in w_slices:
  218.                 img_mask[:, h, w, :] = cnt
  219.                 cnt += 1
  220.         mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
  221.         mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
  222.         attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
  223.         attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
  224.         return attn_mask
  225.     def forward(self, x):
  226.         # reshape x[b c h w] to x[b l c]
  227.         _, _, H_, W_ = x.shape
  228.         Padding = False
  229.         if min(H_, W_) < self.window_size or H_ % self.window_size != 0 or W_ % self.window_size != 0:
  230.             Padding = True
  231.             # print(f'img_size {min(H_, W_)} is less than (or not divided by) window_size {self.window_size}, Padding.')
  232.             pad_r = (self.window_size - W_ % self.window_size) % self.window_size
  233.             pad_b = (self.window_size - H_ % self.window_size) % self.window_size
  234.             x = F.pad(x, (0, pad_r, 0, pad_b))
  235.         # print('2', x.shape)
  236.         B, C, H, W = x.shape
  237.         L = H * W
  238.         x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C)  # b, L, c
  239.         # create mask from init to forward
  240.         if self.shift_size > 0:
  241.             attn_mask = self.create_mask(H, W).to(x.device)
  242.         else:
  243.             attn_mask = None
  244.         shortcut = x
  245.         x = self.norm1(x)
  246.         x = x.view(B, H, W, C)
  247.         # cyclic shift
  248.         if self.shift_size > 0:
  249.             shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
  250.         else:
  251.             shifted_x = x
  252.         # partition windows
  253.         x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
  254.         x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
  255.         # W-MSA/SW-MSA
  256.         attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
  257.         # merge windows
  258.         attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
  259.         shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
  260.         # reverse cyclic shift
  261.         if self.shift_size > 0:
  262.             x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
  263.         else:
  264.             x = shifted_x
  265.         x = x.view(B, H * W, C)
  266.         # FFN
  267.         x = shortcut + self.drop_path(x)
  268.         x = x + self.drop_path(self.mlp(self.norm2(x)))
  269.         x = x.permute(0, 2, 1).contiguous().view(-1, C, H, W)  # b c h w
  270.         if Padding:
  271.             x = x[:, :, :H_, :W_]  # reverse padding
  272.         return x
  273. class SwinTransformerBlock(nn.Module):
  274.     def __init__(self, c1, c2, num_heads, num_layers, window_size=8):
  275.         super().__init__()
  276.         self.conv = None
  277.         if c1 != c2:
  278.             self.conv = Conv(c1, c2)
  279.         # remove input_resolution
  280.         self.blocks = nn.Sequential(*[SwinTransformerLayer(dim=c2, num_heads=num_heads, window_size=window_size,
  281.                                                            shift_size=0 if (i % 2 == 0) else window_size // 2) for i in
  282.                                       range(num_layers)])
  283.     def forward(self, x):
  284.         if self.conv is not None:
  285.             x = self.conv(x)
  286.         x = self.blocks(x)
  287.         return x

3、修改 tasks.py

ultralytics/nn/tasks.py

添加

from ultralytics.nn.SwinTransformer import SwinTransformer

def parse_model函数

  1. if m in (Classify, Conv, ConvTranspose, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, Focus,
  2.                  BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x, RepC3):
  3.             c1, c2 = ch[f], args[0]

改为：

  1. if m in (Classify, Conv, ConvTranspose, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, Focus,
  2.                  BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x, RepC3, SwinTransformer):
  3.             c1, c2 = ch[f], args[0]

即结尾增加 SwinTransformer

4、根目录增加文件

data.yaml

  1. # Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
  2.  
  3. path: ../dataset_yolo # dataset root dir
  4. train: train # train images (relative to 'path') 128 images
  5. val: val # val images (relative to 'path') 128 images
  6. test: # test images (optional)
  7.  
  8. # Classes
  9. names:
  10.   0: yb_text
  11.   1: kk_text
  12.   2: zsd_text
  13.   3: xn_text
  14.   4: controls_text
  15.   5: water_mark
  16.   6: yb
  17.   7: kk
  18.   8: zsd
  19.   9: xn
  20.  
  21. # Download script/URL (optional)
  22. download: https://ultralytics.com/assets/coco128-seg.zip

train.py

不再载入与训练模型

  1. from ultralytics import YOLO 
  2.  
  3. # load a model
  4. # model = YOLO('yolov8m.pt')
  5. model = YOLO('yolov8m.yaml')
  6.  
  7. # Train the model
  8. model.train(data='./data.yaml',epochs=300,batch=64,optimizer='SGD',close_mosaic=10,imgsz=640,device=[4],cache=True)
  9.  
  10.  
  11. # https://blog.csdn.net/apple_59275002/article/details/132181112
  12.  
  13. # from ultralytics import YOLO
  14. # import os
  15. # model = YOLO('yolov8n.yaml')
  16. # model = YOLO('yolov8n.pt')
  17.  
  18. # results = model.train(data='custom.yaml', epochs=80, batch=8, patience=0, augment=True, val=False, degrees=15, translate=0.05, scale=0.05, shear=0.05, perspective=0.0, mosaic=0.0, hsv_h=0.010, hsv_s=0.5, hsv_v=0.2)
  19.  
  20. # results = model.val()
  21.  
  22. # 更多参数见网址

训练文件  my_train.sh

nohup python train.py >>train.log 2>&1 &

训练即可

训练时提示如下表示模块加入成功

refer:

https://blog.csdn.net/weixin_51692073/article/details/132724315

https://www.bilibili.com/video/BV1T8411B7iP/

https://space.bilibili.com/432570190

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/Monodyee/article/detail/371449
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号