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

零基础教程：Yolov5模型改进-添加13种注意力机制_yolo模型添加注意力机制

作者：码创造者 | 2024-07-24 08:54:29

赞

踩

yolo模型添加注意力机制

1.准备工作

先给出13种注意力机制的下载地址：

https://github.com/z1069614715/objectdetection_script

2.加入注意力机制

1.以添加SimAM注意力机制为例（不需要接收通道数的注意力机制）

1.在models文件下新建py文件，取名叫SimAM.py

将以下代码复制到SimAM.py文件种

  1. import torch
  2. import torch.nn as nn
  3.  
  4.  
  5. class SimAM(torch.nn.Module):
  6.  
  7.     # 不需要接收通道数输入
  8.     def __init__(self, e_lambda=1e-4):
  9.         super(SimAM, self).__init__()
  10.  
  11.         self.activaton = nn.Sigmoid()
  12.         self.e_lambda = e_lambda
  13.  
  14.     def __repr__(self):
  15.         s = self.__class__.__name__ + '('
  16.         s += ('lambda=%f)' % self.e_lambda)
  17.         return s
  18.  
  19.     @staticmethod
  20.     def get_module_name():
  21.         return "simam"
  22.  
  23.     def forward(self, x):
  24.         b, c, h, w = x.size()
  25.  
  26.         n = w * h - 1
  27.  
  28.         x_minus_mu_square = (x - x.mean(dim=[2, 3], keepdim=True)).pow(2)
  29.         y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2, 3], keepdim=True) / n + self.e_lambda)) + 0.5
  30.  
  31.         return x * self.activaton(y)

2.在yolo.py头部导入SimAM这个类

3.然后复制yolov5s.yaml到同级目录，取名为yolov5s-SimAM.yaml

在某一层添加注意力机制

[from,number,module,args]

注意：！！！！！！！！！！！！！！！！！！！

添加完一层注意力机制之后，会对后面层数造成影响，记得在检测头那里要改层数

2.添加SE注意力机制（需要接收通道数的注意力机制）

1.新建SE.py

  1. import numpy as np
  2. import torch
  3. from torch import nn
  4. from torch.nn import init
  5.  
  6.  
  7.  
  8. class SEAttention(nn.Module):
  9.  
  10.     def __init__(self, channel=512,reduction=16):
  11.         super().__init__()
  12.         self.avg_pool = nn.AdaptiveAvgPool2d(1)
  13.         self.fc = nn.Sequential(
  14.             nn.Linear(channel, channel // reduction, bias=False),
  15.             nn.ReLU(inplace=True),
  16.             nn.Linear(channel // reduction, channel, bias=False),
  17.             nn.Sigmoid()
  18.         )
  19.  
  20.  
  21.     def init_weights(self):
  22.         for m in self.modules():
  23.             if isinstance(m, nn.Conv2d):
  24.                 init.kaiming_normal_(m.weight, mode='fan_out')
  25.                 if m.bias is not None:
  26.                     init.constant_(m.bias, 0)
  27.             elif isinstance(m, nn.BatchNorm2d):
  28.                 init.constant_(m.weight, 1)
  29.                 init.constant_(m.bias, 0)
  30.             elif isinstance(m, nn.Linear):
  31.                 init.normal_(m.weight, std=0.001)
  32.                 if m.bias is not None:
  33.                     init.constant_(m.bias, 0)
  34.  
  35.     def forward(self, x):
  36.         b, c, _, _ = x.size()
  37.         y = self.avg_pool(x).view(b, c)
  38.         y = self.fc(y).view(b, c, 1, 1)
  39.         return x * y.expand_as(x)

2.修改yolo.py

添加这两行代码

  1.         elif m is SEAttention:
  2.             args = [ch[f]]

3.models下新建yolov5s-SE.yaml

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

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