前言
虽然说yolo-fastestV2在coco数据集上map只达到了24.1，但是应付一些类别少的问题还是可以的。主要是这个速度是真的香！简单来说就是一个快到飞起的模型。
1.关于yolo-fastestV2
github地址如下：yolo-fastestV2
yolo-fastestV2采用了轻量化网络shufflenetV2为backbone，笔者在这里就不详解yolo-fastestV2了，只讲怎么训练自己的yolo-fastestV2模型。
2.修改训练代码
训练部分的代码只修改了一小部分，如下：
import os
import math
import time
import argparse
import numpy as np
from tqdm import tqdm
from numpy.testing._private.utils import print_assert_equalimport torch
from torch import optim
from torch.utils.data import dataset
from numpy.core.fromnumeric import shapefrom torchsummary import summaryimport utils.loss
import utils.utils
import utils.datasets
import model.detectorif __name__ == '__main__':# 指定训练配置文件parser = argparse.ArgumentParser()parser.add_argument('--data', type=str, default='./data/coco.data',help='Specify training profile *.data')opt = parser.parse_args()cfg = utils.utils.load_datafile(opt.data)print("训练配置:")print(cfg)train_dataset = utils.datasets.TensorDataset(cfg["train"], cfg["width"], cfg["height"], imgaug = True)val_dataset = utils.datasets.TensorDataset(cfg["val"], cfg["width"], cfg["height"], imgaug = False)batch_size = int(cfg["batch_size"] / cfg["subdivisions"])nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])print(nw)train_dataloader = torch.utils.data.DataLoader(train_dataset,batch_size=batch_size,shuffle=True,collate_fn=utils.datasets.collate_fn,num_workers=nw,pin_memory=True,drop_last=True,persistent_workers=True)val_dataloader = torch.utils.data.DataLoader(val_dataset,batch_size=batch_size,shuffle=False,collate_fn=utils.datasets.collate_fn,num_workers=nw,pin_memory=True,drop_last=False,persistent_workers=True)device = torch.device("cuda" if torch.cuda.is_available() else "cpu")load_param = Falsepremodel_path = cfg["pre_weights"]if premodel_path != None and os.path.exists(premodel_path):load_param = Truemodel = model.detector.Detector(cfg["classes"], cfg["anchor_num"], load_param).to(device)summary(model, input_size=(3, cfg["height"], cfg["width"]))# 加载预训练模型参数if load_param == True:model_dict = model.state_dict()pretrained_dict = torch.load(premodel_path, map_location=device)pretrained_dict = {
k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}model_dict.update(pretrained_dict)model.load_state_dict(model_dict)print("Load finefune model param: %s" % premodel_path)else:print("Initialize weights: model/backbone/backbone.pth")optimizer = optim.SGD(params=model.parameters(),lr=cfg["learning_rate"],momentum=0.949,weight_decay=0.0005,)scheduler = optim.lr_scheduler.MultiStepLR(optimizer,milestones=cfg["steps"],gamma=0.1)print('Starting training for %g epochs...' % cfg["epochs"])batch_num = 0for epoch in range(cfg["epochs"]):model.train()pbar = tqdm(train_dataloader)for imgs, targets in pbar:imgs = imgs.to(device).float() / 255.0targets = targets.to(device)preds = model(imgs)iou_loss, obj_loss, cls_loss, total_loss = utils.loss.compute_loss(preds, targets, cfg, device)total_loss.backward()for g in optimizer.param_groups:warmup_num = 5 * len(train_dataloader)if batch_num <= warmup_num:scale = math.pow(batch_num/warmup_num, 4)g['lr'] = cfg["learning_rate"] * scalelr = g["lr"]if batch_num % cfg["subdivisions"] == 0:optimizer.step()optimizer.zero_grad()info = "Epoch:%d LR:%f CIou:%f Obj:%f Cls:%f Total:%f" % (epoch, lr, iou_loss, obj_loss, cls_loss, total_loss)pbar.set_description(info)batch_num += 1if epoch % 10 == 0 and epoch > 0:model.eval()print("computer mAP...")_, _, AP, _ = utils.utils.evaluation(val_dataloader, cfg, model, device)print("computer PR...")precision, recall, _, f1 = utils.utils.evaluation(val_dataloader, cfg, model, device, 0.3)print("Precision:%f Recall:%f AP:%f F1:%f"%(precision, recall, AP, f1))torch.save(model.state_dict(), "weights/%s-%d-epoch-%fap-model.pth" %(cfg["model_name"], epoch, AP))scheduler.step()这里笔者修改了加载预训练参数的部分，如果不修改的话，在后面修改训练自己模型的类别会有问题。因为作者的预训练模型是coco数据集80类别，要是不修改简单来说只能让类别数为80，这是肯定不行的。
3.制作自己的目标检测数据集
(1).收集图片
首先使用opencv收集数据集的图片，代码如下：
import cv2cap=cv2.VideoCapture(1)
i=1
while(cap.isOpened()):ret,f=cap.read()c=cv2.waitKey(1)cv2.imshow("f",f)if c==27:breakelif c==ord("q"):print(i)path="./data/"+str(i)+".jpg"cv2.imwrite(path,f)i+=1else:passcap.release()
cv2.destroyAllWindows()
(2).打标签
上面的代码只要修改一下path就能保存到自己指定的路径了，图片建议保存成jpg格式，当然也可以不，但是建议全部图片保存为同一种后缀。当然图片的分辨率也得是一致的。
然后安装标注工具labelimg，一般使用在cmd使用如下命令就能安装了
pip install labelimg
然后在cmd键入labelimg就能进入labelimg了，然后给自己的数据集打上标签。打成的标签xml文件和图片放在同一个文件夹里面。
(3).数据处理
然后运行如下代码：
import xml.etree.ElementTree as et
import osclasses=["a","b"]def find_xml(path):'''找到一个文件夹下所有的xml文件，返回一个装这这些文件路径名的所有文件的列表'''file_list=os.listdir(path)xml_list=[]for file in file_list:if file.endswith(".xml"):xml_list.append(path+"/"+file)return xml_listdef xml2txt(path,x=640,y=480):xml_list=find_xml(path)for xml in xml_list:file=xml[:-4]txt=file+".txt"#xml文件解析器,将xml文件解析成元素树tree=et.parse(xml)#拿到树的根root=tree.getroot()with open(txt,'w') as f:#root.iter创建迭代器，寻找所有object的节点for obj in root.iter('object'):#按找标记名寻找匹配的第一个元素，text返回字符串cls = obj.find('name').textif cls not in classes:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text),int(xmlbox.find('ymax').text))w=float(b[2]-b[0])h=float(b[3]-b[1])center_x=float(b[0]+w/2)center_y=float(b[1]+h/2)d=[center_x/x,center_y/y,w/x,h/y]f.write(str(cls_id)+" "+" ".join([str(a) for a in d])+"\n")if __name__ == '__main__':xml2txt("./xml")
上面代码的作用是将每一个xml文件生成一个txt标签文件，用于后续训练，其中第二个函数中的x，y分别为图片的宽和高。因为作者的代码是要从txt文件里面获取归一化后的框，我们从上面收集的图片来说分辨率都是一样的，所以就直接用我们这个固定的宽高分辨率来归一化，当然如果有不同分辨率的图片的话，可以打标签前做一个resize操作，也可以在解析xml文件的时候输出图片的宽高再用这个宽高归一化。
下面就是生成用于训练的train.txt和test.txt了，这两个文件就是储存了图片的路径，方便在训练或者评估的时候找到图片的路径。代码如下：
import os
import random
random.seed(0)def find_xml(path):'''path用绝对路径找到一个文件夹下所有的xml文件，返回一个装这这些文件路径名的所有文件的列表'''file_list=os.listdir(path)xml_list=[]for file in file_list:if file.endswith(".xml"):xml_list.append(path+"/"+file)return xml_listdef wri_txt(path,tu="jpg",train_path="train.txt",test_path="test.txt",train_precent=0.8):xml_list=find_xml(path)train_len=int(len(xml_list)*train_precent)ftrain=open(train_path,'w')ftest=open(test_path,'w')train_list=random.sample(xml_list,train_len)for xml in xml_list:if xml in train_list:pp = xml[:-4] + "." + tuftrain.write(str(pp) + "\n")else:pp = xml[:-4] + "." + tuftest.write(str(pp) + "\n")print("写入完成")if __name__ == '__main__':wri_txt(path=r"D:\python1\python\pytorch\object-detection\Yolo-FastestV2-main\xml")然后注意第二个函数wri_txt函数中的路径必须用绝对路径。
4.修改超参数
(1).生成先验框
生成anchors的代码在作者的genanchors.py中，我们在main函数中修改traintxt参数的defaul也就是默认参数为我们上面写入的train.txt文件的路径即可。然后就会生成一个anchors6.txt文件，打开它如下：
53.74,73.54, 79.35,106.18, 83.11,168.08, 107.49,135.43, 119.73,183.46, 158.47,204.21
0.847014第一行就是我们的6个先验框的宽高了。第二行可以不用管。
(2).修改超参数
打开data文件夹下的coco.data，把anchors后面的数字改成我们生成的第一行数据即可。
coco.data的数据如下：
[name]
model_name=fruit #模型名字[train-configure]
epochs=300 #批次数
steps=150,250
batch_size=2
subdivisions=1
learning_rate=0.001[model-configure]
pre_weights=./modelzoo/coco2017-0.241078ap-model.pth #预训练参数路径
classes=2 #分类数
width=352
height=352
anchor_num=3 #每个特征层的anchor数
anchors=53.74,73.54, 79.35,106.18, 83.11,168.08, 107.49,135.43, 119.73,183.46, 158.47,204.21
[data-configure]
train=D:\\python1\\python\\pytorch\\object-detection\\Yolo-FastestV2-main\train.txt #训练集路径
val=D:\\python1\\python\\pytorch\\object-detection\\Yolo-FastestV2-main\\test.txt #测试集路径
names=./data/coco.names #保存对应类别名字的文件然后按照上面的注释，把这些超参数改成自己需要的就好了，一些基础的超参数笔者就不解释了。
随后把coco.names文件改成自己数据的类别，一个类别对应一行。
5.开始训练
随后就可以运行train.py文件用自己的数据集训练自己的模型了。在这里笔者建议安装的torch1.9.0，torchvision0.10.0。这是项目作者的环境，不然可能会因为api不兼容报错。
然后万一训练好的模型泛化不到训练场景以外的场景，最简单的方法就是添加这个场景的数据，将预训练权重修改成自己之前训练好的权重继续训练。这样更节省时间。然后在数据集的收集上的话，建议每个类别的数量在1500个框以上吧。