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

yolov8添加注意力机制-MHSA_在yolov8的骨干网络池化金字塔前加入注意力机制

作者：小蓝xlanll | 2024-04-07 14:26:21

赞

踩

在yolov8的骨干网络池化金字塔前加入注意力机制

一 在ultralytics/nn中新建MHSA.py

添加代码

  1. import torch
  2. import torch.nn as nn
  3.  
  4. class MHSA(nn.Module):
  5.     def __init__(self, n_dims, width=14, height=14, heads=4, pos_emb=False):
  6.         super(MHSA, self).__init__()
  7.  
  8.         self.heads = heads
  9.         self.query = nn.Conv2d(n_dims, n_dims, kernel_size=1)
  10.         self.key = nn.Conv2d(n_dims, n_dims, kernel_size=1)
  11.         self.value = nn.Conv2d(n_dims, n_dims, kernel_size=1)
  12.         self.pos = pos_emb
  13.         if self.pos:
  14.             self.rel_h_weight = nn.Parameter(torch.randn([1, heads, (n_dims) // heads, 1, int(height)]),
  15.                                              requires_grad=True)
  16.             self.rel_w_weight = nn.Parameter(torch.randn([1, heads, (n_dims) // heads, int(width), 1]),
  17.                                              requires_grad=True)
  18.         self.softmax = nn.Softmax(dim=-1)
  19.  
  20.     def forward(self, x):
  21.         n_batch, C, width, height = x.size()
  22.         q = self.query(x).view(n_batch, self.heads, C // self.heads, -1)
  23.         k = self.key(x).view(n_batch, self.heads, C // self.heads, -1)
  24.         v = self.value(x).view(n_batch, self.heads, C // self.heads, -1)
  25.         content_content = torch.matmul(q.permute(0, 1, 3, 2), k)  # 1,C,h*w,h*w
  26.         c1, c2, c3, c4 = content_content.size()
  27.         if self.pos:
  28.             content_position = (self.rel_h_weight + self.rel_w_weight).view(1, self.heads, C // self.heads, -1).permute(
  29.                 0, 1, 3, 2)  # 1,4,1024,64
  30.  
  31.             content_position = torch.matmul(content_position, q)  # ([1, 4, 1024, 256])
  32.             content_position = content_position if (
  33.                     content_content.shape == content_position.shape) else content_position[:, :, :c3, ]
  34.             assert (content_content.shape == content_position.shape)
  35.             energy = content_content + content_position
  36.         else:
  37.             energy = content_content
  38.         attention = self.softmax(energy)
  39.         out = torch.matmul(v, attention.permute(0, 1, 3, 2))  # 1,4,256,64
  40.         out = out.view(n_batch, C, width, height)
  41.         return out
  42.  
  43. if __name__ == '__main__':
  44.     input = torch.randn(50, 512, 7, 7)
  45.     mhsa = MHSA(n_dims=512)
  46.     output = mhsa(input)
  47.     print(output.shape)

二 在 ultralytics/nn/tasks.py中导入

from ultralytics.nn.MHSA import MHSA

找到parse_model函数，加入

  1.         elif m in {MHSA}:
  2.             args=[ch[f],*args]

 

三 在配置文件中修改

在ultralytics/models/v8中添加配置文件yolov8n_MHSA.yaml

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

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