赞
踩
随着深度学习技术的不断发展,目标检测作为计算机视觉领域的重要分支,已经在实际应用中取得了显著的成果。YOLO(You Only Look Once)系列模型作为目标检测领域的经典之作,以其高效、快速的特点,赢得了广大研究者和开发者的青睐。YOLOv4更是在前代模型的基础上,通过一系列改进和创新,进一步提升了模型的性能,使得其在处理复杂场景和多样化目标时,表现出色。
本次实验旨在通过实际训练和测试YOLOv4模型,深入了解其内部机制、优化策略以及在实际应用中的表现。在实验中,将重点关注模型的训练过程、参数调整、性能评估等方面。我们将使用自己的标注数据集,通过合理的数据划分和预处理,为模型的训练提供优质的数据支持。最后,我们将对训练好的模型进行测试和评估,验证其在目标检测任务中的准确性和有效性,并展示检测结果。
下载labellmg
整理好需要标注的数据集,放在一个文件夹JPEJImages
进入到labellmg
保存好的标注数据集Annotations文件夹里面的数据都是xml类型。
处理VOC格式数据集的脚本,目的是从数据集的XML标注文件中提取出目标的位置和类别信息,并统计每个类别的目标数量。
先创建一个包含数据集标注标签的class.txt文件。
在代码中导入依赖库、设置参数、获取类别列表、初始化统计变量、函数定义convert_annotation、主程序if_name_=="_main_"
以下是代码voc_annotation部分
- import os
- import random
- import xml.etree.ElementTree as ET
- import numpy as np
- from utils.utils import get_classes
-
- #--------------------------------------------------------------------------------------------------------------------------------#
- # annotation_mode用于指定该文件运行时计算的内容
- # annotation_mode为0代表整个标签处理过程,包括获得VOCdevkit/VOC2008/ImageSets里面的txt以及训练用的2008_train.txt、2008_val.txt
- # annotation_mode为1代表获得VOCdevkit/VOC2008/ImageSets里面的txt
- # annotation_mode为2代表获得训练用的2008_train.txt、2008_val.txt
- #--------------------------------------------------------------------------------------------------------------------------------#
- annotation_mode = 0
- #-------------------------------------------------------------------#
- # 必须要修改,用于生成2008_train.txt、2008_val.txt的目标信息
- # 与训练和预测所用的classes_path一致即可
- # 如果生成的2008_train.txt里面没有目标信息
- # 那么就是因为classes没有设定正确
- # 仅在annotation_mode为0和2的时候有效
- #-------------------------------------------------------------------#
- classes_path = 'model_data/class.txt'
- #--------------------------------------------------------------------------------------------------------------------------------#
- # trainval_percent用于指定(训练集+验证集)与测试集的比例,默认情况下 (训练集+验证集):测试集 = 9:1
- # train_percent用于指定(训练集+验证集)中训练集与验证集的比例,默认情况下 训练集:验证集 = 9:1
- # 仅在annotation_mode为0和1的时候有效
- #--------------------------------------------------------------------------------------------------------------------------------#
- trainval_percent = 0.9
- train_percent = 0.9
- #-------------------------------------------------------#
- # 指向VOC数据集所在的文件夹
- # 默认指向根目录下的VOC数据集
- #-------------------------------------------------------#
- VOCdevkit_path = 'VOCdevkit'
-
- VOCdevkit_sets = [('2008', 'train'), ('2008', 'val')]
- classes, _ = get_classes(classes_path)
-
- #-------------------------------------------------------#
- # 统计目标数量
- #-------------------------------------------------------#
- photo_nums = np.zeros(len(VOCdevkit_sets))
- nums = np.zeros(len(classes))
- def convert_annotation(year, image_id, list_file):
- in_file = open(os.path.join(VOCdevkit_path, 'VOC%s/Annotations/%s.xml'%(year, image_id)), encoding='utf-8')
- tree=ET.parse(in_file)
- root = tree.getroot()
-
- for obj in root.iter('object'):
- difficult = 0
- if obj.find('difficult')!=None:
- difficult = obj.find('difficult').text
- cls = obj.find('name').text
- if cls not in classes or int(difficult)==1:
- continue
- cls_id = classes.index(cls)
- xmlbox = obj.find('bndbox')
- b = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
- list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
-
- nums[classes.index(cls)] = nums[classes.index(cls)] + 1
-
- if __name__ == "__main__":
- random.seed(0)
- if " " in os.path.abspath(VOCdevkit_path):
- raise ValueError("数据集存放的文件夹路径与图片名称中不可以存在空格,否则会影响正常的模型训练,请注意修改。")
-
- if annotation_mode == 0 or annotation_mode == 1:
- print("Generate txt in ImageSets.")
- xmlfilepath = os.path.join(VOCdevkit_path, 'VOC2008/Annotations')
- saveBasePath = os.path.join(VOCdevkit_path, 'VOC2008/ImageSets/Main')
- temp_xml = os.listdir(xmlfilepath)
- total_xml = []
- for xml in temp_xml:
- if xml.endswith(".xml"):
- total_xml.append(xml)
-
- num = len(total_xml)
- list = range(num)
- tv = int(num*trainval_percent)
- tr = int(tv*train_percent)
- trainval= random.sample(list,tv)
- train = random.sample(trainval,tr)
-
- print("train and val size",tv)
- print("train size",tr)
- ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
- ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
- ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
- fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
-
- for i in list:
- name=total_xml[i][:-4]+'\n'
- if i in trainval:
- ftrainval.write(name)
- if i in train:
- ftrain.write(name)
- else:
- fval.write(name)
- else:
- ftest.write(name)
-
- ftrainval.close()
- ftrain.close()
- fval.close()
- ftest.close()
- print("Generate txt in ImageSets done.")
-
- if annotation_mode == 0 or annotation_mode == 2:
- print("Generate 2008_train.txt and 2008_val.txt for train.")
- type_index = 0
- for year, image_set in VOCdevkit_sets:
- image_ids = open(os.path.join(VOCdevkit_path, 'VOC%s/ImageSets/Main/%s.txt'%(year, image_set)), encoding='utf-8').read().strip().split()
- list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
- for image_id in image_ids:
- list_file.write('%s/VOC%s/JPEGImages/%s.jpg'%(os.path.abspath(VOCdevkit_path), year, image_id))
-
- convert_annotation(year, image_id, list_file)
- list_file.write('\n')
- photo_nums[type_index] = len(image_ids)
- type_index += 1
- list_file.close()
- print("Generate 2008_train.txt and 2008_val.txt for train done.")
-
- def printTable(List1, List2):
- for i in range(len(List1[0])):
- print("|", end=' ')
- for j in range(len(List1)):
- print(List1[j][i].rjust(int(List2[j])), end=' ')
- print("|", end=' ')
- print()
-
- str_nums = [str(int(x)) for x in nums]
- tableData = [
- classes, str_nums
- ]
- colWidths = [0]*len(tableData)
- len1 = 0
- for i in range(len(tableData)):
- for j in range(len(tableData[i])):
- if len(tableData[i][j]) > colWidths[i]:
- colWidths[i] = len(tableData[i][j])
- printTable(tableData, colWidths)
-
- if photo_nums[0] <= 500:
- print("训练集数量小于500,属于较小的数据量,请注意设置较大的训练世代(Epoch)以满足足够的梯度下降次数(Step)。")
-
- if np.sum(nums) == 0:
- print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
- print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
- print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
- print("(重要的事情说三遍)。")
该代码的运行结果:
该代码为训练生成2008年的训练集(2008_train.txt
)和验证集(2008_val.txt
)的列表文件。从VOC数据集的现有标注文件中提取图像ID,并将这些ID对应的图像路径写入训练集(2008_train.txt
)和验证集(2008_val.txt
)的列表文件中,同时还对这些图像的标注进行了转换。这些列表文件用于模型训练时的数据加载。
训练的代码main要放在VOCdevkit的同级目录。
主要负责设置训练所需的损失函数、日志记录、混合精度训练、模型配置和数据加载,迭代数据而得到最优的模型,以便后面的检测。监控训练过程中的损失变化,确保模型正常收敛。
以下是代码train部分
- #-------------------------------------#
- # 对数据集进行训练
- #-------------------------------------#
- import datetime
- import os
- os.environ['KMP_DUPLICATE_LIB_OK']='True'
- from functools import partial
-
- import numpy as np
- import torch
- import torch.backends.cudnn as cudnn
- import torch.distributed as dist
- import torch.nn as nn
- import torch.optim as optim
- from torch import nn
- from torch.utils.data import DataLoader
-
- from nets.yolo import YoloBody
- from nets.yolo_training import (YOLOLoss, get_lr_scheduler, set_optimizer_lr,
- weights_init)
- from utils.callbacks import EvalCallback, LossHistory
- from utils.dataloader import YoloDataset, yolo_dataset_collate
- from utils.utils import (get_anchors, get_classes, seed_everything,
- show_config, worker_init_fn)
- from utils.utils_fit import fit_one_epoch
-
- '''
- 训练自己的目标检测模型一定需要注意以下几点:
- 1、训练前仔细检查自己的格式是否满足要求,该库要求数据集格式为VOC格式,需要准备好的内容有输入图片和标签
- 输入图片为.jpg图片,无需固定大小,传入训练前会自动进行resize。
- 灰度图会自动转成RGB图片进行训练,无需自己修改。
- 输入图片如果后缀非jpg,需要自己批量转成jpg后再开始训练。
- 标签为.xml格式,文件中会有需要检测的目标信息,标签文件和输入图片文件相对应。
- 2、损失值的大小用于判断是否收敛,比较重要的是有收敛的趋势,即验证集损失不断下降,如果验证集损失基本上不改变的话,模型基本上就收敛了。
- 损失值的具体大小并没有什么意义,大和小只在于损失的计算方式,并不是接近于0才好。如果想要让损失好看点,可以直接到对应的损失函数里面除上10000。
- 训练过程中的损失值会保存在logs文件夹下的loss_%Y_%m_%d_%H_%M_%S文件夹中
-
- 3、训练好的权值文件保存在logs文件夹中,每个训练世代(Epoch)包含若干训练步长(Step),每个训练步长(Step)进行一次梯度下降。
- 如果只是训练了几个Step是不会保存的,Epoch和Step的概念要捋清楚一下。
- '''
- if __name__ == "__main__":
- #---------------------------------#
- # Cuda 是否使用Cuda
- # 没有GPU可以设置成False
- #---------------------------------#
- Cuda = False
- #----------------------------------------------#
- # Seed 用于固定随机种子
- # 使得每次独立训练都可以获得一样的结果
- #----------------------------------------------#
- seed = 11
- #---------------------------------------------------------------------#
- # distributed 用于指定是否使用单机多卡分布式运行
- # 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
- # Windows系统下默认使用DP模式调用所有显卡,不支持DDP。
- # DP模式:
- # 设置 distributed = False
- # 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py
- # DDP模式:
- # 设置 distributed = True
- # 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
- #---------------------------------------------------------------------#
- distributed = False
- #---------------------------------------------------------------------#
- # sync_bn 是否使用sync_bn,DDP模式多卡可用
- #---------------------------------------------------------------------#
- sync_bn = False
- #---------------------------------------------------------------------#
- # fp16 是否使用混合精度训练
- # 可减少约一半的显存、需要pytorch1.7.1以上
- #---------------------------------------------------------------------#
- fp16 = False
- #---------------------------------------------------------------------#
- # classes_path 指向model_data下的txt,与自己训练的数据集相关
- # 训练前一定要修改classes_path,使其对应自己的数据集
- #---------------------------------------------------------------------#
- classes_path = 'model_data/class.txt'
- #---------------------------------------------------------------------#
- # anchors_path 代表先验框对应的txt文件,一般不修改。
- # anchors_mask 用于帮助代码找到对应的先验框,一般不修改。
- #---------------------------------------------------------------------#
- anchors_path = 'model_data/yolo_anchors.txt'
- anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]
- #----------------------------------------------------------------------------------------------------------------------------#
- # 权值文件的下载请看README,可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的,因为特征是通用的。
- # 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分,用于进行特征提取。
- # 预训练权重对于99%的情况都必须要用,不用的话主干部分的权值太过随机,特征提取效果不明显,网络训练的结果也不会好
- #
- # 如果训练过程中存在中断训练的操作,可以将model_path设置成logs文件夹下的权值文件,将已经训练了一部分的权值再次载入。
- # 同时修改下方的 冻结阶段 或者 解冻阶段 的参数,来保证模型epoch的连续性。
- #
- # 当model_path = ''的时候不加载整个模型的权值。
- #
- # 此处使用的是整个模型的权重,因此是在train.py进行加载的,下面的pretrain不影响此处的权值加载。
- # 如果想要让模型从主干的预训练权值开始训练,则设置model_path = '',下面的pretrain = True,此时仅加载主干。
- # 如果想要让模型从0开始训练,则设置model_path = '',下面的pretrain = Fasle,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
- #
- # 一般来讲,网络从0开始的训练效果会很差,因为权值太过随机,特征提取效果不明显,因此非常、非常、非常不建议大家从0开始训练!
- # 从0开始训练有两个方案:
- # 1、得益于Mosaic数据增强方法强大的数据增强能力,将UnFreeze_Epoch设置的较大(300及以上)、batch较大(16及以上)、数据较多(万以上)的情况下,
- # 可以设置mosaic=True,直接随机初始化参数开始训练,但得到的效果仍然不如有预训练的情况。(像COCO这样的大数据集可以这样做)
- # 2、了解imagenet数据集,首先训练分类模型,获得网络的主干部分权值,分类模型的 主干部分 和该模型通用,基于此进行训练。
- #----------------------------------------------------------------------------------------------------------------------------#
- model_path = 'model_data/yolo4_weights.pth'
- #------------------------------------------------------#
- # input_shape 输入的shape大小,一定要是32的倍数
- #------------------------------------------------------#
- input_shape = [416, 416]
- #----------------------------------------------------------------------------------------------------------------------------#
- # pretrained 是否使用主干网络的预训练权重,此处使用的是主干的权重,因此是在模型构建的时候进行加载的。
- # 如果设置了model_path,则主干的权值无需加载,pretrained的值无意义。
- # 如果不设置model_path,pretrained = True,此时仅加载主干开始训练。
- # 如果不设置model_path,pretrained = False,Freeze_Train = Fasle,此时从0开始训练,且没有冻结主干的过程。
- #----------------------------------------------------------------------------------------------------------------------------#
- pretrained = True
- #------------------------------------------------------------------#
- # mosaic 马赛克数据增强。
- # mosaic_prob 每个step有多少概率使用mosaic数据增强,默认50%。
- #
- # mixup 是否使用mixup数据增强,仅在mosaic=True时有效。
- # 只会对mosaic增强后的图片进行mixup的处理。
- # mixup_prob 有多少概率在mosaic后使用mixup数据增强,默认50%。
- # 总的mixup概率为mosaic_prob * mixup_prob。
- #
- # special_aug_ratio 参考YoloX,由于Mosaic生成的训练图片,远远脱离自然图片的真实分布。
- # 当mosaic=True时,本代码会在special_aug_ratio范围内开启mosaic。
- # 默认为前70%个epoch,100个世代会开启70个世代。
- #
- # 余弦退火算法的参数放到下面的lr_decay_type中设置
- #------------------------------------------------------------------#
- mosaic = True
- mosaic_prob = 0.5
- mixup = True
- mixup_prob = 0.5
- special_aug_ratio = 0.7
- #------------------------------------------------------------------#
- # label_smoothing 标签平滑。一般0.01以下。如0.01、0.005。
- #------------------------------------------------------------------#
- label_smoothing = 0
-
- #----------------------------------------------------------------------------------------------------------------------------#
- # 训练分为两个阶段,分别是冻结阶段和解冻阶段。设置冻结阶段是为了满足机器性能不足的同学的训练需求。
- # 冻结训练需要的显存较小,显卡非常差的情况下,可设置Freeze_Epoch等于UnFreeze_Epoch,此时仅仅进行冻结训练。
- #
- # 在此提供若干参数设置建议,各位训练者根据自己的需求进行灵活调整:
- # (一)从整个模型的预训练权重开始训练:
- # Adam:
- # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 100,Freeze_Train = True,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(冻结)
- # Init_Epoch = 0,UnFreeze_Epoch = 100,Freeze_Train = False,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(不冻结)
- # SGD:
- # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 300,Freeze_Train = True,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(冻结)
- # Init_Epoch = 0,UnFreeze_Epoch = 300,Freeze_Train = False,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(不冻结)
- # 其中:UnFreeze_Epoch可以在100-300之间调整。
- # (二)从主干网络的预训练权重开始训练:
- # Adam:
- # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 100,Freeze_Train = True,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(冻结)
- # Init_Epoch = 0,UnFreeze_Epoch = 100,Freeze_Train = False,optimizer_type = 'adam',Init_lr = 1e-3,weight_decay = 0。(不冻结)
- # SGD:
- # Init_Epoch = 0,Freeze_Epoch = 50,UnFreeze_Epoch = 300,Freeze_Train = True,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(冻结)
- # Init_Epoch = 0,UnFreeze_Epoch = 300,Freeze_Train = False,optimizer_type = 'sgd',Init_lr = 1e-2,weight_decay = 5e-4。(不冻结)
- # 其中:由于从主干网络的预训练权重开始训练,主干的权值不一定适合目标检测,需要更多的训练跳出局部最优解。
- # UnFreeze_Epoch可以在150-300之间调整,YOLOV5和YOLOX均推荐使用300。
- # Adam相较于SGD收敛的快一些。因此UnFreeze_Epoch理论上可以小一点,但依然推荐更多的Epoch。
- # (三)从0开始训练:
- # Init_Epoch = 0,UnFreeze_Epoch >= 300,Unfreeze_batch_size >= 16,Freeze_Train = False(不冻结训练)
- # 其中:UnFreeze_Epoch尽量不小于300。optimizer_type = 'sgd',Init_lr = 1e-2,mosaic = True。
- # (四)batch_size的设置:
- # 在显卡能够接受的范围内,以大为好。显存不足与数据集大小无关,提示显存不足(OOM或者CUDA out of memory)请调小batch_size。
- # 受到BatchNorm层影响,batch_size最小为2,不能为1。
- # 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大,因为关系到学习率的自动调整。
- #----------------------------------------------------------------------------------------------------------------------------#
- #------------------------------------------------------------------#
- # 冻结阶段训练参数
- # 此时模型的主干被冻结了,特征提取网络不发生改变
- # 占用的显存较小,仅对网络进行微调
- # Init_Epoch 模型当前开始的训练世代,其值可以大于Freeze_Epoch,如设置:
- # Init_Epoch = 60、Freeze_Epoch = 50、UnFreeze_Epoch = 100
- # 会跳过冻结阶段,直接从60代开始,并调整对应的学习率。
- # (断点续练时使用)
- # Freeze_Epoch 模型冻结训练的Freeze_Epoch
- # (当Freeze_Train=False时失效)
- # Freeze_batch_size 模型冻结训练的batch_size
- # (当Freeze_Train=False时失效)
- #------------------------------------------------------------------#
- Init_Epoch = 0
- Freeze_Epoch = 20
- Freeze_batch_size = 2
- #------------------------------------------------------------------#
- # 解冻阶段训练参数
- # 此时模型的主干不被冻结了,特征提取网络会发生改变
- # 占用的显存较大,网络所有的参数都会发生改变
- # UnFreeze_Epoch 模型总共训练的epoch
- # SGD需要更长的时间收敛,因此设置较大的UnFreeze_Epoch
- # Adam可以使用相对较小的UnFreeze_Epoch
- # Unfreeze_batch_size 模型在解冻后的batch_size
- #------------------------------------------------------------------#
- UnFreeze_Epoch = 20
- Unfreeze_batch_size = 2
- #------------------------------------------------------------------#
- # Freeze_Train 是否进行冻结训练
- # 默认先冻结主干训练后解冻训练。
- #------------------------------------------------------------------#
- Freeze_Train = True
-
- #------------------------------------------------------------------#
- # 其它训练参数:学习率、优化器、学习率下降有关
- #------------------------------------------------------------------#
- #------------------------------------------------------------------#
- # Init_lr 模型的最大学习率
- # Min_lr 模型的最小学习率,默认为最大学习率的0.01
- #------------------------------------------------------------------#
- Init_lr = 1e-2
- Min_lr = Init_lr * 0.01
- #------------------------------------------------------------------#
- # optimizer_type 使用到的优化器种类,可选的有adam、sgd
- # 当使用Adam优化器时建议设置 Init_lr=1e-3
- # 当使用SGD优化器时建议设置 Init_lr=1e-2
- # momentum 优化器内部使用到的momentum参数
- # weight_decay 权值衰减,可防止过拟合
- # adam会导致weight_decay错误,使用adam时建议设置为0。
- #------------------------------------------------------------------#
- optimizer_type = "sgd"
- momentum = 0.937
- weight_decay = 5e-4
- #------------------------------------------------------------------#
- # lr_decay_type 使用到的学习率下降方式,可选的有step、cos
- #------------------------------------------------------------------#
- lr_decay_type = "cos"
- #------------------------------------------------------------------#
- # focal_loss 是否使用Focal Loss平衡正负样本
- # focal_alpha Focal Loss的正负样本平衡参数
- # focal_gamma Focal Loss的难易分类样本平衡参数
- #------------------------------------------------------------------#
- focal_loss = False
- focal_alpha = 0.25
- focal_gamma = 2
- #------------------------------------------------------------------#
- # iou_type 使用什么iou损失,ciou或者siou
- #------------------------------------------------------------------#
- iou_type = 'ciou'
- #------------------------------------------------------------------#
- # save_period 多少个epoch保存一次权值
- #------------------------------------------------------------------#
- save_period = 10
- #------------------------------------------------------------------#
- # save_dir 权值与日志文件保存的文件夹
- #------------------------------------------------------------------#
- save_dir = 'logs'
- #------------------------------------------------------------------#
- # eval_flag 是否在训练时进行评估,评估对象为验证集
- # 安装pycocotools库后,评估体验更佳。
- # eval_period 代表多少个epoch评估一次,不建议频繁的评估
- # 评估需要消耗较多的时间,频繁评估会导致训练非常慢
- # 此处获得的mAP会与get_map.py获得的会有所不同,原因有二:
- # (一)此处获得的mAP为验证集的mAP。
- # (二)此处设置评估参数较为保守,目的是加快评估速度。
- #------------------------------------------------------------------#
- eval_flag = True
- eval_period = 10
- #------------------------------------------------------------------#
- # num_workers 用于设置是否使用多线程读取数据
- # 开启后会加快数据读取速度,但是会占用更多内存
- # 内存较小的电脑可以设置为2或者0
- #------------------------------------------------------------------#
- num_workers = 4
-
- #------------------------------------------------------#
- # train_annotation_path 训练图片路径和标签
- # val_annotation_path 验证图片路径和标签
- #------------------------------------------------------#
- train_annotation_path = '2008_train.txt'
- val_annotation_path = '2008_val.txt'
-
- seed_everything(seed)
- #------------------------------------------------------#
- # 设置用到的显卡
- #------------------------------------------------------#
- ngpus_per_node = torch.cuda.device_count()
- if distributed:
- dist.init_process_group(backend="nccl")
- local_rank = int(os.environ["LOCAL_RANK"])
- rank = int(os.environ["RANK"])
- device = torch.device("cuda", local_rank)
- if local_rank == 0:
- print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...")
- print("Gpu Device Count : ", ngpus_per_node)
- else:
- device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
- local_rank = 0
- rank = 0
-
- #------------------------------------------------------#
- # 获取classes和anchor
- #------------------------------------------------------#
- class_names, num_classes = get_classes(classes_path)
- anchors, num_anchors = get_anchors(anchors_path)
-
- #------------------------------------------------------#
- # 创建yolo模型
- #------------------------------------------------------#
- model = YoloBody(anchors_mask, num_classes, pretrained = pretrained)
- if not pretrained:
- weights_init(model)
- if model_path != '':
- #------------------------------------------------------#
- # 权值文件请看README,百度网盘下载
- #------------------------------------------------------#
- if local_rank == 0:
- print('Load weights {}.'.format(model_path))
-
- #------------------------------------------------------#
- # 根据预训练权重的Key和模型的Key进行加载
- #------------------------------------------------------#
- model_dict = model.state_dict()
- pretrained_dict = torch.load(model_path, map_location = device)
- load_key, no_load_key, temp_dict = [], [], {}
- for k, v in pretrained_dict.items():
- if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
- temp_dict[k] = v
- load_key.append(k)
- else:
- no_load_key.append(k)
- model_dict.update(temp_dict)
- model.load_state_dict(model_dict)
- #------------------------------------------------------#
- # 显示没有匹配上的Key
- #------------------------------------------------------#
- if local_rank == 0:
- print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
- print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))
- print("\n\033[1;33;44m温馨提示,head部分没有载入是正常现象,Backbone部分没有载入是错误的。\033[0m")
-
- #----------------------#
- # 获得损失函数
- #----------------------#
- yolo_loss = YOLOLoss(anchors, num_classes, input_shape, Cuda, anchors_mask, label_smoothing, focal_loss, focal_alpha, focal_gamma, iou_type)
- #----------------------#
- # 记录Loss
- #----------------------#
- if local_rank == 0:
- time_str = datetime.datetime.strftime(datetime.datetime.now(),'%Y_%m_%d_%H_%M_%S')
- log_dir = os.path.join(save_dir, "loss_" + str(time_str))
- loss_history = LossHistory(log_dir, model, input_shape=input_shape)
- else:
- loss_history = None
-
- #------------------------------------------------------------------#
- # torch 1.2不支持amp,建议使用torch 1.7.1及以上正确使用fp16
- # 因此torch1.2这里显示"could not be resolve"
- #------------------------------------------------------------------#
- if fp16:
- from torch.cuda.amp import GradScaler as GradScaler
- scaler = GradScaler()
- else:
- scaler = None
-
- model_train = model.train()
- #----------------------------#
- # 多卡同步Bn
- #----------------------------#
- if sync_bn and ngpus_per_node > 1 and distributed:
- model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train)
- elif sync_bn:
- print("Sync_bn is not support in one gpu or not distributed.")
-
- if Cuda:
- if distributed:
- #----------------------------#
- # 多卡平行运行
- #----------------------------#
- model_train = model_train.cuda(local_rank)
- model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank], find_unused_parameters=True)
- else:
- model_train = torch.nn.DataParallel(model)
- cudnn.benchmark = True
- model_train = model_train.cuda()
-
- #---------------------------#
- # 读取数据集对应的txt
- #---------------------------#
- with open(train_annotation_path, encoding='utf-8') as f:
- train_lines = f.readlines()
- with open(val_annotation_path, encoding='utf-8') as f:
- val_lines = f.readlines()
- num_train = len(train_lines)
- num_val = len(val_lines)
-
- if local_rank == 0:
- show_config(
- classes_path = classes_path, anchors_path = anchors_path, anchors_mask = anchors_mask, model_path = model_path, input_shape = input_shape, \
- Init_Epoch = Init_Epoch, Freeze_Epoch = Freeze_Epoch, UnFreeze_Epoch = UnFreeze_Epoch, Freeze_batch_size = Freeze_batch_size, Unfreeze_batch_size = Unfreeze_batch_size, Freeze_Train = Freeze_Train, \
- Init_lr = Init_lr, Min_lr = Min_lr, optimizer_type = optimizer_type, momentum = momentum, lr_decay_type = lr_decay_type, \
- save_period = save_period, save_dir = save_dir, num_workers = num_workers, num_train = num_train, num_val = num_val
- )
- #---------------------------------------------------------#
- # 总训练世代指的是遍历全部数据的总次数
- # 总训练步长指的是梯度下降的总次数
- # 每个训练世代包含若干训练步长,每个训练步长进行一次梯度下降。
- # 此处仅建议最低训练世代,上不封顶,计算时只考虑了解冻部分
- #----------------------------------------------------------#
- wanted_step = 5e4 if optimizer_type == "sgd" else 1.5e4
- total_step = num_train // Unfreeze_batch_size * UnFreeze_Epoch
- if total_step <= wanted_step:
- if num_train // Unfreeze_batch_size == 0:
- raise ValueError('数据集过小,无法进行训练,请扩充数据集。')
- wanted_epoch = wanted_step // (num_train // Unfreeze_batch_size) + 1
- print("\n\033[1;33;44m[Warning] 使用%s优化器时,建议将训练总步长设置到%d以上。\033[0m"%(optimizer_type, wanted_step))
- print("\033[1;33;44m[Warning] 本次运行的总训练数据量为%d,Unfreeze_batch_size为%d,共训练%d个Epoch,计算出总训练步长为%d。\033[0m"%(num_train, Unfreeze_batch_size, UnFreeze_Epoch, total_step))
- print("\033[1;33;44m[Warning] 由于总训练步长为%d,小于建议总步长%d,建议设置总世代为%d。\033[0m"%(total_step, wanted_step, wanted_epoch))
-
- #------------------------------------------------------#
- # 主干特征提取网络特征通用,冻结训练可以加快训练速度
- # 也可以在训练初期防止权值被破坏。
- # Init_Epoch为起始世代
- # Freeze_Epoch为冻结训练的世代
- # UnFreeze_Epoch总训练世代
- # 提示OOM或者显存不足请调小Batch_size
- #------------------------------------------------------#
- if True:
- UnFreeze_flag = False
- #------------------------------------#
- # 冻结一定部分训练
- #------------------------------------#
- if Freeze_Train:
- for param in model.backbone.parameters():
- param.requires_grad = False
-
- #-------------------------------------------------------------------#
- # 如果不冻结训练的话,直接设置batch_size为Unfreeze_batch_size
- #-------------------------------------------------------------------#
- batch_size = Freeze_batch_size if Freeze_Train else Unfreeze_batch_size
-
- #-------------------------------------------------------------------#
- # 判断当前batch_size,自适应调整学习率
- #-------------------------------------------------------------------#
- nbs = 64
- lr_limit_max = 1e-3 if optimizer_type in ['adam', 'adamw'] else 5e-2
- lr_limit_min = 3e-4 if optimizer_type in ['adam', 'adamw'] else 5e-4
- Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
- Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
-
- #---------------------------------------#
- # 根据optimizer_type选择优化器
- #---------------------------------------#
- pg0, pg1, pg2 = [], [], []
- for k, v in model.named_modules():
- if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter):
- pg2.append(v.bias)
- if isinstance(v, nn.BatchNorm2d) or "bn" in k:
- pg0.append(v.weight)
- elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter):
- pg1.append(v.weight)
- optimizer = {
- 'adam' : optim.Adam(pg0, Init_lr_fit, betas = (momentum, 0.999)),
- 'adamw' : optim.AdamW(pg0, Init_lr_fit, betas = (momentum, 0.999)),
- 'sgd' : optim.SGD(pg0, Init_lr_fit, momentum = momentum, nesterov=True)
- }[optimizer_type]
- optimizer.add_param_group({"params": pg1, "weight_decay": weight_decay})
- optimizer.add_param_group({"params": pg2})
-
- #---------------------------------------#
- # 获得学习率下降的公式
- #---------------------------------------#
- lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
-
- #---------------------------------------#
- # 判断每一个世代的长度
- #---------------------------------------#
- epoch_step = num_train // batch_size
- epoch_step_val = num_val // batch_size
-
- if epoch_step == 0 or epoch_step_val == 0:
- raise ValueError("数据集过小,无法继续进行训练,请扩充数据集。")
-
- #---------------------------------------#
- # 构建数据集加载器。
- #---------------------------------------#
- train_dataset = YoloDataset(train_lines, input_shape, num_classes, epoch_length = UnFreeze_Epoch, \
- mosaic=mosaic, mixup=mixup, mosaic_prob=mosaic_prob, mixup_prob=mixup_prob, train=True, special_aug_ratio=special_aug_ratio)
- val_dataset = YoloDataset(val_lines, input_shape, num_classes, epoch_length = UnFreeze_Epoch, \
- mosaic=False, mixup=False, mosaic_prob=0, mixup_prob=0, train=False, special_aug_ratio=0)
-
- if distributed:
- train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True,)
- val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False,)
- batch_size = batch_size // ngpus_per_node
- shuffle = False
- else:
- train_sampler = None
- val_sampler = None
- shuffle = True
-
- gen = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
- drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler,
- worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
- gen_val = DataLoader(val_dataset , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
- drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler,
- worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
-
- #----------------------#
- # 记录eval的map曲线
- #----------------------#
- if local_rank == 0:
- eval_callback = EvalCallback(model, input_shape, anchors, anchors_mask, class_names, num_classes, val_lines, log_dir, Cuda, \
- eval_flag=eval_flag, period=eval_period)
- else:
- eval_callback = None
-
- #---------------------------------------#
- # 开始模型训练
- #---------------------------------------#
- for epoch in range(Init_Epoch, UnFreeze_Epoch):
- #---------------------------------------#
- # 如果模型有冻结学习部分
- # 则解冻,并设置参数
- #---------------------------------------#
- if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train:
- batch_size = Unfreeze_batch_size
-
- #-------------------------------------------------------------------#
- # 判断当前batch_size,自适应调整学习率
- #-------------------------------------------------------------------#
- nbs = 64
- lr_limit_max = 1e-3 if optimizer_type in ['adam', 'adamw'] else 5e-2
- lr_limit_min = 3e-4 if optimizer_type in ['adam', 'adamw'] else 5e-4
- Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
- Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
- #---------------------------------------#
- # 获得学习率下降的公式
- #---------------------------------------#
- lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch)
-
- for param in model.backbone.parameters():
- param.requires_grad = True
-
- epoch_step = num_train // batch_size
- epoch_step_val = num_val // batch_size
-
- if epoch_step == 0 or epoch_step_val == 0:
- raise ValueError("数据集过小,无法继续进行训练,请扩充数据集。")
-
- if distributed:
- batch_size = batch_size // ngpus_per_node
-
- gen = DataLoader(train_dataset, shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
- drop_last=True, collate_fn=yolo_dataset_collate, sampler=train_sampler,
- worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
- gen_val = DataLoader(val_dataset , shuffle = shuffle, batch_size = batch_size, num_workers = num_workers, pin_memory=True,
- drop_last=True, collate_fn=yolo_dataset_collate, sampler=val_sampler,
- worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed))
-
- UnFreeze_flag = True
-
- gen.dataset.epoch_now = epoch
- gen_val.dataset.epoch_now = epoch
-
- if distributed:
- train_sampler.set_epoch(epoch)
-
- set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
-
- fit_one_epoch(model_train, model, yolo_loss, loss_history, eval_callback, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, UnFreeze_Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rank)
-
- if distributed:
- dist.barrier()
-
- if local_rank == 0:
- loss_history.writer.close()
这里迭代20次,得到一个最优的模型并保存在best_epoch_weights.pth文件。
该代码的运行结果:
用于执行不同对象检测任务的脚本,它允许用户根据所选模式执行不同的操作。
使用训练好的模型对测试图片进行对象检测,将检测结果可视化,并展示在测试图片上。
以下是代码predict部分
- #-----------------------------------------------------------------------#
- # predict.py将单张图片预测、摄像头检测、FPS测试和目录遍历检测等功能
- # 整合到了一个py文件中,通过指定mode进行模式的修改。
- #-----------------------------------------------------------------------#
- import time
-
- import cv2
- import numpy as np
- from PIL import Image
-
- from yolo import YOLO, YOLO_ONNX
-
- if __name__ == "__main__":
- #----------------------------------------------------------------------------------------------------------#
- # mode用于指定测试的模式:
- # 'predict' 表示单张图片预测,如果想对预测过程进行修改,如保存图片,截取对象等,可以先看下方详细的注释
- # 'video' 表示视频检测,可调用摄像头或者视频进行检测,详情查看下方注释。
- # 'fps' 表示测试fps,使用的图片是img里面的street.jpg,详情查看下方注释。
- # 'dir_predict' 表示遍历文件夹进行检测并保存。默认遍历img文件夹,保存img_out文件夹,详情查看下方注释。
- # 'heatmap' 表示进行预测结果的热力图可视化,详情查看下方注释。
- # 'export_onnx' 表示将模型导出为onnx,需要pytorch1.7.1以上。
- # 'predict_onnx' 表示利用导出的onnx模型进行预测,相关参数的修改在yolo.py_416行左右处的YOLO_ONNX
- #----------------------------------------------------------------------------------------------------------#
- mode = "predict"
- #-------------------------------------------------------------------------#
- # crop 指定了是否在单张图片预测后对目标进行截取
- # count 指定了是否进行目标的计数
- # crop、count仅在mode='predict'时有效
- #-------------------------------------------------------------------------#
- crop = False
- count = False
- #----------------------------------------------------------------------------------------------------------#
- # video_path 用于指定视频的路径,当video_path=0时表示检测摄像头
- # 想要检测视频,则设置如video_path = "xxx.mp4"即可,代表读取出根目录下的xxx.mp4文件。
- # video_save_path 表示视频保存的路径,当video_save_path=""时表示不保存
- # 想要保存视频,则设置如video_save_path = "yyy.mp4"即可,代表保存为根目录下的yyy.mp4文件。
- # video_fps 用于保存的视频的fps
- #
- # video_path、video_save_path和video_fps仅在mode='video'时有效
- # 保存视频时需要ctrl+c退出或者运行到最后一帧才会完成完整的保存步骤。
- #----------------------------------------------------------------------------------------------------------#
- video_path = 0
- video_save_path = ""
- video_fps = 25.0
- #----------------------------------------------------------------------------------------------------------#
- # test_interval 用于指定测量fps的时候,图片检测的次数。理论上test_interval越大,fps越准确。
- # fps_image_path 用于指定测试的fps图片
- #
- # test_interval和fps_image_path仅在mode='fps'有效
- #----------------------------------------------------------------------------------------------------------#
- test_interval = 100
- fps_image_path = "img/street.jpg"
- #-------------------------------------------------------------------------#
- # dir_origin_path 指定了用于检测的图片的文件夹路径
- # dir_save_path 指定了检测完图片的保存路径
- #
- # dir_origin_path和dir_save_path仅在mode='dir_predict'时有效
- #-------------------------------------------------------------------------#
- dir_origin_path = "img/"
- dir_save_path = "img_out/"
- #-------------------------------------------------------------------------#
- # heatmap_save_path 热力图的保存路径,默认保存在model_data下
- #
- # heatmap_save_path仅在mode='heatmap'有效
- #-------------------------------------------------------------------------#
- heatmap_save_path = "model_data/heatmap_vision.png"
- #-------------------------------------------------------------------------#
- # simplify 使用Simplify onnx
- # onnx_save_path 指定了onnx的保存路径
- #-------------------------------------------------------------------------#
- simplify = True
- onnx_save_path = "model_data/models.onnx"
-
- if mode != "predict_onnx":
- yolo = YOLO()
- else:
- yolo = YOLO_ONNX()
-
- if mode == "predict":
- '''
- 1、如果想要进行检测完的图片的保存,利用r_image.save("img.jpg")即可保存,直接在predict.py里进行修改即可。
- 2、如果想要获得预测框的坐标,可以进入yolo.detect_image函数,在绘图部分读取top,left,bottom,right这四个值。
- 3、如果想要利用预测框截取下目标,可以进入yolo.detect_image函数,在绘图部分利用获取到的top,left,bottom,right这四个值
- 在原图上利用矩阵的方式进行截取。
- 4、如果想要在预测图上写额外的字,比如检测到的特定目标的数量,可以进入yolo.detect_image函数,在绘图部分对predicted_class进行判断,
- 比如判断if predicted_class == 'car': 即可判断当前目标是否为车,然后记录数量即可。利用draw.text即可写字。
- '''
- while True:
- img = input('Input image filename:')
- try:
- image = Image.open(img)
- except:
- print('Open Error! Try again!')
- continue
- else:
- r_image = yolo.detect_image(image, crop = crop, count=count)
- r_image.show()
-
- elif mode == "video":
- capture = cv2.VideoCapture(video_path)
- if video_save_path!="":
- fourcc = cv2.VideoWriter_fourcc(*'XVID')
- size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
- out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
-
- ref, frame = capture.read()
- if not ref:
- raise ValueError("未能正确读取摄像头(视频),请注意是否正确安装摄像头(是否正确填写视频路径)。")
-
- fps = 0.0
- while(True):
- t1 = time.time()
- # 读取某一帧
- ref, frame = capture.read()
- if not ref:
- break
- # 格式转变,BGRtoRGB
- frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
- # 转变成Image
- frame = Image.fromarray(np.uint8(frame))
- # 进行检测
- frame = np.array(yolo.detect_image(frame))
- # RGBtoBGR满足opencv显示格式
- frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
-
- fps = ( fps + (1./(time.time()-t1)) ) / 2
- print("fps= %.2f"%(fps))
- frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
- cv2.imshow("video",frame)
- c= cv2.waitKey(1) & 0xff
- if video_save_path!="":
- out.write(frame)
-
- if c==27:
- capture.release()
- break
-
- print("Video Detection Done!")
- capture.release()
- if video_save_path!="":
- print("Save processed video to the path :" + video_save_path)
- out.release()
- cv2.destroyAllWindows()
-
- elif mode == "fps":
- img = Image.open(fps_image_path)
- tact_time = yolo.get_FPS(img, test_interval)
- print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1')
-
- elif mode == "dir_predict":
- import os
-
- from tqdm import tqdm
-
- img_names = os.listdir(dir_origin_path)
- for img_name in tqdm(img_names):
- if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
- image_path = os.path.join(dir_origin_path, img_name)
- image = Image.open(image_path)
- r_image = yolo.detect_image(image)
- if not os.path.exists(dir_save_path):
- os.makedirs(dir_save_path)
- r_image.save(os.path.join(dir_save_path, img_name.replace(".jpg", ".png")), quality=95, subsampling=0)
-
- elif mode == "heatmap":
- while True:
- img = input('Input image filename:')
- try:
- image = Image.open(img)
- except:
- print('Open Error! Try again!')
- continue
- else:
- yolo.detect_heatmap(image, heatmap_save_path)
-
- elif mode == "export_onnx":
- yolo.convert_to_onnx(simplify, onnx_save_path)
-
- elif mode == "predict_onnx":
- while True:
- img = input('Input image filename:')
- try:
- image = Image.open(img)
- except:
- print('Open Error! Try again!')
- continue
- else:
- r_image = yolo.detect_image(image)
- r_image.show()
- else:
- raise AssertionError("Please specify the correct mode: 'predict', 'video', 'fps', 'heatmap', 'export_onnx', 'dir_predict'.")
根据调节confidence的参数来显示正确框出之前标注的数据。
该代码的运行结果:
本实验的主要目标是使用YOLOv4算法,通过自己标注的数据集进行训练,得到一个训练好的模型。随后,我们将使用这个训练好的模型对测试图片进行对象检测,并展示检测结果。
经过训练,我们得到了一个针对特定对象的YOLOv4模型。在测试阶段,该模型能够准确地检测出测试图片中的对象,并给出相应的类别和位置信息。通过可视化展示,我们可以直观地看到模型的检测效果。
本实验仍存在一定的局限性。例如,在处理复杂背景和遮挡目标时,模型的检测性能可能会受到一定影响。未来,将进一步探索如何提升模型在复杂场景下的检测能力,并尝试将模型应用于更多的实际场景中。
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。