YOLOv5改进——添加SimAM注意力机制

目录

一、SimAM注意力机制核心代码

二、修改common.py

三、修改yolo.py

​三、建立yaml文件

四、验证


一、SimAM注意力机制核心代码

在models文件夹下新建modules文件夹,在modules文件夹下新建一个py文件。这里为simam.py。复制以下代码到文件里面。

python 复制代码
import torch
import torch.nn as nn


class SimAM(torch.nn.Module):
    def __init__(self, channels = None, e_lambda = 1e-4):
        super(SimAM, self).__init__()

        self.activaton = nn.Sigmoid()
        self.e_lambda = e_lambda

    def __repr__(self):
        s = self.__class__.__name__ + '('
        s += ('lambda=%f)' % self.e_lambda)
        return s

    @staticmethod
    def get_module_name():
        return "simam"

    def forward(self, x):

        b, c, h, w = x.size()
        
        n = w * h - 1

        x_minus_mu_square = (x - x.mean(dim=[2,3], keepdim=True)).pow(2)
        y = x_minus_mu_square / (4 * (x_minus_mu_square.sum(dim=[2,3], keepdim=True) / n + self.e_lambda)) + 0.5

        return x * self.activaton(y)

注意:很多改进教程都是将代码直接复制到common.py文件,如果改进机制多了容易造成混乱。建议创建一个modules文件夹,将改进机制放里面方便管理。

二、修改common.py

在common.py文件中,在前面的部分添加以下代码,导入GhostV2.py的内容:

python 复制代码
from models.modules.simam import *

三、修改yolo.py

在yolo.py文件中,在导入common模块的上面一行添加以下代码,导入GhostV2.py的内容:

python 复制代码
from models.modules.simam import *

注意:这里位置不要搞错,不然可能会找不到导入的模块。

如下图所示:

找到parse_model函数,将SimAM模块加入,如下图所示:

​三、建立yaml文件

在models文件夹下,复制yolov5s.yaml文件,粘贴并重命名为yolov5s-simam.yaml。

这里将SimAM注意力机制加在backbone最末端。这样可以使注意力机制看到整个backbone部分的特征图,将具有全局视野,类似于一个小的transformer结构。

如图所示,将SimAM注意力机制加在SPPF的下一层:

加了层数后,后面的head部分也得修改,如下图所示:

修改前:

修改后:

加一层就在参数上加一,若加多个机制,依此类推。

yolov5s-simam.yaml完整代码如下:

python 复制代码
# Ultralytics YOLOv5 , AGPL-3.0 license

# Parameters
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
anchors:
  - [10, 13, 16, 30, 33, 23] # P3/8
  - [30, 61, 62, 45, 59, 119] # P4/16
  - [116, 90, 156, 198, 373, 326] # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [
    [-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
    [-1, 1, Conv, [128, 3, 2]], # 1-P2/4
    [-1, 3, C3, [128]],
    [-1, 1, Conv, [256, 3, 2]], # 3-P3/8
    [-1, 6, C3, [256]],
    [-1, 1, Conv, [512, 3, 2]], # 5-P4/16
    [-1, 9, C3, [512]],
    [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
    [-1, 3, C3, [1024]],
    [-1, 1, SPPF, [1024, 5]], # 9
    [-1, 3, SimAM, [1024]], # 10
  ]

# YOLOv5 v6.0 head
head: [
    [-1, 1, Conv, [512, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 6], 1, Concat, [1]], # cat backbone P4
    [-1, 3, C3, [512, False]], # 14

    [-1, 1, Conv, [256, 1, 1]],
    [-1, 1, nn.Upsample, [None, 2, "nearest"]],
    [[-1, 4], 1, Concat, [1]], # cat backbone P3
    [-1, 3, C3, [256, False]], # 18 (P3/8-small)

    [-1, 1, Conv, [256, 3, 2]],
    [[-1, 15], 1, Concat, [1]], # cat head P4
    [-1, 3, C3, [512, False]], # 21 (P4/16-medium)

    [-1, 1, Conv, [512, 3, 2]],
    [[-1, 11], 1, Concat, [1]], # cat head P5
    [-1, 3, C3, [1024, False]], # 24 (P5/32-large)

    [[18, 21, 24], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
  ]

至此更改完成。

四、验证

在yolo.py 文件里面的配置文件改为刚才自定义的yolov5s-simam.yaml

ctrl+shift+p 在弹出框窗口搜索Python:选择解释器,选择自己创建的Python虚拟环境,这里是yolo。

运行yolo.py,出现刚刚加入的SimAM注意力机制,说明加入成功。

相关推荐
只是甲2 小时前
Text2SQL 系列博客 02:技术原理深度剖析 - 从自然语言到 SQL 的完整链路
人工智能·text2sql·nl2sql·ai agent·自助数据分析·data agent·agentic 数据洞察
懷淰メ2 小时前
【AI赋能】基于PyQt+YOLO+DeepSeek水上漂浮物检测系统(详细介绍)
人工智能·yolo·目标检测·计算机视觉·pyqt·漂浮物·水上漂浮物
EQUINOX12 小时前
【论文精读】| CLIP精读
大数据·人工智能
PC2005-cloud2 小时前
博文已索引但不展示:一次 Bing SEO 问题排查记录
服务器·人工智能
weixin_408099672 小时前
2026 图片去文字 API 实战:AI一键去除图片文字(附 Python / Java / PHP / JS 完整代码)
人工智能·api接口·图片去水印·石榴智能·ai图像修复·图片去文字
幻影123!3 小时前
从零训练一个会下五子棋的AI
python·深度学习·神经网络·强化学习·五子棋·alpha zero·mokugo
greenbbLV3 小时前
2026节日慰问品政策解读:合规标准升级与福利行业新趋势
人工智能·节日
奈斯先生Vector3 小时前
告别工具碎片化:基于 Nano Banana 全模态 AI 聚合架构搭建“文本-图像-视频”自动化协同生产线
运维·数据库·人工智能·架构·自动化·aigc·音视频
乌萨奇也要立志学C++3 小时前
【Linux】网络基础 从零理解网络通信:局域网、跨网传输与 Socket 预备知识
linux·服务器·网络
领麦微红外3 小时前
国产MEMS红外测温传感器:传感器+算法+产品级出厂标定一站式交付模式深度解读
人工智能·硬件工程·产品经理