替换YOLOv5主干网络

首先写一个 demo.py 文件来查看 torchvision下面都有哪些模型

替换YOLOv5主干网络_第1张图片

如图所示,可以看到 torchvision.models 下面有很多模型, 比如 vgg16,resnet50,AlexNet,resnet18等等,在这里我选择 vgg16 模型 来替换YOLOv5的主干网络

下面先来看一下 vgg16 模型的结构:

替换YOLOv5主干网络_第2张图片

替换YOLOv5主干网络_第3张图片

替换YOLOv5主干网络_第4张图片 可以看到 vgg16 模型主要由 feature (特征提取层:做特征提取) 、avgpool (池化层) 以及 classifier (分类层) 组成,要替换YOLOv5的主干网络,其实就是要将YOLOv5的特征提取层用vgg16的特征提取层进行替换

下面给该网络模型一个输入数据,来看一下该网络模型每一层的输出大小是多少,才能去进行主干网络的替换操作

打开 yolov5s.yaml 文件,如图所示,输出要有 80x80,40x40

替换YOLOv5主干网络_第5张图片

你可能感兴趣的:(YOLO)