前言

雖然說yolo-fastestV2在coco數(shù)據(jù)集上map只達到了24.1，但是應(yīng)付一些類別少的問題還是可以的。主要是這個速度是真的香！簡單來說就是一個快到飛起的模型。

1.關(guān)于yolo-fastestV2

github地址如下：yolo-fastestV2
yolo-fastestV2采用了輕量化網(wǎng)絡(luò)shufflenetV2為backbone，筆者在這里就不詳解yolo-fastestV2了，只講怎么訓(xùn)練自己的yolo-fastestV2模型。

2.修改訓(xùn)練代碼

訓(xùn)練部分的代碼只修改了一小部分，如下：

import os
import math
import time
import argparse
import numpy as np
from tqdm import tqdm
from numpy.testing._private.utils import print_assert_equal

import torch
from torch import optim
from torch.utils.data import dataset
from numpy.core.fromnumeric import shape

from torchsummary import summary

import utils.loss
import utils.utils
import utils.datasets
import model.detector

if __name__ == '__main__':
    # 指定訓(xùn)練配置文件
    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("訓(xùn)練配置:")
    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 = False
    premodel_path = cfg["pre_weights"]
    if premodel_path != None and os.path.exists(premodel_path):
        load_param = True

    model = model.detector.Detector(cfg["classes"], cfg["anchor_num"], load_param).to(device)
    summary(model, input_size=(3, cfg["height"], cfg["width"]))

    # 加載預(yù)訓(xùn)練模型參數(shù)
    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 = 0
    for epoch in range(cfg["epochs"]):
        model.train()
        pbar = tqdm(train_dataloader)

        for imgs, targets in pbar:
            imgs = imgs.to(device).float() / 255.0
            targets = 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"] * scale
                    
                lr = 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 += 1

        if 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()

這里筆者修改了加載預(yù)訓(xùn)練參數(shù)的部分，如果不修改的話，在后面修改訓(xùn)練自己模型的類別會有問題。因為作者的預(yù)訓(xùn)練模型是coco數(shù)據(jù)集80類別，要是不修改簡單來說只能讓類別數(shù)為80，這是肯定不行的。

3.制作自己的目標檢測數(shù)據(jù)集

(1).收集圖片

首先使用opencv收集數(shù)據(jù)集的圖片，代碼如下：

import cv2

cap=cv2.VideoCapture(1)
i=1
while(cap.isOpened()):
    ret,f=cap.read()
    c=cv2.waitKey(1)
    cv2.imshow("f",f)
    if c==27:
        break
    elif c==ord("q"):
        print(i)
        path="./data/"+str(i)+".jpg"
        cv2.imwrite(path,f)
        i+=1
    else:
        pass

cap.release()
cv2.destroyAllWindows()

(2).打標簽

上面的代碼只要修改一下path就能保存到自己指定的路徑了，圖片建議保存成jpg格式，當(dāng)然也可以不，但是建議全部圖片保存為同一種后綴。當(dāng)然圖片的分辨率也得是一致的。
然后安裝標注工具labelimg，一般使用在cmd使用如下命令就能安裝了

pip install labelimg

然后在cmd鍵入labelimg就能進入labelimg了，然后給自己的數(shù)據(jù)集打上標簽。打成的標簽xml文件和圖片放在同一個文件夾里面。

(3).數(shù)據(jù)處理

然后運行如下代碼：

import xml.etree.ElementTree as et
import os

classes=["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_list

def 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創(chuàng)建迭代器，尋找所有object的節(jié)點
            for obj in root.iter('object'):
                #按找標記名尋找匹配的第一個元素，text返回字符串
                cls = obj.find('name').text
                if cls not in classes:
                    continue
                cls_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ù)訓(xùn)練，其中第二個函數(shù)中的x，y分別為圖片的寬和高。因為作者的代碼是要從txt文件里面獲取歸一化后的框，我們從上面收集的圖片來說分辨率都是一樣的，所以就直接用我們這個固定的寬高分辨率來歸一化，當(dāng)然如果有不同分辨率的圖片的話，可以打標簽前做一個resize操作，也可以在解析xml文件的時候輸出圖片的寬高再用這個寬高歸一化。
下面就是生成用于訓(xùn)練的train.txt和test.txt了，這兩個文件就是儲存了圖片的路徑，方便在訓(xùn)練或者評估的時候找到圖片的路徑。代碼如下：

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_list

def 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] + "." + tu
            ftrain.write(str(pp) + "\n")
        else:
            pp = xml[:-4] + "." + tu
            ftest.write(str(pp) + "\n")
    print("寫入完成")


if __name__ == '__main__':
    wri_txt(path=r"D:\python1\python\pytorch\object-detection\Yolo-FastestV2-main\xml")

然后注意第二個函數(shù)wri_txt函數(shù)中的路徑必須用絕對路徑。

4.修改超參數(shù)

(1).生成先驗框

生成anchors的代碼在作者的genanchors.py中，我們在main函數(shù)中修改traintxt參數(shù)的defaul也就是默認參數(shù)為我們上面寫入的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).修改超參數(shù)

打開data文件夾下的coco.data，把anchors后面的數(shù)字改成我們生成的第一行數(shù)據(jù)即可。
coco.data的數(shù)據(jù)如下：

[name]
model_name=fruit   #模型名字

[train-configure]
epochs=300      #批次數(shù)
steps=150,250
batch_size=2
subdivisions=1
learning_rate=0.001

[model-configure]
pre_weights=./modelzoo/coco2017-0.241078ap-model.pth  #預(yù)訓(xùn)練參數(shù)路徑
classes=2   #分類數(shù)
width=352
height=352
anchor_num=3   #每個特征層的anchor數(shù)
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   #訓(xùn)練集路徑
val=D:\\python1\\python\\pytorch\\object-detection\\Yolo-FastestV2-main\\test.txt    #測試集路徑
names=./data/coco.names   #保存對應(yīng)類別名字的文件

然后按照上面的注釋，把這些超參數(shù)改成自己需要的就好了，一些基礎(chǔ)的超參數(shù)筆者就不解釋了。
隨后把coco.names文件改成自己數(shù)據(jù)的類別，一個類別對應(yīng)一行。

5.開始訓(xùn)練

隨后就可以運行train.py文件用自己的數(shù)據(jù)集訓(xùn)練自己的模型了。在這里筆者建議安裝的torch1.9.0，torchvision0.10.0。這是項目作者的環(huán)境，不然可能會因為api不兼容報錯。
然后萬一訓(xùn)練好的模型泛化不到訓(xùn)練場景以外的場景，最簡單的方法就是添加這個場景的數(shù)據(jù)，將預(yù)訓(xùn)練權(quán)重修改成自己之前訓(xùn)練好的權(quán)重繼續(xù)訓(xùn)練。這樣更節(jié)省時間。然后在數(shù)據(jù)集的收集上的話，建議每個類別的數(shù)量在1500個框以上吧。文章來源地址http://www.zghlxwxcb.cn/news/detail-461900.html

到了這里，關(guān)于深度學(xué)習(xí)-yolo-fastestV2使用自己的數(shù)據(jù)集訓(xùn)練自己的模型的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！