修改1：Conv替换LAE
0）LSM-YOLO提出的轻量化自适应特征提取(LAE)模块，旨在多尺度、自适应的提取特征，全局与局部特征融合，并降低计算成本。
核心代码（fxlae.py，可复制下面代码）：
| import torch import torch.nn as nn from einops import rearrange
__all__ = ['LAE', 'MSFM']
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 LAE(nn.Module): # Light-weight Adaptive Extraction def __init__(self, ch, group=16) -> None: super().__init__()
self.softmax = nn.Softmax(dim=-1) self.attention = nn.Sequential( nn.AvgPool2d(kernel_size=3, stride=1, padding=1), Conv(ch, ch, k=1) )
self.ds_conv = Conv(ch, ch * 4, k=3, s=2, g=(ch // group))
def forward(self, x): # bs, ch, 2*h, 2*w => bs, ch, h, w, 4 att = rearrange(self.attention(x), 'bs ch (s1 h) (s2 w) -> bs ch h w (s1 s2)', s1=2, s2=2) att = self.softmax(att)
# bs, 4 * ch, h, w => bs, ch, h, w, 4 x = rearrange(self.ds_conv(x), 'bs (s ch) h w -> bs ch h w s', s=4) x = torch.sum(x * att, dim=-1) return x
class MatchNeck_Inner(nn.Module): def __init__(self, channels) -> None: super().__init__()
self.gap = nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), Conv(channels, channels) ) self.pool_h = nn.AdaptiveAvgPool2d((None, 1)) self.pool_w = nn.AdaptiveAvgPool2d((1, None)) self.conv_hw = Conv(channels, channels, (3, 1)) self.conv_pool_hw = Conv(channels, channels, 1)
def forward(self, x): _, _, h, w = x.size() x_pool_h, x_pool_w, x_pool_ch = self.pool_h(x), self.pool_w(x).permute(0, 1, 3, 2), self.gap(x) x_pool_hw = torch.cat([x_pool_h, x_pool_w], dim=2) x_pool_h, x_pool_w = torch.split(x_pool_hw, [h, w], dim=2) x_pool_hw_weight = x_pool_hw.sigmoid() x_pool_h_weight, x_pool_w_weight = torch.split(x_pool_hw_weight, [h, w], dim=2) x_pool_h, x_pool_w = x_pool_h * x_pool_h_weight, x_pool_w * x_pool_w_weight x_pool_ch = x_pool_ch * torch.mean(x_pool_hw_weight, dim=2, keepdim=True) return x * x_pool_h.sigmoid() * x_pool_w.permute(0, 1, 3, 2).sigmoid() * x_pool_ch.sigmoid()
class MatchNeck(nn.Module): def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5): super().__init__() c_ = int(c2 * e) # hidden channels self.cv1 = Conv(c1, c_, k[0], 1) self.cv2 = Conv(c_, c2, k[1], 1, g=g) self.add = shortcut and c1 == c2 self.MN = MatchNeck_Inner(c2)
def forward(self, x): return x + self.MN(self.cv2(self.cv1(x))) if self.add else self.MN(self.cv2(self.cv1(x)))
class MSFM(nn.Module): def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): super().__init__() self.c = int(c2 * e) # hidden channels self.cv1 = Conv(c1, 2 * self.c, 1, 1) self.cv2 = Conv((2 + n) * self.c, c2, 1) self.m = nn.ModuleList(MatchNeck(self.c, self.c, shortcut, g, k=(3, 3), e=1.0) for _ in range(n))
def forward(self, x): 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): 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))
if __name__ == "__main__": # Generating Sample image image_size = (1, 64, 224, 224) image = torch.rand(*image_size) # Model model = LAE(64) out = model(image) print(out.size()) |
执行上面程序，可以看出：(1, 64, 224, 224)→(1, 64, 112, 112)，说明通道数没变，h、w减小1/2。
1）将fxlae.py文件放到ultralytics/nn/modules文件夹中。
2）修改ultralytics/nn/modules/__init__.py，导入LAE模块，即增加88行和99行的内容。
3）修改ultralytics/nn/tasks.py，有两处，即增加19行、1038-1039行的内容。
4）以yolo11.yaml文件为基础重新写一个yolo11-fxlae.yaml文件，位置如下：
ultralytics/cfg/models/11/yolo11-fxlae.yaml
| # Parameters nc: 80 # number of classes scales: # model compound scaling constants # [depth, width, max_channels] n: [0.50, 0.25, 1024] s: [0.50, 0.50, 1024] m: [0.50, 1.00, 512] l: [1.00, 1.00, 512] x: [1.00, 1.50, 512] # YOLO11n-fxlae 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, LAE, []] # 3-P3/8 - [-1, 2, C3k2, [512, False, 0.25]] - [-1, 1, LAE, []] # 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]] # 9 - [-1, 2, C2PSA, [1024]] # 10 # YOLO11n-fxlae head head: - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 6], 1, Concat, [1]] # cat backbone P4 - [-1, 2, C3k2, [512, False]] # 13 - [-1, 1, nn.Upsample, [None, 2, "nearest"]] - [[-1, 4], 1, Concat, [1]] # cat backbone P3 - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small) - [-1, 1, LAE, []] - [[-1, 13], 1, Concat, [1]] # cat head P4 - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium) - [-1, 1, LAE, []] - [[-1, 10], 1, Concat, [1]] # cat head P5 - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large) - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5) |
注意，替代的通道数不变的Conv。
5）准备好数据集，并针对数据集写一个neu-det.yaml文件，如下：
| train: NEU-DET/train.txt val: NEU-DET/val.txt test: NEU-DET/test.txt #number of classes nc: 6 #class names names: ["crazing", "inclusion", "patches", "pitted_surface", "rolled-in_scale", "scratches"] |
NEU-DET文件夹下有：image文件夹、label文件夹、train.txt、val.txt、test.txt，image文件夹里是数据集的所有图片，labels文件夹里是每一个图片的txt文件（1个类别+4个目标框信息），3个txt文件里存储的是用于train、val、test的图片绝对路径。
NEU-DET和neu-det.yaml都放在仓库根目录下data文件夹里。
说明：这一部分内容也可以看我的另一篇博文（1. YOLO11首次使用）的第五部分。
6）在仓库根目录下写一个train.py，可以修改ultralytics/cfg/default.yaml里的设置与参数。
| import warnings warnings.filterwarnings('ignore') from ultralytics import YOLO
if __name__ == '__main__': model = YOLO('yolo11-fxlae.yaml') model.load('yolo11n.pt') # 加载预训练权重 model.train(data='./data/neu-det.yaml', cache=False, imgsz=640, epochs=300, single_cls=False, batch=16, close_mosaic=0, workers=0, device='0', optimizer='SGD', resume=False, amp=False, # 如果出现训练损失为Nan可以关闭amp project='runs/train', name='exp-fxlae', ) |
python train.py即可训练。
下面是一个val.py的代码：
| import warnings warnings.filterwarnings('ignore') from ultralytics import YOLO if __name__ == '__main__': model = YOLO('runs/train/exp-fxlae/weights/best.pt') model.val(data='./data/neu-det.yaml', split='val', # 测试可以改为 split='test', imgsz=640, batch=16, mode='val', # 测试可以改为 mode='test', save=True, project='runs/val', # 测试可以改为 project='runs/test', name='exp-fxlae', ) |
python val.py即可验证或测试。
7）对比修改前后的网络参数：
修改前：
修改后：
-------------------------------------------------------------------------------------------------------------------
修改2：Head改为AFPN
0）AFPN引入了一种渐进融合的策略，且采用了4个检测头，AFPN包含了Neck部分；
核心代码（fxafpn.py，可复制下面代码）：
| import math from collections import OrderedDict import torch import torch.nn as nn import torch.nn.functional as F from ultralytics.nn.modules import DFL from ultralytics.nn.modules.conv import Conv from ultralytics.utils.tal import dist2bbox, make_anchors
__all__ =['Detect_AFPN']
def BasicConv(filter_in, filter_out, kernel_size, stride=1, pad=None): if not pad: pad = (kernel_size - 1) // 2 if kernel_size else 0 else: pad = pad return nn.Sequential(OrderedDict([ ("conv", nn.Conv2d(filter_in, filter_out, kernel_size=kernel_size, stride=stride, padding=pad, bias=False)), ("bn", nn.BatchNorm2d(filter_out)), ("relu", nn.ReLU(inplace=True)), ]))
class BasicBlock(nn.Module): expansion = 1
def __init__(self, filter_in, filter_out): super(BasicBlock, self).__init__() self.conv1 = nn.Conv2d(filter_in, filter_out, 3, padding=1) self.bn1 = nn.BatchNorm2d(filter_out, momentum=0.1) self.relu = nn.ReLU(inplace=True) self.conv2 = nn.Conv2d(filter_out, filter_out, 3, padding=1) self.bn2 = nn.BatchNorm2d(filter_out, momentum=0.1)
def forward(self, x): residual = x
out = self.conv1(x) out = self.bn1(out) out = self.relu(out)
out = self.conv2(out) out = self.bn2(out)
out += residual out = self.relu(out)
return out
class Upsample(nn.Module): def __init__(self, in_channels, out_channels, scale_factor=2): super(Upsample, self).__init__()
self.upsample = nn.Sequential( BasicConv(in_channels, out_channels, 1), nn.Upsample(scale_factor=scale_factor, mode='bilinear') )
def forward(self, x): x = self.upsample(x)
return x
class Downsample_x2(nn.Module): def __init__(self, in_channels, out_channels): super(Downsample_x2, self).__init__()
self.downsample = nn.Sequential( BasicConv(in_channels, out_channels, 2, 2, 0) )
def forward(self, x, ): x = self.downsample(x)
return x
class Downsample_x4(nn.Module): def __init__(self, in_channels, out_channels): super(Downsample_x4, self).__init__()
self.downsample = nn.Sequential( BasicConv(in_channels, out_channels, 4, 4, 0) )
def forward(self, x, ): x = self.downsample(x)
return x
class Downsample_x8(nn.Module): def __init__(self, in_channels, out_channels): super(Downsample_x8, self).__init__()
self.downsample = nn.Sequential( BasicConv(in_channels, out_channels, 8, 8, 0) )
def forward(self, x, ): x = self.downsample(x)
return x
class ASFF_2(nn.Module): def __init__(self, inter_dim=512): super(ASFF_2, self).__init__()
self.inter_dim = inter_dim compress_c = 8
self.weight_level_1 = BasicConv(self.inter_dim, compress_c, 1, 1) self.weight_level_2 = BasicConv(self.inter_dim, compress_c, 1, 1)
self.weight_levels = nn.Conv2d(compress_c * 2, 2, kernel_size=1, stride=1, padding=0)
self.conv = BasicConv(self.inter_dim, self.inter_dim, 3, 1)
def forward(self, input1, input2): level_1_weight_v = self.weight_level_1(input1) level_2_weight_v = self.weight_level_2(input2)
levels_weight_v = torch.cat((level_1_weight_v, level_2_weight_v), 1) levels_weight = self.weight_levels(levels_weight_v) levels_weight = F.softmax(levels_weight, dim=1)
fused_out_reduced = input1 * levels_weight[:, 0:1, :, :] + \ input2 * levels_weight[:, 1:2, :, :]
out = self.conv(fused_out_reduced)
return out
class ASFF_3(nn.Module): def __init__(self, inter_dim=512): super(ASFF_3, self).__init__()
self.inter_dim = inter_dim compress_c = 8
self.weight_level_1 = BasicConv(self.inter_dim, compress_c, 1, 1) self.weight_level_2 = BasicConv(self.inter_dim, compress_c, 1, 1) self.weight_level_3 = BasicConv(self.inter_dim, compress_c, 1, 1)
self.weight_levels = nn.Conv2d(compress_c * 3, 3, kernel_size=1, stride=1, padding=0)
self.conv = BasicConv(self.inter_dim, self.inter_dim, 3, 1)
def forward(self, input1, input2, input3): level_1_weight_v = self.weight_level_1(input1) level_2_weight_v = self.weight_level_2(input2) level_3_weight_v = self.weight_level_3(input3)
levels_weight_v = torch.cat((level_1_weight_v, level_2_weight_v, level_3_weight_v), 1) levels_weight = self.weight_levels(levels_weight_v) levels_weight = F.softmax(levels_weight, dim=1)
fused_out_reduced = input1 * levels_weight[:, 0:1, :, :] + \ input2 * levels_weight[:, 1:2, :, :] + \ input3 * levels_weight[:, 2:, :, :]
out = self.conv(fused_out_reduced)
return out
class ASFF_4(nn.Module): def __init__(self, inter_dim=512): super(ASFF_4, self).__init__()
self.inter_dim = inter_dim compress_c = 8
self.weight_level_0 = BasicConv(self.inter_dim, compress_c, 1, 1) self.weight_level_1 = BasicConv(self.inter_dim, compress_c, 1, 1) self.weight_level_2 = BasicConv(self.inter_dim, compress_c, 1, 1) self.weight_level_3 = BasicConv(self.inter_dim, compress_c, 1, 1)
self.weight_levels = nn.Conv2d(compress_c * 4, 4, kernel_size=1, stride=1, padding=0)
self.conv = BasicConv(self.inter_dim, self.inter_dim, 3, 1)
def forward(self, input0, input1, input2, input3): level_0_weight_v = self.weight_level_0(input0) level_1_weight_v = self.weight_level_1(input1) level_2_weight_v = self.weight_level_2(input2) level_3_weight_v = self.weight_level_3(input3)
levels_weight_v = torch.cat((level_0_weight_v, level_1_weight_v, level_2_weight_v, level_3_weight_v), 1) levels_weight = self.weight_levels(levels_weight_v) levels_weight = F.softmax(levels_weight, dim=1)
fused_out_reduced = input0 * levels_weight[:, 0:1, :, :] + \ input1 * levels_weight[:, 1:2, :, :] + \ input2 * levels_weight[:, 2:3, :, :] + \ input3 * levels_weight[:, 3:, :, :]
out = self.conv(fused_out_reduced)
return out
class BlockBody(nn.Module): def __init__(self, channels=[64, 128, 256, 512]): super(BlockBody, self).__init__()
self.blocks_scalezero1 = nn.Sequential( BasicConv(channels[0], channels[0], 1), ) self.blocks_scaleone1 = nn.Sequential( BasicConv(channels[1], channels[1], 1), ) self.blocks_scaletwo1 = nn.Sequential( BasicConv(channels[2], channels[2], 1), ) self.blocks_scalethree1 = nn.Sequential( BasicConv(channels[3], channels[3], 1), )
self.downsample_scalezero1_2 = Downsample_x2(channels[0], channels[1]) self.upsample_scaleone1_2 = Upsample(channels[1], channels[0], scale_factor=2)
self.asff_scalezero1 = ASFF_2(inter_dim=channels[0]) self.asff_scaleone1 = ASFF_2(inter_dim=channels[1])
self.blocks_scalezero2 = nn.Sequential( BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), ) self.blocks_scaleone2 = nn.Sequential( BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), )
self.downsample_scalezero2_2 = Downsample_x2(channels[0], channels[1]) self.downsample_scalezero2_4 = Downsample_x4(channels[0], channels[2]) self.downsample_scaleone2_2 = Downsample_x2(channels[1], channels[2]) self.upsample_scaleone2_2 = Upsample(channels[1], channels[0], scale_factor=2) self.upsample_scaletwo2_2 = Upsample(channels[2], channels[1], scale_factor=2) self.upsample_scaletwo2_4 = Upsample(channels[2], channels[0], scale_factor=4)
self.asff_scalezero2 = ASFF_3(inter_dim=channels[0]) self.asff_scaleone2 = ASFF_3(inter_dim=channels[1]) self.asff_scaletwo2 = ASFF_3(inter_dim=channels[2])
self.blocks_scalezero3 = nn.Sequential( BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), ) self.blocks_scaleone3 = nn.Sequential( BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), ) self.blocks_scaletwo3 = nn.Sequential( BasicBlock(channels[2], channels[2]), BasicBlock(channels[2], channels[2]), BasicBlock(channels[2], channels[2]), BasicBlock(channels[2], channels[2]), )
self.downsample_scalezero3_2 = Downsample_x2(channels[0], channels[1]) self.downsample_scalezero3_4 = Downsample_x4(channels[0], channels[2]) self.downsample_scalezero3_8 = Downsample_x8(channels[0], channels[3]) self.upsample_scaleone3_2 = Upsample(channels[1], channels[0], scale_factor=2) self.downsample_scaleone3_2 = Downsample_x2(channels[1], channels[2]) self.downsample_scaleone3_4 = Downsample_x4(channels[1], channels[3]) self.upsample_scaletwo3_4 = Upsample(channels[2], channels[0], scale_factor=4) self.upsample_scaletwo3_2 = Upsample(channels[2], channels[1], scale_factor=2) self.downsample_scaletwo3_2 = Downsample_x2(channels[2], channels[3]) self.upsample_scalethree3_8 = Upsample(channels[3], channels[0], scale_factor=8) self.upsample_scalethree3_4 = Upsample(channels[3], channels[1], scale_factor=4) self.upsample_scalethree3_2 = Upsample(channels[3], channels[2], scale_factor=2)
self.asff_scalezero3 = ASFF_4(inter_dim=channels[0]) self.asff_scaleone3 = ASFF_4(inter_dim=channels[1]) self.asff_scaletwo3 = ASFF_4(inter_dim=channels[2]) self.asff_scalethree3 = ASFF_4(inter_dim=channels[3])
self.blocks_scalezero4 = nn.Sequential( BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), BasicBlock(channels[0], channels[0]), ) self.blocks_scaleone4 = nn.Sequential( BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), BasicBlock(channels[1], channels[1]), ) self.blocks_scaletwo4 = nn.Sequential( BasicBlock(channels[2], channels[2]), BasicBlock(channels[2], channels[2]), BasicBlock(channels[2], channels[2]), BasicBlock(channels[2], channels[2]), ) self.blocks_scalethree4 = nn.Sequential( BasicBlock(channels[3], channels[3]), BasicBlock(channels[3], channels[3]), BasicBlock(channels[3], channels[3]), BasicBlock(channels[3], channels[3]), )
def forward(self, x): x0, x1, x2, x3 = x
x0 = self.blocks_scalezero1(x0) x1 = self.blocks_scaleone1(x1) x2 = self.blocks_scaletwo1(x2) x3 = self.blocks_scalethree1(x3)
scalezero = self.asff_scalezero1(x0, self.upsample_scaleone1_2(x1)) scaleone = self.asff_scaleone1(self.downsample_scalezero1_2(x0), x1)
x0 = self.blocks_scalezero2(scalezero) x1 = self.blocks_scaleone2(scaleone)
scalezero = self.asff_scalezero2(x0, self.upsample_scaleone2_2(x1), self.upsample_scaletwo2_4(x2)) scaleone = self.asff_scaleone2(self.downsample_scalezero2_2(x0), x1, self.upsample_scaletwo2_2(x2)) scaletwo = self.asff_scaletwo2(self.downsample_scalezero2_4(x0), self.downsample_scaleone2_2(x1), x2)
x0 = self.blocks_scalezero3(scalezero) x1 = self.blocks_scaleone3(scaleone) x2 = self.blocks_scaletwo3(scaletwo)
scalezero = self.asff_scalezero3(x0, self.upsample_scaleone3_2(x1), self.upsample_scaletwo3_4(x2), self.upsample_scalethree3_8(x3)) scaleone = self.asff_scaleone3(self.downsample_scalezero3_2(x0), x1, self.upsample_scaletwo3_2(x2), self.upsample_scalethree3_4(x3)) scaletwo = self.asff_scaletwo3(self.downsample_scalezero3_4(x0), self.downsample_scaleone3_2(x1), x2, self.upsample_scalethree3_2(x3)) scalethree = self.asff_scalethree3(self.downsample_scalezero3_8(x0), self.downsample_scaleone3_4(x1), self.downsample_scaletwo3_2(x2), x3)
scalezero = self.blocks_scalezero4(scalezero) scaleone = self.blocks_scaleone4(scaleone) scaletwo = self.blocks_scaletwo4(scaletwo) scalethree = self.blocks_scalethree4(scalethree)
return scalezero, scaleone, scaletwo, scalethree
class AFPN(nn.Module): def __init__(self, in_channels=[256, 512, 1024, 2048], out_channels=128): super(AFPN, self).__init__()
self.fp16_enabled = False
self.conv0 = BasicConv(in_channels[0], in_channels[0] // 8, 1) self.conv1 = BasicConv(in_channels[1], in_channels[1] // 8, 1) self.conv2 = BasicConv(in_channels[2], in_channels[2] // 8, 1) self.conv3 = BasicConv(in_channels[3], in_channels[3] // 8, 1)
self.body = nn.Sequential( BlockBody([in_channels[0] // 8, in_channels[1] // 8, in_channels[2] // 8, in_channels[3] // 8]) )
self.conv00 = BasicConv(in_channels[0] // 8, out_channels, 1) self.conv11 = BasicConv(in_channels[1] // 8, out_channels, 1) self.conv22 = BasicConv(in_channels[2] // 8, out_channels, 1) self.conv33 = BasicConv(in_channels[3] // 8, out_channels, 1) self.conv44 = nn.MaxPool2d(kernel_size=1, stride=2)
# init weight for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.xavier_normal_(m.weight, gain=0.02) elif isinstance(m, nn.BatchNorm2d): torch.nn.init.normal_(m.weight.data, 1.0, 0.02) torch.nn.init.constant_(m.bias.data, 0.0)
def forward(self, x): x0, x1, x2, x3 = x
x0 = self.conv0(x0) x1 = self.conv1(x1) x2 = self.conv2(x2) x3 = self.conv3(x3)
out0, out1, out2, out3 = self.body([x0, x1, x2, x3])
out0 = self.conv00(out0) out1 = self.conv11(out1) out2 = self.conv22(out2) out3 = self.conv33(out3)
return out0, out1, out2, out3
class Detect_AFPN(nn.Module): """YOLOv8 Detect head for detection models.""" dynamic = False # force grid reconstruction export = False # export mode shape = None anchors = torch.empty(0) # init strides = torch.empty(0) # init
def __init__(self, nc=80, channel=128, ch=()): """Initializes the YOLOv8 detection layer with specified number of classes and channels.""" super().__init__() self.nc = nc # number of classes self.nl = len(ch) # number of detection layers self.reg_max = 16 # DFL channels (ch[0] // 16 to scale 4/8/12/16/20 for n/s/m/l/x) self.no = nc + self.reg_max * 4 # number of outputs per anchor self.stride = torch.zeros(self.nl) # strides computed during build c2, c3 = max((16, ch[0] // 4, self.reg_max * 4)), max(ch[0], min(self.nc, 100)) # channels self.cv2 = nn.ModuleList( nn.Sequential(Conv(channel, c2, 3), Conv(c2, c2, 3), nn.Conv2d(c2, 4 * self.reg_max, 1)) for x in ch) self.cv3 = nn.ModuleList(nn.Sequential(Conv(channel, c3, 3), Conv(c3, c3, 3), nn.Conv2d(c3, self.nc, 1)) for x in ch) self.dfl = DFL(self.reg_max) if self.reg_max > 1 else nn.Identity() self.AFPN = AFPN(ch)
def forward(self, x): """Concatenates and returns predicted bounding boxes and class probabilities.""" x = list(self.AFPN(x)) shape = x[0].shape # BCHW for i in range(self.nl): x[i] = torch.cat((self.cv2[i](x[i]), self.cv3[i](x[i])), 1) if self.training: return x elif self.dynamic or self.shape != shape: self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5)) self.shape = shape
x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2) if self.export and self.format in ('saved_model', 'pb', 'tflite', 'edgetpu', 'tfjs'): # avoid TF FlexSplitV ops box = x_cat[:, :self.reg_max * 4] cls = x_cat[:, self.reg_max * 4:] else: box, cls = x_cat.split((self.reg_max * 4, self.nc), 1) dbox = dist2bbox(self.dfl(box), self.anchors.unsqueeze(0), xywh=True, dim=1) * self.strides
if self.export and self.format in ('tflite', 'edgetpu'): # Normalize xywh with image size to mitigate quantization error of TFLite integer models as done in YOLOv5: # https://github.com/ultralytics/yolov5/blob/0c8de3fca4a702f8ff5c435e67f378d1fce70243/models/tf.py#L307-L309 # See this PR for details: https://github.com/ultralytics/ultralytics/pull/1695 img_h = shape[2] * self.stride[0] img_w = shape[3] * self.stride[0] img_size = torch.tensor([img_w, img_h, img_w, img_h], device=dbox.device).reshape(1, 4, 1) dbox /= img_size
y = torch.cat((dbox, cls.sigmoid()), 1) return y if self.export else (y, x)
def bias_init(self): """Initialize Detect() biases, WARNING: requires stride availability.""" m = self # self.model[-1] # Detect() module # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1 # ncf = math.log(0.6 / (m.nc - 0.999999)) if cf is None else torch.log(cf / cf.sum()) # nominal class frequency for a, b, s in zip(m.cv2, m.cv3, m.stride): # from a[-1].bias.data[:] = 1.0 # box b[-1].bias.data[:m.nc] = math.log(5 / m.nc / (640 / s) ** 2) # cls (.01 objects, 80 classes, 640 img) |
1）将fxafpn.py文件放到ultralytics/nn/modules文件夹中；
2）修改ultralytics/nn/modules/__init__.py，导入Detect_AFPN模块，即增加92行和100行的内容；
3）修改ultralytics/nn/tasks.py，有5处需要添加Detect_AFPN：
最后1处需要添加 else：return “detect”；
4）以yolo11.yaml文件为基础重新写一个yolo11-fxafpn.yaml文件，位置如下：
ultralytics/cfg/models/11/yolo11-fxafpn.yaml
| # Parameters nc: 80 # number of classes scales: # model compound scaling constants, i.e. 'model=yolo11n.yaml' will call yolo11.yaml with scale 'n' # [depth, width, max_channels] n: [0.50, 0.25, 1024] # summary: 319 layers, 2624080 parameters, 2624064 gradients, 6.6 GFLOPs s: [0.50, 0.50, 1024] # summary: 319 layers, 9458752 parameters, 9458736 gradients, 21.7 GFLOPs m: [0.50, 1.00, 512] # summary: 409 layers, 20114688 parameters, 20114672 gradients, 68.5 GFLOPs l: [1.00, 1.00, 512] # summary: 631 layers, 25372160 parameters, 25372144 gradients, 87.6 GFLOPs x: [1.00, 1.50, 512] # summary: 631 layers, 56966176 parameters, 56966160 gradients, 196.0 GFLOPs # YOLO11n-fxafpn 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]] # 9 - [-1, 2, C2PSA, [1024]] # 10 # YOLO11n-fxafpn head head: - [[2, 4, 6, 10], 1, Detect_AFPN, [nc, 128]] # Detect(P2, P3, P4, P5) |
5）其他步骤可以参照前面的“修改1：Conv替换LAE”。
扫一扫
7271
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
