【YOLO改进】换遍MMPretrain TIMM主干网络之EfficientNet-B1（基于MMYOLO）

最新推荐文章于 2024-07-01 15:36:24 发布

五山一胖

最新推荐文章于 2024-07-01 15:36:24 发布

阅读量516

点赞数 8

文章标签： YOLO

本文链接：https://blog.csdn.net/vlone_pp/article/details/139006906

版权

TIMM库

TIMM (PyTorch Image Models) 是一个广泛使用的深度学习库，专注于图像模型的实现和优化。

TIMM库的优点

丰富的模型库： TIMM库包含了大量的预训练图像模型，包括ResNet、EfficientNet、ViT、ConvNeXt等。这些模型在各自的任务上都经过了广泛的验证，具有高性能和可靠性。
高可配置性： TIMM库提供了丰富的配置选项，用户可以根据具体需求自定义模型的各个参数，如输入尺寸、卷积核大小、层数等。这种灵活性使得用户能够针对特定任务进行优化。
集成预训练权重： TIMM库中的模型大多提供了预训练权重，能够在ImageNet等大型数据集上进行初始化。这不仅节省了训练时间，还能提升模型的初始性能。
优化和性能调优： TIMM库对每个模型都进行了精心的优化，包括高效的实现和内置的性能调优工具。这使得模型在实际应用中能够达到较高的推理速度和效率。
活跃的社区和持续更新： TIMM库由活跃的开发者社区维护，持续更新和优化新模型，确保其与最新的研究进展和技术保持同步。

作为YOLO系列网络主干网络的可行性

YOLO（You Only Look Once）系列网络对主干网络的要求是既要有高效的特征提取能力，又要在推理速度上具有较高的实时性。TIMM库中的模型具备这些特点，因此可以作为YOLO系列网络的主干网络。具体可行性分析如下：

多样化选择： TIMM库提供了多种高性能的卷积神经网络和Transformer架构，可以根据实际需求选择最合适的主干网络。例如，可以选择高效的EfficientNet来优化推理速度，或者选择性能强大的ResNet或ConvNeXt来提高检测精度。
预训练模型：使用TIMM库中的预训练模型作为YOLO的主干网络，可以显著加速训练过程，减少所需计算资源，并提高模型初始性能。这对于需要快速部署的应用场景尤为重要。
性能优化： TIMM库中的模型经过了精心的优化，具有较高的计算效率和较低的推理延迟。这与YOLO网络的实时性需求高度契合，可以在保持高检测精度的同时，保证足够快的推理速度。
易于集成： TIMM库的模型设计兼容性好，可以较为容易地集成到YOLO架构中。用户可以根据需求调整模型的参数和配置，以适应YOLO网络的输入和输出需求。

替换TIMM 库中的主干网络以 `EfficientNet-B1为例子`

由于 MMPretrain提供了 TIMM 主干网络的封装，用户也可以通过 MMClassification 直接使用 timm 中的主干网络。假设想将 EfficientNet-B1作为 YOLOv5 的主干网络，则配置文件如下：

_base_ = './yolov5_s-v61_syncbn_8xb16-300e_coco.py'


# pip install timm，安装 timm
# 导入 mmpretrain.models 使得可以调用 mmcls 中注册的模块
custom_imports = dict(imports=['mmpretrain.models'], allow_failed_imports=False)

deepen_factor = _base_.deepen_factor
widen_factor = 1.0
channels = [40, 112, 320]

model = dict(
    backbone=dict(
        _delete_=True,  # 将 _base_ 中关于 backbone 的字段删除
        type='mmpretrain.TIMMBackbone',  # 使用 mmcls 中的 timm 主干网络
        model_name='efficientnet_b1',  # 使用 TIMM 中的 efficientnet_b1
        features_only=True,
        pretrained=True,
        out_indices=(2, 3, 4)),
    neck=dict(
        type='YOLOv5PAFPN',
        deepen_factor=deepen_factor,
        widen_factor=widen_factor,
        in_channels=channels,  # 注意：EfficientNet-B1 输出的3个通道是 [40, 112, 320]，和原先的 yolov5-s neck 不匹配，需要更改
        out_channels=channels),
    bbox_head=dict(
        type='YOLOv5Head',
        head_module=dict(
            type='YOLOv5HeadModule',
            in_channels=channels,  # head 部分输入通道也要做相应更改
            widen_factor=widen_factor))
)