基于pytorch实现的yolo v3

一、introduction

论文地址:https://arxiv.org/abs/1804.02767
项目代码:https://download.csdn.net/download/qq_40356092/13122771
本文是对yolo v3的复现,在前人的基础上对代码进行了修改,尽可能方便读者进行训练和测试,同时易于修改。

二、网络结构

yolo v3使用了Darknet53作为主干网络(backbone),Darknet53的结构来源于Resnet,同样使用了resnet中的残差网络结构和shortcut连接,不同的是全部使用3×3和1×1的卷积层(已经被证明,用多个小的卷积层替代一个大的卷积层效果更好),将池化层全部用大小为3×3、步长为2的卷积层代替(因为池化层会丢弃掉一部分特征信息)。
基于pytorch实现的yolo v3_第1张图片
yolo v3使用了三个检测头来检测物体,下采样系数分别为8、16、32,对应不同大小的特征图(52×52)、(26×26)、(13×13),分别用来检测小、中、大目标。

文件夹nets下的darknet.py是主干网络模型,yolo v3.py是yolo v3完整的模型,可以通过运行主函数来查看网络结构。

三、测试图片

1、测试图片代码predict.py
需要测试的图片存放在文件夹img下,直接运行predict.py即可对图片进行测试。测试代码通过调用yolo.py文件中的YOLO()类来实现图片的检测:
 初始化对象yolo = YOLO()(YOLO类来自于yolo.py文件)
 调用yolo.detect_image(image)方法获取检测之后的图片并显示出来

2、参数设置:

    _defaults = {
        "model_path": 'model_data/yolo_weights.pth',
        #"model_path": 'checkpoints/Epoch32-Total_Loss6.9099-Val_Loss10.9076.pth',
        "classes_path": 'model_data/voc_classes.txt',
        "model_image_size" : (416, 416, 3),
        "confidence": 0.5,
        "iou" : 0.3,
        "cuda": True
    }

model_path是模型权重保存的位置
classes_path是类别文件,一共有80个类。在训练你自己的数据集的时候需要根据实际的类别数进行修改。

3、检测图片的流程def detect_image(self, image):
 图片预处理过程:获取图片的宽和高,将图片resize到(416,416)并加上灰条,归一化操作,转换成tensor存入cuda
 将图片传入网络,得到三个维度的输出,将其解码并用torch.cat连接到一起,使用NMS算法去除类别置信度过低的框和IOU重叠过大的框,返回图片存入到cpu中
 后续处理:将NMS返回的框的置信度和种类的置信度相乘作为最终的置信度,获取每个框的四个坐标(在最终生成的图片中,每个框上边的类别和类别的置信度来源于此)。去掉图片的灰条,调整回原来的图片尺寸大小。定义框上边字体的格式和框的样式,最后利用for循环依次画框。

4、解码class DecodeBox(nn.Module):
这里t是网络的直接输出值,我们需要通过一系列变换将其解码成b(中心坐标是相对于图像的,宽高是绝对值)(在NMS算法中会将中心坐标+宽高的形式再转换成左上角和右下角两个点的形式)
变换公式如下图所示:
基于pytorch实现的yolo v3_第2张图片
5、NMS算法def non_max_suppression(prediction, num_classes, conf_thres=0.5, nms_thres=0.4):

:param prediction:输出的预测结果
:param num_classes: 类别数
:param conf_thres: 置信度阈值
:param nms_thres: NMS阈值,重叠程度
:return:[output[image_i], max_detections],二维变量,图片的id和框(x1, y1, x2, y2,obj_conf, class_conf)

(1)prediction的格式是中心和宽高(cx,cy,w,h),先将其转换成左上角和右下角格式(x1,y1,x2,y2)
(2)遍历所有图片,每张图片遍历所有的网格点grid cell,对于每个网格点,选择num_classes个类别(这里一共有20个类别)中最大的值作为当前网格点负责目标的类别。
(3)利用每个grid cell中框的置信度(obj_conf参数)来进行第一轮筛选,去掉小于conf_thres阈值的grid cell,将结果重新连接成(x1, y1, x2, y2, obj_conf, class_conf, class_pred)的形式
(4)获取当前图片中预测出来的所有类别,对于每一个类别,用官方自带的NMS算法来去除重叠程度过大的框。
(5)官方NMS算法流程:对于此类别下所有的框,先找到置信度最大的框A,然后其余所有框和A计算IOU,如果IOU大于阈值,则看做其余框和A框到了同一个物体,去掉这些IOU过大的框,只保留框A。然后在剩余的框中找到置信度最大的框,重复以上过程直到没有框剩下。

            keep = nms(
                detections_class[:, :4],
                detections_class[:, 4]*detections_class[:, 5],
                nms_thres
            )
            max_detections = detections_class[keep]

总结:测试图片的流程其实非常简单,可以分为以下六步。
(1)输入图片
(2)调整图片大小为416×416,归一化操作,转换成tensor格式。(预处理操作)
(3)feed进网络,得到输出,将输出解码并连接
(4)使用NMS算法去重
(5)将图片和预测框转换成原始格式
(6)在图片上画框并显示

四、训练模型

训练模型的代码在train.py文件中
1、在训练模型之前我们首先要对数据集进行处理,生成一个txt文件用来存储图片和标注框的信息。(make_dataset.py)

2、训练
(1)加载预训练模型和权重,权重文件保存路径在model_data/yolo_weights.pth。设置net = model.train(),使用GPU进行训练
(2)建立loss函数,其中loss函数的定义在yolo_training.py文件的类函数class YOLOLoss(nn.Module)中。(关于loss函数的详细讲解后续会有更新)
(3)按照9:1的比例分割训练集和验证集
(4)定义超参数,优化器和学习率下降方式
(5)自定义YoloDataset类函数来完成预处理(包括多种数据增强方法),YoloDataset类继承了Dataset类。使用DataLoader返回数据集的迭代,批量feed进网络。

训练的一些细节:前50个epoch冻结了主干网络的参数,后50个epoch对这些参数进行解冻训练。在YoloDataset类函数中对数据进行了随机打乱,而DataLoader中没有设置shuffle(默认为False),所以每一代的训练集和验证集都是相同的。

3、调用函数def fit_ont_epoch(net,yolo_losses,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda)进行每一个epoch进行训练。
模型保存在文件夹checkpoints里面。

五、writing in the end

此博客是对我上传的yolo v3代码的一个解释。如果你下载了我的代码或者对以上内容有任何不理解的地方,都可以在下方进行留言和评论。

下面是一些对你理解yolo v3很有帮助的文章:
[1]:https://blog.csdn.net/weixin_44791964/article/details/105310627
[2]:https://zhuanlan.zhihu.com/p/37201615

你可能感兴趣的:(目标检测,pytorch,深度学习)