基于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的卷积层代替（因为池化层会丢弃掉一部分特征信息）。

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算法中会将中心坐标+宽高的形式再转换成左上角和右下角两个点的形式）
变换公式如下图所示：

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里面。

淘先锋技术网

基于pytorch实现的yolo v3

一、introduction

二、网络结构

三、测试图片

四、训练模型

五、writing in the end