当前位置：首页 > news >正文

yolov26改进 | 添加注意力机制篇 | 添加TripletAttention三重注意力机制（附代码+机制原理+添加教程+网络结构图）

news 2026/6/1 2:14:54

开始讲解之前推荐一下我的专栏，本专栏的内容支持(分类、检测、分割、追踪、关键点检测),专栏目前为限时折扣，欢迎大家订阅本专栏，本专栏每周更新5-7篇最新机制，更有包含我所有改进的文件和交流群提供给大家，本人定期在群内分享发表论文方法和经验。

一、本文介绍

本文给大家带来的改进是Triplet Attention三重注意力机制。这个机制，它通过三个不同的视角来分析输入的数据，就好比三个人从不同的角度来观察同一幅画，然后共同决定哪些部分最值得注意。三重注意力机制的主要思想是在网络中引入了一种新的注意力模块，这个模块包含三个分支，分别关注图像的不同维度。比如说，一个分支可能专注于图像的宽度，另一个分支专注于高度，第三个分支则聚焦于图像的深度，即色彩和纹理等特征。这样一来，网络就能够更全面地理解图像内容，就像是得到了一副三维眼镜，能够看到图片的立体效果一样。

专栏链接：YOLOv26有效涨点专栏包含：Conv、注意力机制、主干/Backbone、损失函数、优化器、后处理等改进机制

一、本文介绍

二、Triplet Attention机制原理

2.1 Triplet Attention的基本原理

2.2 Triplet Attention和其它简单注意力机制的对比

2.3 Triplet Attention的实现流程

三、Triplet Attention的完整代码

四、手把手教你添加Triplet Attention

4.1 修改一

4.2 修改二

4.3 修改三

4.4 修改四

4.5 修改五

4.6 修改六

五、正式训练

5.1 yaml文件

5.1.1 yaml文件1

5.1.2 yaml文件2

5.2 训练代码

5.3 训练过程截图

五、本文总结

二、Triplet Attention机制原理

论文地址：官方论文地址

代码地址：官方代码地址

2.1 Triplet Attention的基本原理

三重注意力（Triplet Attention）的基本原理是利用三支结构捕获输入数据的跨维度交互，从而计算注意力权重。这个方法能够构建输入通道或空间位置之间的相互依赖性，而且计算代价小。三重注意力由三个分支组成，每个分支负责捕获空间维度H或W与通道维度C之间的交互特征。通过对每个分支中的输入张量进行排列变换，然后通过Z池操作和一个大小为k×k的卷积层，生成注意力权重。这些权重是通过一个S形激活层生成的，然后应用于排列变换后的输入张量，再变换回原来的输入形状

三重注意力（Triplet Attention）的主要改进点包括：

跨维度的注意力权重计算: 通过一个创新的三支结构捕获通道、高度、宽度三个维度之间的交互关系来计算注意力权重。
旋转操作和残差变换: 通过旋转输入张量和应用残差变换来建立不同维度间的依赖，这是三重注意力机制中的关键步骤。
维度间依赖性的重要性: 强调在计算注意力权重时，捕获跨维度依赖性的重要性，这是三重注意力的核心直觉和设计理念。

下面的图片是三重注意力的一个抽象表示图，展示了三个分支如何捕获跨维度交互。图中的每个子图表示三重注意力中的一个分支：

1. 分支（a）:这个分支直接处理输入张量，没有进行旋转，然后通过残差变换来提取特征。

2. 分支（b）:这个分支首先沿着宽度（W）和通道（C）的维度旋转输入张量，然后进行残差变换。

3. 分支（c）:这个分支沿着高度（H）和通道（C）的维度旋转输入张量，之后同样进行残差变换。

总结：通过这样的设计，三重注意力模型能够有效地捕获输入张量中的空间和通道维度之间的交互关系。这种方法使模型能够构建通道与空间位置之间的相互依赖性，提高模型对特征的理解能力。

2.2 Triplet Attention和其它简单注意力机制的对比

下面的图片是论文中三重注意力机制和其它注意力机制的一个对比大家有兴趣可以看看，横向扩展以下自己的知识库。

这张图片是一幅对比不同注意力模块的图示，其中包括：

1.Squeeze Excitation (SE) Module:
这个模块使用全局平均池化 (Global Avg Pool) 生成通道描述符，接着通过两个全连接层（1x1 Conv），中间使用ReLU激活函数，最后通过Sigmoid函数生成每个通道的权重。

2. Convolutional Block Attention Module (CBAM):
首先使用全局平均池化和全局最大池化（GAP + GMP）结合，再通过一个卷积层和ReLU激活函数，最后经过另一个卷积层和Sigmoid函数生成注意力权重。

3. Global Context (GC) Module:
从一个1x1卷积层开始，经过Softmax函数进行归一化，接着进行另一个1x1卷积，然后使用LayerNorm和最终的1x1卷积，通过广播加法结合原始特征图。

4. Triplet Attention (我们的方法):
分为三个分支，每个分支进行不同的处理：通道池化后的7x7卷积，Z池化，再接一个7x7卷积，然后是批量归一化和Sigmoid函数。每个分支都有一个Permute操作来调整维度。最后，三个分支的结果通过平均池化聚合起来生成最终的注意力权重。

每种模块都设计用于处理特征图（C x H x W），其中C是通道数，H是高度，W是宽度。这些模块通过不同方式计算注意力权重，增强网络对特征的重要部分的关注度，从而在各种视觉任务中提高性能。图片中的符号⊗代表矩阵乘法，⊕代表广播元素级加法。

2.3 Triplet Attention的实现流程

下面的图片是三重注意力（Triplet Attention）的具体实现流程图。图中详细展示了三个分支如何处理输入张量，并最终合成三重注意力。下面是对这个过程的描述：

上部分支: 负责计算通道维度C和空间维度W的注意力权重。这个分支对输入张量进行Z池化（Z-Pool）操作，然后通过一个卷积层（Conv），接着用Sigmoid函数生成注意力权重。
中部分支: 负责捕获通道维度C与空间维度H和W之间的依赖性。这个分支首先进行相同的Z池化和卷积操作，然后同样通过Sigmoid函数生成注意力权重。
下部分支: 用于捕获空间维度之间的依赖性。这个分支保持输入的身份（Identity，即不改变输入），执行Z池化和卷积操作，之后也通过Sigmoid函数生成注意力权重。

每个分支在生成注意力权重后，会对输入进行排列（Permutation），然后将三个分支的输出进行平均聚合（Avg），最终得到三重注意力输出。

这种结构通过不同的旋转和排列操作，能够综合不同维度上的信息，更好地捕获数据的内在特征，同时这种方法在计算上是高效的，并且可以作为一个模块加入到现有的网络架构中，增强网络对复杂数据结构的理解和处理能力。

三、Triplet Attention的完整代码

核心代码的使用方式看章节四！

import torch import torch.nn as nn __all__ = ['C2PSA_TripleAttention', 'TripletAttention'] class BasicConv(nn.Module): def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False): super(BasicConv, self).__init__() self.out_channels = out_planes self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias) self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None self.relu = nn.ReLU() if relu else None def forward(self, x): x = self.conv(x) if self.bn is not None: x = self.bn(x) if self.relu is not None: x = self.relu(x) return x class ZPool(nn.Module): def forward(self, x): return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1) class AttentionGate(nn.Module): def __init__(self): super(AttentionGate, self).__init__() kernel_size = 7 self.compress = ZPool() self.conv = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size - 1) // 2, relu=False) def forward(self, x): x_compress = self.compress(x) x_out = self.conv(x_compress) scale = torch.sigmoid_(x_out) return x * scale class TripletAttention(nn.Module): def __init__(self, no_spatial=False): super(TripletAttention, self).__init__() self.cw = AttentionGate() self.hc = AttentionGate() self.no_spatial = no_spatial if not no_spatial: self.hw = AttentionGate() def forward(self, x): x_perm1 = x.permute(0, 2, 1, 3).contiguous() x_out1 = self.cw(x_perm1) x_out11 = x_out1.permute(0, 2, 1, 3).contiguous() x_perm2 = x.permute(0, 3, 2, 1).contiguous() x_out2 = self.hc(x_perm2) x_out21 = x_out2.permute(0, 3, 2, 1).contiguous() if not self.no_spatial: x_out = self.hw(x) x_out = 1 / 3 * (x_out + x_out11 + x_out21) else: x_out = 1 / 2 * (x_out11 + x_out21) return x_out def autopad(k, p=None, d=1): # kernel, padding, dilation """Pad to 'same' shape outputs.""" if d > 1: k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size if p is None: p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad return p class Conv(nn.Module): """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation).""" default_act = nn.SiLU() # default activation def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True): """Initialize Conv layer with given arguments including activation.""" super().__init__() self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False) self.bn = nn.BatchNorm2d(c2) self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity() def forward(self, x): """Apply convolution, batch normalization and activation to input tensor.""" return self.act(self.bn(self.conv(x))) def forward_fuse(self, x): """Perform transposed convolution of 2D data.""" return self.act(self.conv(x)) def autopad(k, p=None, d=1): # kernel, padding, dilation """Pad to 'same' shape outputs.""" if d > 1: k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k] # actual kernel-size if p is None: p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad return p class Conv(nn.Module): """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation).""" default_act = nn.SiLU() # default activation def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True): """Initialize Conv layer with given arguments including activation.""" super().__init__() self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False) self.bn = nn.BatchNorm2d(c2) self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity() def forward(self, x): """Apply convolution, batch normalization and activation to input tensor.""" return self.act(self.bn(self.conv(x))) def forward_fuse(self, x): """Perform transposed convolution of 2D data.""" return self.act(self.conv(x)) class PSABlock(nn.Module): """ PSABlock class implementing a Position-Sensitive Attention block for neural networks. This class encapsulates the functionality for applying multi-head attention and feed-forward neural network layers with optional shortcut connections. Attributes: attn (Attention): Multi-head attention module. ffn (nn.Sequential): Feed-forward neural network module. add (bool): Flag indicating whether to add shortcut connections. Methods: forward: Performs a forward pass through the PSABlock, applying attention and feed-forward layers. Examples: Create a PSABlock and perform a forward pass >>> psablock = PSABlock(c=128, attn_ratio=0.5, num_heads=4, shortcut=True) >>> input_tensor = torch.randn(1, 128, 32, 32) >>> output_tensor = psablock(input_tensor) """ def __init__(self, c, attn_ratio=0.5, num_heads=4, shortcut=True) -> None: """Initializes the PSABlock with attention and feed-forward layers for enhanced feature extraction.""" super().__init__() self.attn = TripletAttention() self.ffn = nn.Sequential(Conv(c, c * 2, 1), Conv(c * 2, c, 1, act=False)) self.add = shortcut def forward(self, x): """Executes a forward pass through PSABlock, applying attention and feed-forward layers to the input tensor.""" x = x + self.attn(x) if self.add else self.attn(x) x = x + self.ffn(x) if self.add else self.ffn(x) return x class C2PSA_TripleAttention(nn.Module): """ C2PSA module with attention mechanism for enhanced feature extraction and processing. This module implements a convolutional block with attention mechanisms to enhance feature extraction and processing capabilities. It includes a series of PSABlock modules for self-attention and feed-forward operations. Attributes: c (int): Number of hidden channels. cv1 (Conv): 1x1 convolution layer to reduce the number of input channels to 2*c. cv2 (Conv): 1x1 convolution layer to reduce the number of output channels to c. m (nn.Sequential): Sequential container of PSABlock modules for attention and feed-forward operations. Methods: forward: Performs a forward pass through the C2PSA module, applying attention and feed-forward operations. Notes: This module essentially is the same as PSA module, but refactored to allow stacking more PSABlock modules. Examples: >>> c2psa = C2PSA(c1=256, c2=256, n=3, e=0.5) >>> input_tensor = torch.randn(1, 256, 64, 64) >>> output_tensor = c2psa(input_tensor) """ def __init__(self, c1, c2, n=1, e=0.5): """Initializes the C2PSA module with specified input/output channels, number of layers, and expansion ratio.""" super().__init__() assert c1 == c2 self.c = int(c1 * e) self.cv1 = Conv(c1, 2 * self.c, 1, 1) self.cv2 = Conv(2 * self.c, c1, 1) self.m = nn.Sequential(*(PSABlock(self.c, attn_ratio=0.5, num_heads=self.c // 64) for _ in range(n))) def forward(self, x): """Processes the input tensor 'x' through a series of PSA blocks and returns the transformed tensor.""" a, b = self.cv1(x).split((self.c, self.c), dim=1) b = self.m(b) return self.cv2(torch.cat((a, b), 1)) if __name__ == "__main__": # Generating Sample image image_size = (1, 64, 224, 224) image = torch.rand(*image_size) # Model model = C2PSA_TripleAttention(64, 64) out = model(image) print(out.size())

四、手把手教你添加Triplet Attention

下面的步骤如果你不会或者不想麻烦操作，可以联系作者获得本专栏添加所有项目文件的源代码，可直接训练.

4.1 修改一

第一还是建立文件，我们找到如下ultralytics/nn文件夹下建立一个目录名字呢就是'Addmodules'文件夹！

4.2 修改二

然后在Addmodules文件夹内建立一个新的py文件，将本文章节三中的“核心代码"复制粘贴进去。

4.3 修改三

第二步我们在该目录下创建一个新的py文件名字为'__init__.py'，然后在其内部导入我们的文件，如下图所示。

4.4 修改四

第三步我门中到如下文件'ultralytics/nn/tasks.py'进行导入和注册我们的模块(此处只需要添加一次即可，如果你用我其它的改进机制这里的步骤只需要添加一次)！

4.5 修改五

在'ultralytics/nn/tasks.py'文件内的parse_model方法函数内（位置大概在1500+行左右），按照图示位置添加即可（此处需要自己有一定的判别能力，如果不会可联系作者获得视频教程）。

4.6 修改六

在'ultralytics/nn/tasks.py'文件内的parse_model方法函数内（位置大概在1550+行左右），按照图示位置添加即可，此处一定要对应好位置和缩进否则很容易报错。

elif m in {此处填写本章代码的名字.}: c2 = ch[f] args = [c2, *args]

五、正式训练

到此就修改完成了，大家可以复制下面的yaml文件运行，更多使用方式可以联系作者获得使用视频，本文仅列出常见的使用方式。

5.1 yaml文件

5.1.1 yaml文件1

训练信息：YOLO26-C2PSA-TripletAtt summary: 261 layers, 2,455,112 parameters, 2,455,112 gradients, 5.7 GFLOPs

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license # Ultralytics YOLO26 object detection model with P3/8 - P5/32 outputs # Model docs: https://docs.ultralytics.com/models/yolo26 # Task docs: https://docs.ultralytics.com/tasks/detect # Parameters nc: 80 # number of classes end2end: True # whether to use end-to-end mode reg_max: 1 # DFL bins scales: # model compound scaling constants, i.e. 'model=yolo26n.yaml' will call yolo26.yaml with scale 'n' # [depth, width, max_channels] n: [0.50, 0.25, 1024] # summary: 260 layers, 2,572,280 parameters, 2,572,280 gradients, 6.1 GFLOPs s: [0.50, 0.50, 1024] # summary: 260 layers, 10,009,784 parameters, 10,009,784 gradients, 22.8 GFLOPs m: [0.50, 1.00, 512] # summary: 280 layers, 21,896,248 parameters, 21,896,248 gradients, 75.4 GFLOPs l: [1.00, 1.00, 512] # summary: 392 layers, 26,299,704 parameters, 26,299,704 gradients, 93.8 GFLOPs x: [1.00, 1.50, 512] # summary: 392 layers, 58,993,368 parameters, 58,993,368 gradients, 209.5 GFLOPs # YOLO26n backbone backbone: # [from, repeats, module, args] - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 - [-1, 2, C3k2, [256, False, 0.25]] - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 - [-1, 2, C3k2, [512, False, 0.25]] - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 - [-1, 2, C3k2, [512, True]] - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 - [-1, 2, C3k2, [1024, True]] - [-1, 1, SPPF, [1024, 5, 3, True]] # 9 - [-1, 2, C2PSA_TripleAttention, [1024]] # 10 # YOLO26n head head: - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 6], 1, Concat, [1]] # cat backbone P4 - [-1, 2, C3k2, [512, True]] # 13 - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 4], 1, Concat, [1]] # cat backbone P3 - [-1, 2, C3k2, [256, True]] # 16 (P3/8-small) - [-1, 1, Conv, [256, 3, 2]] - [[-1, 13], 1, Concat, [1]] # cat head P4 - [-1, 2, C3k2, [512, True]] # 19 (P4/16-medium) - [-1, 1, Conv, [512, 3, 2]] - [[-1, 10], 1, Concat, [1]] # cat head P5 - [-1, 1, C3k2, [1024, True, 0.5, True]] # 22 (P5/32-large) - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.1.2 yaml文件2

训练信息：YOLO26-Att-TripletAtt summary: 266 layers, 2,506,340 parameters, 2,506,340 gradients, 5.8 GFLOPs

# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license # Ultralytics YOLO26 object detection model with P3/8 - P5/32 outputs # Model docs: https://docs.ultralytics.com/models/yolo26 # Task docs: https://docs.ultralytics.com/tasks/detect # Parameters nc: 80 # number of classes end2end: True # whether to use end-to-end mode reg_max: 1 # DFL bins scales: # model compound scaling constants, i.e. 'model=yolo26n.yaml' will call yolo26.yaml with scale 'n' # [depth, width, max_channels] n: [0.50, 0.25, 1024] # summary: 260 layers, 2,572,280 parameters, 2,572,280 gradients, 6.1 GFLOPs s: [0.50, 0.50, 1024] # summary: 260 layers, 10,009,784 parameters, 10,009,784 gradients, 22.8 GFLOPs m: [0.50, 1.00, 512] # summary: 280 layers, 21,896,248 parameters, 21,896,248 gradients, 75.4 GFLOPs l: [1.00, 1.00, 512] # summary: 392 layers, 26,299,704 parameters, 26,299,704 gradients, 93.8 GFLOPs x: [1.00, 1.50, 512] # summary: 392 layers, 58,993,368 parameters, 58,993,368 gradients, 209.5 GFLOPs # YOLO26n backbone backbone: # [from, repeats, module, args] - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 - [-1, 2, C3k2, [256, False, 0.25]] - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 - [-1, 2, C3k2, [512, False, 0.25]] - [-1, 1, Conv, [512, 3, 2]] # 5-P4/16 - [-1, 2, C3k2, [512, True]] - [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32 - [-1, 2, C3k2, [1024, True]] - [-1, 1, SPPF, [1024, 5, 3, True]] # 9 - [-1, 2, C2PSA, [1024]] # 10 # YOLO26n head head: - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 6], 1, Concat, [1]] # cat backbone P4 - [-1, 2, C3k2, [512, True]] # 13 - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 4], 1, Concat, [1]] # cat backbone P3 - [-1, 2, C3k2, [256, True]] # 16 (P3/8-small) - [-1, 1, Conv, [256, 3, 2]] - [[-1, 13], 1, Concat, [1]] # cat head P4 - [-1, 2, C3k2, [512, True]] # 19 (P4/16-medium) - [-1, 1, Conv, [512, 3, 2]] - [[-1, 10], 1, Concat, [1]] # cat head P5 - [-1, 1, C3k2, [1024, True, 0.5, True]] # 22 (P5/32-large) - [16, 1, TripletAttention, []] # 23 # - [19, 1, TripletAttention, []] # 24 # - [22, 1, TripletAttention, []] # 25 # 此处的使用说法注释: 其中上面的三个注意力机制目前仅使用了23层，如果你想使用24层那么就取消掉代码注释， # 并将下面检测头中的19改为24,如果想使用第25层注意力机制同理，将下面检测头中的22改为25即可。 # 此处用法比较复杂如过不会联系Snu77博主获取视频教程 - [[23, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)

5.2 训练代码

大家可以创建一个py文件将我给的代码复制粘贴进去，配置好自己的文件路径即可运行。

import warnings warnings.filterwarnings('ignore') from ultralytics import YOLO if __name__ == '__main__': model = YOLO('模型配置文件地址,也就是5.1你保存到本地文件的地址') # 如何切换模型版本, 上面的ymal文件可以改为 yolo26s.yaml就是使用的26s, # 类似某个改进的yaml文件名称为yolo26-XXX.yaml那么如果想使用其它版本就把上面的名称改为yolo26l-XXX.yaml即可（改的是上面YOLO中间的名字不是配置文件的）！ # model.load('yolo26n.pt') # 是否加载预训练权重,科研不建议大家加载否则很难提升精度 model.train( data=r"数据集文件地址", # 如果大家任务是其它的'ultralytics/cfg/default.yaml'找到这里修改task可以改成detect, segment, classify, pose cache=False, imgsz=640, epochs=20, single_cls=False, # 是否是单类别检测 batch=16, close_mosaic=0, workers=0, device='0', optimizer='MuSGD', # using SGD/MuSGD # resume=, # 这里是填写last.pt地址 amp=True, # 如果出现训练损失为Nan可以关闭amp project='runs/train', name='exp', )