首先写一个 demo.py 文件来查看 torchvision下面都有哪些模型
如图所示,可以看到 torchvision.models 下面有很多模型, 比如 vgg16,resnet50,AlexNet,resnet18等等,在这里我选择 vgg16 模型 来替换YOLOv5的主干网络
下面先来看一下 vgg16 模型的结构:
可以看到 vgg16 模型主要由 feature (特征提取层:做特征提取) 、avgpool (池化层) 以及 classifier (分类层) 组成,要替换YOLOv5的主干网络,其实就是要将YOLOv5的特征提取层用vgg16的特征提取层进行替换
下面给该网络模型一个输入数据,来看一下该网络模型每一层的输出大小是多少,才能去进行主干网络的替换操作
打开 yolov5s.yaml 文件,如图所示,输出要有 80x80,40x40