• YOLOv8+swin_transfomerv2

    测试环境：cuda11.3  pytorch1.11 rtx3090  wsl2 ubuntu20.04

    踩了很多坑，网上很多博主的代码根本跑不通，自己去github仓库复现修改的

    网上博主的代码日常出现cpu,gpu混合，或许是人家分布式训练了，哈哈哈

    下面上干货吧，宝子们点个关注，点个赞，没有废话

    yolov8_yaml文件修改(目标检测就修改

    1. # Ultralytics YOLO 🚀, AGPL-3.0 license
    2. # YOLOv8-seg instance segmentation model. For Usage examples see https://docs.ultralytics.com/tasks/segment
    3.  
    4. # Parameters
    5. nc: 1  # number of classes
    6. scales: # model compound scaling constants, i.e. 'model=yolov8n-seg.yaml' will call yolov8-seg.yaml with scale 'n'
    7.   # [depth, width, max_channels]
    8.   n: [0.33, 0.25, 1024]
    9.   s: [0.33, 0.50, 1024]
    10.   m: [0.67, 0.75, 768]
    11.   l: [1.00, 1.00, 512]
    12.   x: [1.00, 1.25, 512]
    13.  
    14. # YOLOv8.0n backbone
    15. backbone:
    16.   # [from, repeats, module, args]
    17.   - [-1, 1, Conv, [64, 3, 2]]  # 0-P1/2
    18.   - [-1, 1, Conv, [128, 3, 2]]  # 1-P2/4
    19.   - [-1, 3, C2f, [128, True]]
    20.   - [-1, 1, Conv, [256, 3, 2]]  # 3-P3/8
    21.   - [-1, 6, C2f, [256, True]]
    22.   - [-1, 1, Conv, [512, 3, 2]]  # 5-P4/16
    23.   - [-1, 9, SwinV2_CSPB, [512, 512]]
    24.   - [-1, 1, Conv, [1024, 3, 2]]  # 7-P5/32
    25.   - [-1, 3, SwinV2_CSPB, [1024, 1024]]
    26.   - [-1, 1, SPPF, [1024, 5]]  # 9
    27.  
    28. # YOLOv8.0n head
    29. head:
    30.   - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
    31.   - [[-1, 6], 1, Concat, [1]]  # cat backbone P4
    32.   - [-1, 3, C2f, [512]]  # 12
    33.  
    34.   - [-1, 1, nn.Upsample, [None, 2, 'nearest']]
    35.   - [[-1, 4], 1, Concat, [1]]  # cat backbone P3
    36.   - [-1, 3, C2f, [256]]  # 15 (P3/8-small)
    37.  
    38.   - [-1, 1, Conv, [256, 3, 2]]
    39.   - [[-1, 12], 1, Concat, [1]]  # cat head P4
    40.   - [-1, 3, C2f, [512]]  # 18 (P4/16-medium)
    41.  
    42.   - [-1, 1, Conv, [512, 3, 2]]
    43.   - [[-1, 9], 1, Concat, [1]]  # cat head P5
    44.   - [-1, 3, C2f, [1024]]  # 21 (P5/32-large)
    45.  
    46.   - [[15, 18, 21], 1, Segment, [nc, 32, 256]]  # Segment(P3, P4, P5)

    在nn/modules/block.py最下面加入

    1. import torch
    2. import torch.nn as nn
    3. import torch.nn.functional as F
    4. from timm.models.layers import DropPath, to_2tuple, trunc_normal_
    5. from .conv import Conv, DWConv, GhostConv, LightConv, RepConv
    6. from .transformer import TransformerBlock
    7. import numpy as np
    8. class WindowAttention(nn.Module):
    9.  
    10.     def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
    11.  
    12.         super().__init__()
    13.         self.dim = dim
    14.         self.window_size = window_size  # Wh, Ww
    15.         self.num_heads = num_heads
    16.         head_dim = dim // num_heads
    17.         self.scale = qk_scale or head_dim ** -0.5
    18.  
    19.         # define a parameter table of relative position bias
    20.         self.relative_position_bias_table = nn.Parameter(
    21.             torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
    22.  
    23.         # get pair-wise relative position index for each token inside the window
    24.         coords_h = torch.arange(self.window_size[0])
    25.         coords_w = torch.arange(self.window_size[1])
    26.         coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
    27.         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
    28.         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
    29.         relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
    30.         relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
    31.         relative_coords[:, :, 1] += self.window_size[1] - 1
    32.         relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
    33.         relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
    34.         self.register_buffer("relative_position_index", relative_position_index)
    35.  
    36.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
    37.         self.attn_drop = nn.Dropout(attn_drop)
    38.         self.proj = nn.Linear(dim, dim)
    39.         self.proj_drop = nn.Dropout(proj_drop)
    40.  
    41.         nn.init.normal_(self.relative_position_bias_table, std=.02)
    42.         self.softmax = nn.Softmax(dim=-1)
    43.  
    44.     def forward(self, x, mask=None):
    45.  
    46.         B_, N, C = x.shape
    47.         qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
    48.         q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
    49.  
    50.         q = q * self.scale
    51.         attn = (q @ k.transpose(-2, -1))
    52.  
    53.         relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
    54.             self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
    55.         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
    56.         attn = attn + relative_position_bias.unsqueeze(0)
    57.  
    58.         if mask is not None:
    59.             nW = mask.shape[0]
    60.             attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
    61.             attn = attn.view(-1, self.num_heads, N, N)
    62.             attn = self.softmax(attn)
    63.         else:
    64.             attn = self.softmax(attn)
    65.  
    66.         attn = self.attn_drop(attn)
    67.  
    68.         # print(attn.dtype, v.dtype)
    69.         try:
    70.             x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
    71.         except:
    72.             # print(attn.dtype, v.dtype)
    73.             x = (attn.half() @ v).transpose(1, 2).reshape(B_, N, C)
    74.         x = self.proj(x)
    75.         x = self.proj_drop(x)
    76.         return x
    77.  
    78.  
    79. class WindowAttention(nn.Module):
    80.  
    81.     def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
    82.  
    83.         super().__init__()
    84.         self.dim = dim
    85.         self.window_size = window_size  # Wh, Ww
    86.         self.num_heads = num_heads
    87.         head_dim = dim // num_heads
    88.         self.scale = qk_scale or head_dim ** -0.5
    89.  
    90.         # define a parameter table of relative position bias
    91.         self.relative_position_bias_table = nn.Parameter(
    92.             torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
    93.  
    94.         # get pair-wise relative position index for each token inside the window
    95.         coords_h = torch.arange(self.window_size[0])
    96.         coords_w = torch.arange(self.window_size[1])
    97.         coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
    98.         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
    99.         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
    100.         relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
    101.         relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
    102.         relative_coords[:, :, 1] += self.window_size[1] - 1
    103.         relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
    104.         relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
    105.         self.register_buffer("relative_position_index", relative_position_index)
    106.  
    107.         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
    108.         self.attn_drop = nn.Dropout(attn_drop)
    109.         self.proj = nn.Linear(dim, dim)
    110.         self.proj_drop = nn.Dropout(proj_drop)
    111.  
    112.         nn.init.normal_(self.relative_position_bias_table, std=.02)
    113.         self.softmax = nn.Softmax(dim=-1)
    114.  
    115.     def forward(self, x, mask=None):
    116.  
    117.         B_, N, C = x.shape
    118.         qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
    119.         q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
    120.  
    121.         q = q * self.scale
    122.         attn = (q @ k.transpose(-2, -1))
    123.  
    124.         relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
    125.             self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
    126.         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
    127.         attn = attn + relative_position_bias.unsqueeze(0)
    128.  
    129.         if mask is not None:
    130.             nW = mask.shape[0]
    131.             attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
    132.             attn = attn.view(-1, self.num_heads, N, N)
    133.             attn = self.softmax(attn)
    134.         else:
    135.             attn = self.softmax(attn)
    136.  
    137.         attn = self.attn_drop(attn)
    138.  
    139.         # print(attn.dtype, v.dtype)
    140.         try:
    141.             x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
    142.         except:
    143.             # print(attn.dtype, v.dtype)
    144.             x = (attn.half() @ v).transpose(1, 2).reshape(B_, N, C)
    145.         x = self.proj(x)
    146.         x = self.proj_drop(x)
    147.         return x
    148.  
    149.  
    150. class WindowAttention_v2(nn.Module):
    151.  
    152.     def __init__(self, dim, window_size, num_heads, qkv_bias=True, attn_drop=0., proj_drop=0.,
    153.                  pretrained_window_size=[0, 0]):
    154.  
    155.         super().__init__()
    156.         self.dim = dim
    157.         self.window_size = window_size  # Wh, Ww
    158.         self.pretrained_window_size = pretrained_window_size
    159.         self.num_heads = num_heads
    160.  
    161.         self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)
    162.  
    163.         # mlp to generate continuous relative position bias
    164.         self.cpb_mlp = nn.Sequential(nn.Linear(2, 512, bias=True),
    165.                                      nn.ReLU(inplace=True),
    166.                                      nn.Linear(512, num_heads, bias=False))
    167.  
    168.         # get relative_coords_table
    169.         relative_coords_h = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)
    170.         relative_coords_w = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)
    171.         relative_coords_table = torch.stack(
    172.             torch.meshgrid([relative_coords_h,
    173.                             relative_coords_w])).permute(1, 2, 0).contiguous().unsqueeze(0)  # 1, 2*Wh-1, 2*Ww-1, 2
    174.         if pretrained_window_size[0] > 0:
    175.             relative_coords_table[:, :, :, 0] /= (pretrained_window_size[0] - 1)
    176.             relative_coords_table[:, :, :, 1] /= (pretrained_window_size[1] - 1)
    177.         else:
    178.             relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)
    179.             relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)
    180.         relative_coords_table *= 8  # normalize to -8, 8
    181.         relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
    182.             torch.abs(relative_coords_table) + 1.0) / np.log2(8)
    183.  
    184.         self.register_buffer("relative_coords_table", relative_coords_table)
    185.  
    186.         # get pair-wise relative position index for each token inside the window
    187.         coords_h = torch.arange(self.window_size[0])
    188.         coords_w = torch.arange(self.window_size[1])
    189.         coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
    190.         coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
    191.         relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
    192.         relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
    193.         relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
    194.         relative_coords[:, :, 1] += self.window_size[1] - 1
    195.         relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
    196.         relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
    197.         self.register_buffer("relative_position_index", relative_position_index)
    198.  
    199.         self.qkv = nn.Linear(dim, dim * 3, bias=False)
    200.         if qkv_bias:
    201.             self.q_bias = nn.Parameter(torch.zeros(dim))
    202.             self.v_bias = nn.Parameter(torch.zeros(dim))
    203.         else:
    204.             self.q_bias = None
    205.             self.v_bias = None
    206.         self.attn_drop = nn.Dropout(attn_drop)
    207.         self.proj = nn.Linear(dim, dim)
    208.         self.proj_drop = nn.Dropout(proj_drop)
    209.         self.softmax = nn.Softmax(dim=-1)
    210.  
    211.     def forward(self, x, mask=None):
    212.  
    213.         B_, N, C = x.shape
    214.         qkv_bias = None
    215.         if self.q_bias is not None:
    216.             qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
    217.         qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
    218.         qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
    219.         q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
    220.  
    221.         # cosine attention
    222.         attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
    223.         max_tensor = torch.log(torch.tensor(1. / 0.01)).to(self.logit_scale.device)
    224.         logit_scale = torch.clamp(self.logit_scale, max=max_tensor).exp()
    225.         attn = attn * logit_scale
    226.  
    227.         relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)
    228.         relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
    229.             self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
    230.         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
    231.         relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
    232.         attn = attn + relative_position_bias.unsqueeze(0)
    233.  
    234.         if mask is not None:
    235.             nW = mask.shape[0]
    236.             attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
    237.             attn = attn.view(-1, self.num_heads, N, N)
    238.             attn = self.softmax(attn)
    239.         else:
    240.             attn = self.softmax(attn)
    241.  
    242.         attn = self.attn_drop(attn)
    243.  
    244.         try:
    245.             x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
    246.         except:
    247.             x = (attn.half() @ v).transpose(1, 2).reshape(B_, N, C)
    248.  
    249.         x = self.proj(x)
    250.         x = self.proj_drop(x)
    251.         return x
    252.  
    253.     def extra_repr(self) -> str:
    254.         return f'dim={self.dim}, window_size={self.window_size}, ' \
    255.                f'pretrained_window_size={self.pretrained_window_size}, num_heads={self.num_heads}'
    256.  
    257.     def flops(self, N):
    258.         # calculate flops for 1 window with token length of N
    259.         flops = 0
    260.         # qkv = self.qkv(x)
    261.         flops += N * self.dim * 3 * self.dim
    262.         # attn = (q @ k.transpose(-2, -1))
    263.         flops += self.num_heads * N * (self.dim // self.num_heads) * N
    264.         #  x = (attn @ v)
    265.         flops += self.num_heads * N * N * (self.dim // self.num_heads)
    266.         # x = self.proj(x)
    267.         flops += N * self.dim * self.dim
    268.         return flops
    269.  
    270.  
    271. class Mlp_v2(nn.Module):
    272.     def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0.):
    273.         super().__init__()
    274.         out_features = out_features or in_features
    275.         hidden_features = hidden_features or in_features
    276.         self.fc1 = nn.Linear(in_features, hidden_features)
    277.         self.act = act_layer()
    278.         self.fc2 = nn.Linear(hidden_features, out_features)
    279.         self.drop = nn.Dropout(drop)
    280.  
    281.     def forward(self, x):
    282.         x = self.fc1(x)
    283.         x = self.act(x)
    284.         x = self.drop(x)
    285.         x = self.fc2(x)
    286.         x = self.drop(x)
    287.         return x
    288.  
    289.  
    290. # add 2 functions
    291. class SwinTransformerLayer_v2(nn.Module):
    292.  
    293.     def __init__(self, dim, num_heads, window_size=7, shift_size=0,
    294.                  mlp_ratio=4., qkv_bias=True, drop=0., attn_drop=0., drop_path=0.,
    295.                  act_layer=nn.SiLU, norm_layer=nn.LayerNorm, pretrained_window_size=0):
    296.         super().__init__()
    297.         self.dim = dim
    298.         # self.input_resolution = input_resolution
    299.         self.num_heads = num_heads
    300.         self.window_size = window_size
    301.         self.shift_size = shift_size
    302.         self.mlp_ratio = mlp_ratio
    303.         # if min(self.input_resolution) <= self.window_size:
    304.         #    # if window size is larger than input resolution, we don't partition windows
    305.         #    self.shift_size = 0
    306.         #    self.window_size = min(self.input_resolution)
    307.         assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
    308.         self.norm1 = norm_layer(dim)
    309.         self.attn = WindowAttention_v2(
    310.             dim, window_size=(self.window_size, self.window_size), num_heads=num_heads,
    311.             qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop,
    312.             pretrained_window_size=(pretrained_window_size, pretrained_window_size))
    313.         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
    314.         self.norm2 = norm_layer(dim)
    315.         mlp_hidden_dim = int(dim * mlp_ratio)
    316.         self.mlp = Mlp_v2(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
    317.  
    318.  
    319.     def create_mask(self, H, W):
    320.         # calculate attention mask for SW-MSA
    321.         img_mask = torch.zeros((1, H, W, 1))  # 1 H W 1
    322.         h_slices = (slice(0, -self.window_size),
    323.                     slice(-self.window_size, -self.shift_size),
    324.                     slice(-self.shift_size, None))
    325.         w_slices = (slice(0, -self.window_size),
    326.                     slice(-self.window_size, -self.shift_size),
    327.                     slice(-self.shift_size, None))
    328.         cnt = 0
    329.         for h in h_slices:
    330.             for w in w_slices:
    331.                 img_mask[:, h, w, :] = cnt
    332.                 cnt += 1
    333.  
    334.         def window_partition(x, window_size):
    335.             """
    336.             Args:
    337.                 x: (B, H, W, C)
    338.                 window_size (int): window size
    340.             Returns:
    341.                 windows: (num_windows*B, window_size, window_size, C)
    342.             """
    343.             B, H, W, C = x.shape
    344.             x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    345.             windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    346.             return windows
    347.  
    348.         mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
    349.         mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
    350.         attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
    351.         attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
    352.         return attn_mask
    353.  
    354.     def forward(self, x):
    355.         # reshape x[b c h w] to x[b l c]
    356.         _, _, H_, W_ = x.shape
    357.         Padding = False
    358.         if min(H_, W_) < self.window_size or H_ % self.window_size != 0 or W_ % self.window_size != 0:
    359.             Padding = True
    360.             # print(f'img_size {min(H_, W_)} is less than (or not divided by) window_size {self.window_size}, Padding.')
    361.             pad_r = (self.window_size - W_ % self.window_size) % self.window_size
    362.             pad_b = (self.window_size - H_ % self.window_size) % self.window_size
    363.             x = F.pad(x, (0, pad_r, 0, pad_b))
    364.         # print('2', x.shape)
    365.         B, C, H, W = x.shape
    366.         L = H * W
    367.         x = x.permute(0, 2, 3, 1).contiguous().view(B, L, C)  # b, L, c
    368.         # create mask from init to forward
    369.         if self.shift_size > 0:
    370.             attn_mask = self.create_mask(H, W).to(x.device)
    371.         else:
    372.             attn_mask = None
    373.         shortcut = x
    374.         x = x.view(B, H, W, C)
    375.         # cyclic shift
    376.         if self.shift_size > 0:
    377.             shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
    378.         else:
    379.             shifted_x = x
    380.         # partition windows
    381.  
    382.         def window_partition(x, window_size):
    383.             """
    384.             Args:
    385.                 x: (B, H, W, C)
    386.                 window_size (int): window size
    388.             Returns:
    389.                 windows: (num_windows*B, window_size, window_size, C)
    390.             """
    391.             B, H, W, C = x.shape
    392.             x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
    393.             windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
    394.             return windows
    395.  
    396.         def window_reverse(windows, window_size, H, W):
    397.             """
    398.             Args:
    399.                 windows: (num_windows*B, window_size, window_size, C)
    400.                 window_size (int): Window size
    401.                 H (int): Height of image
    402.                 W (int): Width of image
    404.             Returns:
    405.                 x: (B, H, W, C)
    406.             """
    407.             B = int(windows.shape[0] / (H * W / window_size / window_size))
    408.             x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
    409.             x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
    410.             return x
    411.  
    412.  
    413.         x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
    414.         x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
    415.         # W-MSA/SW-MSA
    416.         attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
    417.         # merge windows
    418.         attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
    419.         shifted_x = window_reverse(attn_windows, self.window_size, H, W)  # B H' W' C
    420.         # reverse cyclic shift
    421.         if self.shift_size > 0:
    422.             x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
    423.         else:
    424.             x = shifted_x
    425.         x = x.view(B, H * W, C)
    426.         x = shortcut + self.drop_path(self.norm1(x))
    427.         # FFN
    428.         x = x + self.drop_path(self.norm2(self.mlp(x)))
    429.         x = x.permute(0, 2, 1).contiguous().view(-1, C, H, W)  # b c h w
    430.  
    431.         if Padding:
    432.             x = x[:, :, :H_, :W_]  # reverse padding
    433.         return x
    434.  
    435.     def extra_repr(self) -> str:
    436.         return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
    437.                f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
    438.  
    439.     def flops(self):
    440.         flops = 0
    441.         H, W = self.input_resolution
    442.         # norm1
    443.         flops += self.dim * H * W
    444.         # W-MSA/SW-MSA
    445.         nW = H * W / self.window_size / self.window_size
    446.         flops += nW * self.attn.flops(self.window_size * self.window_size)
    447.         # mlp
    448.         flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
    449.         # norm2
    450.         flops += self.dim * H * W
    451.         return flops
    452.  
    453.  
    454. class SwinTransformer2Block(nn.Module):
    455.     def __init__(self, c1, c2, num_heads, num_layers, window_size=7):
    456.         super().__init__()
    457.         self.conv = None
    458.         if c1 != c2:
    459.             self.conv = Conv(c1, c2)
    460.         # remove input_resolution
    461.         self.blocks = nn.Sequential(*[SwinTransformerLayer_v2(dim=c2, num_heads=num_heads, window_size=window_size,
    462.                                                               shift_size=0 if (i % 2 == 0) else window_size // 2) for i
    463.                                       in range(num_layers)])
    464.  
    465.     def forward(self, x):
    466.         if self.conv is not None:
    467.             x = self.conv(x)
    468.         x = self.blocks(x)
    469.         return x
    470.  
    471.  
    472. class SwinV2_CSPB(nn.Module):
    473.     # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
    474.     def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
    475.         super(SwinV2_CSPB, self).__init__()
    476.         c_ = int(c2)  # hidden channels
    477.         self.cv1 = Conv(c1, c_, 1, 1)
    478.         self.cv2 = Conv(c_, c_, 1, 1)
    479.         self.cv3 = Conv(2 * c_, c2, 1, 1)
    480.         num_heads = c_ // 32
    481.         self.m = SwinTransformer2Block(c_, c_, num_heads, n)
    482.         # self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])
    483.  
    484.     def forward(self, x):
    485.         x1 = self.cv1(x)
    486.         y1 = self.m(x1)
    487.         y2 = self.cv2(x1)
    488.         return self.cv3(torch.cat((y1, y2), dim=1))

    随后在这个目录的init.py里面把你的新的东西放进去

    1. from .block import (C1, C2, C3, C3TR, DFL, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x, GhostBottleneck,
    2.                     HGBlock, HGStem, Proto, RepC3,GAM_Attention,ResBlock_CBAM,GCT,C3STR,SwinV2_CSPB)
    3. from .conv import (CBAM, ChannelAttention, Concat, Conv, Conv2, ConvTranspose, DWConv, DWConvTranspose2d, Focus,
    4.                    GhostConv, LightConv, RepConv, SpatialAttention)
    5. from .head import Classify, Detect, Pose, RTDETRDecoder, Segment
    6. from .transformer import (AIFI, MLP, DeformableTransformerDecoder, DeformableTransformerDecoderLayer, LayerNorm2d,
    7.                           MLPBlock, MSDeformAttn, TransformerBlock, TransformerEncoderLayer, TransformerLayer)
    8.  
    9. __all__ = ('Conv', 'Conv2', 'LightConv', 'RepConv', 'DWConv', 'DWConvTranspose2d', 'ConvTranspose', 'Focus',
    10.            'GhostConv', 'ChannelAttention', 'SpatialAttention', 'CBAM', 'Concat', 'TransformerLayer',
    11.            'TransformerBlock', 'MLPBlock', 'LayerNorm2d', 'DFL', 'HGBlock', 'HGStem', 'SPP', 'SPPF', 'C1', 'C2', 'C3',
    12.            'C2f', 'C3x', 'C3TR', 'C3Ghost', 'GhostBottleneck', 'Bottleneck', 'BottleneckCSP', 'Proto', 'Detect',
    13.            'Segment', 'Pose', 'Classify', 'TransformerEncoderLayer', 'RepC3', 'RTDETRDecoder', 'AIFI',
    14.            'DeformableTransformerDecoder', 'DeformableTransformerDecoderLayer', 'MSDeformAttn', 'MLP','ResBlock_CBAM','CBAM','GAM_Attention',
    15.            'GCT','C3STR','SwinV2_CSPB')

    然后去task.py

    在导包那里，把这个导入进去

    1. rom ultralytics.nn.modules import (AIFI, C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x,
    2.                                     Classify, Concat, Conv, Conv2, ConvTranspose, Detect, DWConv, DWConvTranspose2d,
    3.                                     Focus, GhostBottleneck, GhostConv, HGBlock, HGStem, Pose, RepC3, RepConv,
    4.                                     RTDETRDecoder ,Segment,CBAM,GAM_Attention,ResBlock_CBAM,GCT,C3STR,SwinV2_CSPB)

    然后同样在task.py里面搜索if m in

    把你的新的swin模块放进去

    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.                  ,CBAM , GAM_Attention ,ResBlock_CBAM,GCT,C3STR,SwinV2_CSPB):

    最后

    命令行：python setup.py install 注册

    命令行：你的训练命令

  • 相关阅读:
    如何还原回收站已经清空的文件
    想学设计模式、想搞架构设计，先学学UML系统建模吧您
    MySQL高级10-InnoDB引擎存储架构
    解决mybatis case when 报错的问题
    你们关心的问题：产品经理面试中的职业规划及项目经历要怎么说？
    【Java Web】论坛——我收到的赞
    7 Spring Boot 整合 Spring Data JPA
    Qt编写物联网管理平台47-通用数据库设置
    【Redis】数据结构---String
    常用工具类commons-io的学习使用
  • 原文地址：https://blog.csdn.net/xty123abc/article/details/133428517