首页 > 其他分享 >YOLOv10改进教程|C2f-CIB加入注意力机制

YOLOv10改进教程|C2f-CIB加入注意力机制

时间:2024-07-01 22:00:41浏览次数:3  
标签:__ CIB nn C2fCIBAttention self C2f YOLOv10 True


  一、 导读

        论文链接:https://arxiv.org/abs/2311.11587

        代码链接:GitHub - CV-ZhangXin/AKConv

 YOLOv10训练、验证及推理教程


二、 C2f-CIB加入注意力机制

2.1 复制代码

        打开ultralytics->nn->modules->block.py文件,复制SE注意力机制(也可以自行换成别的)代码,并创建C2fCIBAttention代码,如下图所示:

class SE(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channel, channel // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channel // reduction, channel, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)


class C2fCIBAttention(nn.Module):
    """Faster Implementation of CSP Bottleneck with 2 convolutions."""

    def __init__(self, c1, c2, n=1, shortcut=False, lk=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 channels
        self.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(CIB(self.c, self.c, shortcut, e=1.0, lk=lk) for _ in range(n))
        self.atten = SE(C2)

    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.atten(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))

        并在上方声明C2fCIBAttention类。

        在nn.models.__init__.py中声明 C2fCIBAttention。

2.2 修改tasks.py 

       打开ultralytics->nn->tasks.py,如图所示操作。

​2.3 修改yolov10n.yaml

        将yolov10n.yaml文件中的C2fCIB替换为C2fCIBAttention。

# Ultralytics YOLO 

标签:__,CIB,nn,C2fCIBAttention,self,C2f,YOLOv10,True
From: https://blog.csdn.net/StopAndGoyyy/article/details/140110212

相关文章