1. 工程简介 🚀

本工程基于 Ultralytics 框架扩展,面向语义分割与 YOLO 系列模型改进实验。核心特点不是只支持单一模型,而是支持通过切换 YAML 配置文件,快速完成不同网络结构的训练、验证与对比实验。

当前已支持的主要模型家族 🧩

  • 语义分割模型:UNetUNet++DeepLabV3+DPTFPNPSPNetMAnetPANLinknetUPerNetSegformer
  • YOLO 系列模型:YOLOv8YOLOv10YOLO11YOLO12YOLO26

2. 本工程的优势 ✨

只需要替换 ultralytics/cfg/models/... 下的模型 YAML,就可以在相同数据集、相同训练入口和相同评估流程下完成不同结构的横向对比。

🌟 本框架最突出的特点,就是支持通过切换 YAML 快速完成不同结构对比实验。

3. 模块信息卡片 🧾

项目 内容
对应模型 YOLO26
改进模块 ACAB
方案编号 1
YAML 文件 ultralytics/cfg/models/improve/attention/yolo26/yolo26-ACAB-1.yaml
模块源码 ultralytics/nn/extra_modules/attention/ACAB.py
适用目标 在分割头输入前增强 P3 / P4 / P5 多尺度特征的通道与空间表达

4. 论文介绍 📚

文章地址:

论文核心内容 🧠

ISGLNet: Infrared Small Target Detection with Intrinsic Sensitivity and Guided Learning 聚焦红外小目标检测中“目标极弱、纹理极少、易被背景淹没”的难题。论文作者认为,这类任务的难点不只是提取更强特征,还包括让网络更稳定地区分真实弱目标与复杂背景噪声。

这篇工作强调两件事:一是提升网络对小目标固有敏感特征的响应能力,二是通过引导式学习机制减少误检和漏检。从整体思路上看,它更关心如何让模型在弱纹理、弱边界和高杂波环境中,依然保持较好的目标识别鲁棒性。

需要说明的是,当前工程中的 ACAB 并不是整篇论文网络的完整复现,而是将其中更适合插件化接入的“通道 + 空间重标定”思路封装成独立注意力模块,用于接入 YOLO26。这样做的好处很直接:你可以在统一训练入口下,仅通过切换 YAML,快速验证这类混合注意力机制对不同结构位置的增益。📌

模块价值总结 🌈

  • 先做通道重标定,再做空间重标定,结构清晰直接。
  • 适合作为轻量插件插入主干、颈部或输出前特征层。
  • 非常适合在本框架中通过切换 YAML 做结构对比实验。

改进步骤 🛠️

步骤1:定位并加入 ACAB 原始模块代码 📦

这一步需要先确认 ACAB 在当前工程中的真实源码位置,再将其作为可复用注意力模块接入 YOLO 体系,供后续导出、注册和 YAML 调用。

  • 原始代码位置:ultralytics/nn/extra_modules/attention/ACAB.py
  • 本步作用:将 ACAB 纳入 attention 目录,作为后续结构改进实验的基础组件。
import warnings
warnings.filterwarnings('ignore')
from calflops import calculate_flops

import math
import torch
import torch.nn as nn

class CA(nn.Module):
    def __init__(self, channel, b=1, gamma=2):
        super(CA, self).__init__()
        kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1

        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        y1 = self.avg_pool(x)
        y1 = self.conv(y1.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y2 = self.max_pool(x)
        y2 = self.conv(y2.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y = self.sigmoid(y1 + y2)
        return y.expand_as(x)

class SA(nn.Module):
    def __init__(self, kernel_size=7):
        super(SA, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv1(x)
        return self.sigmoid(x)

class ACAB(nn.Module):
    def __init__(self, channel, kernel_size=7, b=1, gamma=2):
        super().__init__()
        self.ca = CA(channel, b, gamma)
        self.sa = SA(kernel_size=kernel_size)

    def forward(self, x):
        x = self.ca(x) * x
        x = self.sa(x) * x
        return x

步骤2:在 ultralytics/nn/extra_modules/__init__.py 中导入模块 🧩

from .attention.ACAB import ACAB

步骤3:在 ultralytics/nn/tasks.py 中注册模块 ⚙️

extra_modules.ACAB,

步骤4:新建 YOLO26 对应的改进 YAML 文件 📄

当前方案对应的配置文件为 ultralytics/cfg/models/improve/attention/yolo26/yolo26-ACAB-1.yaml。这一版的做法,是在分割输出前分别对 P3P4P5 三个尺度特征加入 ACAB,用于增强多尺度特征的通道与空间选择能力。

head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, Concat, [1]]
  - [-1, 2, C3k2, [512, True]]

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]]
  - [-1, 2, C3k2, [256, True]]

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 13], 1, Concat, [1]]
  - [-1, 2, C3k2, [512, True]]

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 10], 1, Concat, [1]]
  - [-1, 1, C3k2, [1024, True, 0.5, True]]

  - [16, 1, ACAB, []]
  - [19, 1, ACAB, []]
  - [22, 1, ACAB, []]

  - [[23, 24, 25], 1, SemanticSegmentHead, [nc]]

步骤5:开始训练 🚀

# -*- coding: utf-8 -*-
"""
@Auth :AICurator
@File :train.py
"""
import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO(model=r'G:\improve\segment\ultralytics-main\ultralytics\cfg\models\improve\attention\yolo26\yolo26-ACAB-1.yaml')
  # model.load('')
    model.train(data=r'dataset\data.yaml',
                imgsz=640,
                epochs=50,
                batch=4,
                workers=0,
                device='0',
                optimizer='SGD',
                close_mosaic=10,
                resume=False,
                project='runs/train',
                name='exp',
                single_cls=False,
                cache=False,
                )

6. 总结 🎯

这一版 yolo26-ACAB-1 的重点,是在分割头前对三个尺度特征分别做一次 ACAB 通道空间混合增强。你可以直接通过切换 YAML,快速比较 ACAB 接入前后的分割表现差异。✅

订阅专栏后添加博主微信,领取完整代码 📮

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐