• YOLOv5、YOLOv8改进：Swin Transformer-V2

    1.介绍

    论文地址：https://arxiv.org/abs/2111.09883

    综述

    该论文作者提出了缩放 Swin Transformer 的技术 多达 30 亿个参数，使其能够使用多达 1,536 个图像进行训练1,536 分辨率。通过扩大容量和分辨率，Swin Transformer 在四个具有代表性的视觉基准上创造了新记录：ImageNet-V2 图像分类的84.0% top-1 准确率，COCO 对象检测的63.1 / 54.4 box / mask mAP，ADE20K 语义分割的59.9 mIoU，和86.8%Kinetics-400 视频动作分类的前 1 准确率。我们的技术通常适用于扩大视觉模型，但尚未像 NLP 语言模型那样被广泛探索，部分原因是在训练和应用方面存在以下困难：1）视觉模型经常面临大规模的不稳定性问题和 2）许多下游视觉任务需要高分辨率图像或窗口，目前尚不清楚如何有效地将低分辨率预训练的模型转移到更高分辨率的模型。当图像分辨率很高时，GPU 内存消耗也是一个问题。为了解决这些问题，我们提出了几种技术，并通过使用 Swin Transformer 作为案例研究来说明：1）后归一化技术和缩放余弦注意方法，以提高大型视觉模型的稳定性；2) 一种对数间隔的连续位置偏差技术，可有效地将在低分辨率图像和窗口上预训练的模型转移到其更高分辨率的对应物上。此外，我们分享了我们的关键实现细节，这些细节可以显着节省 GPU 内存消耗，从而使使用常规 GPU 训练大型视觉模型变得可行。使用这些技术和自我监督的预训练，我们成功训练了一个强大的 30 亿个 Swin Transformer 模型，并有效地将其转移到涉及高分辨率图像或窗口的各种视觉任务中，在各种的基准。代码将在 我们分享了我们的关键实现细节，这些细节可以显着节省 GPU 内存消耗，从而使使用常规 GPU 训练大型视觉模型变得可行。使用这些技术和自我监督的预训练，我们成功训练了一个强大的 30 亿个 Swin Transformer 模型，并有效地将其转移到涉及高分辨率图像或窗口的各种视觉任务中，在各种的基准。代码将在 我们分享了我们的关键实现细节，这些细节可以显着节省 GPU 内存消耗，从而使使用常规 GPU 训练大型视觉模型变得可行。使用这些技术和自我监督的预训练，我们成功训练了一个强大的 30 亿个 Swin Transformer 模型，并有效地将其转移到涉及高分辨率图像或窗口的各种视觉任务中，在各种的基准。代码将在 我们成功训练了一个强大的 30 亿个 Swin Transformer 模型，并将其有效地转移到涉及高分辨率图像或窗口的各种视觉任务中，在各种基准测试中达到了最先进的精度。代码将在 我们成功训练了一个强大的 30 亿个 Swin Transformer 模型，并将其有效地转移到涉及高分辨率图像或窗口的各种视觉任务中，在各种基准测试中达到了最先进的精度。

    要解决的问题
    视觉模型通常面临尺度不稳定问题；

    下游任务需要高分辨率图像，尚不明确如何将低分辨率预训练模型迁移为高分辨率版本 ；当图像分辨率非常大时，GPU显存占用也是个问题。

    改进方案
    提出后规范化(Post Normalization)技术与可缩放(Scaled)cosine注意力提升大视觉模型的稳定性;
    提出log空间连续位置偏置技术进行低分辨率预训练模型向高分辨率模型迁移；
    我们还共享了至关重要的实现细节 ，它可以大幅节省GPU显存占用以使得大视觉模型训练变得可行。

    2.  YOLOv5改进方法

    2.1 YOLOv5的yaml配置文件

    首先增加以下yolov5_swin_transfomrer.yaml文件

    1. # YOLOv5 🚀 by Ultralytics, GPL-3.0 license
    2.  
    3. # Parameters
    4. nc: 80  # number of classes
    5. depth_multiple: 0.33  # model depth multiple
    6. width_multiple: 0.50  # layer channel 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
    11.  
    12. # YOLOv5 v6.0 backbone by yoloair
    13. backbone:
    14.   # [from, number, module, args]
    15.   [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
    16.    [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
    17.    [-1, 3, C3, [128]],
    18.    [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
    19.    [-1, 6, C3, [256]],
    20.    [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
    21.    [-1, 9, SwinV2_CSPB, [256, 256]],
    22.    [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
    23.    [-1, 3, SwinV2_CSPB, [512, 512]],  # 9 <--- ST2CSPB() Transformer module
    24.    [-1, 1, SPPF, [1024, 5]],  # 9
    25.   ]
    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
    33.  
    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)
    38.  
    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)
    42.  
    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)
    46.  
    47.    [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
    48.   ]
    49.

    2.2 common.py配置

    在./models/common.py文件中增加以下模块，直接复制即可

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

    2.3 yolo.py配置

    不需要

    修改完成

  • 相关阅读:
    将自定义 GitHub 徽章添加到您的代码库
    数据化运营15 活跃（上）：如何通过运营手法提升⽤户活跃度？
    HK32F030MF4P6 红外遥控接收例程
    Git 的暂存区(staging area)理解
    linux之cpu模拟负载程序
    同步辐射散射测试中影响效果的原因有哪些？
    互斥量互斥锁
    Origin中如何上标？R\+(2)即 — R的2上标
    element-plus打开Dialog、图片预览导致页面抖动
    伺服丝杠系统常用运算功能块
  • 原文地址：https://blog.csdn.net/weixin_45303602/article/details/133216041