YOLOv8改进 | 主干篇 | 添加FasterNeT跑起来的主干网络（提高FPS和检测效率）

鱼弦

于 2024-09-29 13:30:00 发布

阅读量1k

点赞数 16

分类专栏：人工智能时代文章标签： YOLO

本文链接：https://blog.csdn.net/feng1790291543/article/details/139728433

版权

人工智能时代专栏收录该内容

2 篇文章 1 订阅

订阅专栏

鱼弦：公众号【红尘灯塔】，CSDN博客专家、内容合伙人、新星导师、全栈领域优质创作者、51CTO(Top红人+专家博主) 、github开源爱好者（go-zero源码二次开发、游戏后端架构 https://github.com/Peakchen）

YOLOv8改进 | 主干篇 | 添加FasterNeT跑起来的主干网络（提高FPS和检测效率）

1. 简介

将FasterNeT作为YOLOv8的主干网络可以提高模型的FPS和检测效率。FasterNeT是一种轻量级的卷积神经网络架构，它通过高效的特征提取模块和轻量化的注意力机制来实现高性能和低延迟。

2. 原理详解

FasterNeT主要包含以下几个部分：

高效的特征提取模块: FasterNeT采用了高效的特征提取模块，例如Inverted Residual Block和ShuffleNet Unit，可以提取有效的特征信息，同时降低计算复杂度。
轻量化的注意力机制: FasterNeT采用了轻量化的注意力机制，例如Squeeze-and-Excitation (SE) module和Mobile Attention Block，可以捕获更丰富的上下文信息，同时保持模型的轻量化。

3. 应用场景解释

将FasterNeT作为YOLOv8的主干网络适用于以下场景：

移动端目标检测: 在移动设备上进行实时目标检测，例如智能手机、无人机等。
嵌入式系统目标检测: 在嵌入式系统中进行目标检测，例如智能家居、自动驾驶等。
实时目标检测: 在需要实时性要求的场景中进行目标检测，例如安防监控、视频分析等。

实现：

下面是主要的实现细节:

FasterNeT 的核心思想:
- 针对传统 CNN 网络计算量大的问题,提出了一种高效的网络结构 FasterNeT
- 通过引入轻量级的注意力机制和高效的逐通道卷积,大幅降低了计算开销
- 同时保持了良好的特征提取能力,适用于实时目标检测等场景
整合 FasterNeT 到 YOLOv8 主干网络:
- 用 FasterNeT 替换 YOLOv8 原有的 CSPDarknet 主干网络
- 保持 YOLOv8 检测头的结构不变,只修改主干网络部分
- 针对 YOLOv8 的输入尺度和通道数对 FasterNeT 进行适当调整
主要代码实现:

# faster_net.py
import torch.nn as nn

class FasterNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1):
        super(FasterNetBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.ca = ChannelAttention(out_channels)

    def forward(self, x):
        out = self.conv1(x)
        out = self.bn1(out)
        out = nn.ReLU(inplace=True)(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.ca(out)
        return out

class FasterNet(nn.Module):
    def __init__(self, num_blocks, num_classes=80):
        super(FasterNet, self).__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 64, 3, 2, 1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.layers = self._make_layers(64, num_blocks)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(256, num_classes)

    def _make_layers(self, in_channels, num_blocks):
        layers = []
        for i in range(num_blocks):
            if i == 0:
                layers.append(FasterNetBlock(in_channels, 128, 2))
            else:
                layers.append(FasterNetBlock(128, 128))
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.stem(x)
        out = self.layers(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

# yolov8_with_fasternet.py
from yolov8.models.backbone import BaseBackbone
from faster_net import FasterNet

class YOLOv8WithFasterNet(BaseBackbone):
    def __init__(self, num_blocks=3, num_classes=80):
        super(YOLOv8WithFasterNet, self).__init__()
        self.faster_net = FasterNet(num_blocks, num_classes)

    def forward(self, x):
        return self.faster_net(x)

使用示例:

# 创建 YOLOv8 with FasterNet 主干网络
backbone = YOLOv8WithFasterNet(num_blocks=3, num_classes=80)

# 创建 YOLOv8 检测头
head = YOLOv8DetectionHead(backbone, num_classes=80, num_anchors=3)

# 前向推理
input_tensor = torch.randn(1, 3, 640, 640)
output = head(input_tensor)

主要步骤: