YOLOv8改进 | 细节创新篇 | iAFF迭代注意力特征融合助力多目标细节涨点

import torch
import torch.nn as nndef 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-sizeif p is None:p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-padreturn pclass Conv(nn.Module):"""Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""default_act = nn.SiLU()  # default activationdef __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 iAFF(nn.Module):'''多特征融合 iAFF'''def __init__(self, channels=64, r=2):super(iAFF, self).__init__()inter_channels = int(channels // r)# 本地注意力self.local_att = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)# 全局注意力self.global_att = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),)# 第二次本地注意力self.local_att2 = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)# 第二次全局注意力self.global_att2 = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.sigmoid = nn.Sigmoid()def forward(self, x, residual):xa = x + residualxl = self.local_att(xa)xg = self.global_att(xa)xlg = xl + xgwei = self.sigmoid(xlg)xi = x * wei + residual * (1 - wei)xl2 = self.local_att2(xi)xg2 = self.global_att(xi)xlg2 = xl2 + xg2wei2 = self.sigmoid(xlg2)xo = x * wei2 + residual * (1 - wei2)return xoclass AFF(nn.Module):'''多特征融合 AFF'''def __init__(self, channels=64, r=4):super(AFF, self).__init__()inter_channels = int(channels // r)self.local_att = nn.Sequential(nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.global_att = nn.Sequential(nn.AdaptiveAvgPool2d(1),nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(inter_channels),nn.ReLU(inplace=True),nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),nn.BatchNorm2d(channels),)self.sigmoid = nn.Sigmoid()def forward(self, x, residual):xa = x + residualxl = self.local_att(xa)xg = self.global_att(xa)xlg = xl + xgwei = self.sigmoid(xlg)xo = 2 * x * wei + 2 * residual * (1 - wei)return xoclass C2f_iAFF(nn.Module):"""Faster Implementation of CSP Bottleneck with 2 convolutions."""def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):"""Initialize CSP bottleneck layer with two convolutions with arguments ch_in, ch_out, number, shortcut, groups,expansion."""super().__init__()self.c = int(c2 * e)  # hidden channelsself.cv1 = Conv(c1, 2 * self.c, 1, 1)self.cv2 = Conv((2 + n) * self.c, c2, 1)  # optional act=FReLU(c2)self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n))def forward(self, x):"""Forward pass through C2f layer."""y = list(self.cv1(x).chunk(2, 1))y.extend(m(y[-1]) for m in self.m)return self.cv2(torch.cat(y, 1))def forward_split(self, x):"""Forward pass using split() instead of chunk()."""y = list(self.cv1(x).split((self.c, self.c), 1))y.extend(m(y[-1]) for m in self.m)return self.cv2(torch.cat(y, 1))class Bottleneck(nn.Module):"""Standard bottleneck."""def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):"""Initializes a bottleneck module with given input/output channels, shortcut option, group, kernels, andexpansion."""super().__init__()c_ = int(c2 * e)  # hidden channelsself.cv1 = Conv(c1, c_, k[0], 1)self.cv2 = Conv(c_, c2, k[1], 1, g=g)self.add = shortcut and c1 == c2self.iAFF = iAFF(c2)def forward(self, x):"""'forward()' applies the YOLO FPN to input data."""if self.add:results =  self.iAFF(x , self.cv2(self.cv1(x)))else:results = self.cv2(self.cv1(x))return resultsif __name__ == '__main__':x = torch.ones(8, 64, 32, 32)channels = x.shape[1]model = C2f_iAFF(channels, channels, True)output = model(x)print(output.shape)