一步一步帶你訓練自己的SSD檢測算法

一、前言

  隨着深度學習的快速發展，越來越多的場景都開始應用這麼技術，但是當前的深度學習還是具有一些門檻，厲害的大佬們都在設計各種各樣的網絡，發各種各樣的頂會論文，他們關注的是算法的設計。然而在現實場景中，對於很多人而言，他們更多的是關注如何去使用這個算法，比如我如何快速的使用現有的工具訓練一個滿足自己需求的模型，當前網上淺顯易懂的帶小白構建自己的模型的教程太少了，我決定來彌補一下這個空缺，爲自己爲大家做一點點小貢獻。本文專注於使用SSD網絡從頭訓練自己的數據集，我們詳細講解裏面的細節，具體的內容如下所示感興趣的小夥伴可以掃描下方的二維碼加羣，我們可以一起交流更多的問題。

二、實現細節

1、前提條件

  在執行下面的操作之前，請保證你當前已經具有以下的這些環境。

• Ubuntu 16.04系統
• NVIDIA 顯卡，建議不低於GTX 1080
• CUDA驅動，CUDA8.0\CUDA9.0\CUDA10.0均可，建議設置爲可以隨時切換的模式
• Anaconda3

2、數據標註

  對於一個完整的深度學習流程而言，數據標註是必不可少的一環，數據標註質量會極大的影響算法的性能。當前如果你使用別人已經標註好的數據或者公有的數據集，請直接跳過該步驟。
  下面就來簡單的介紹幾種數據標註工具，當前已經有很多開源的數據標註工具，這些工具基本上可以滿足我們98%的需求，工具沒有好壞之分，合適自己的工具就是好工具，所以大家需要根據自己的情況去選擇合適的標註工具。

2.1 Labelme

Github鏈接

2.1.1 工具特點簡介

• 同時支持Ubuntu / macOS / Windows系統
• 可以使用多邊形、矩形、圓、線和點進行圖像註釋
• 支持VOC和COCO數據格式
• 支持分類、檢測和分割等多個計算機視覺任務

2.1.2 工具安裝

方案1：

conda create --name=labelme python=3.6
source activate labelme
pip install labelme

方案2：

sudo apt-get install python3-pyqt5  # PyQt5
sudo pip3 install labelme

2.1.3 工具使用簡介

  上圖展示了Labelme工具的標註界面，工具的使用留給大家自己去熟悉，在這裏我只會強調幾個需要注意的地方，具體的內容如下所示：

• 建議把需要標註的數據集存放在一個文件夾下面，並以相應的名字來命名
• 該工具默認的標註方式是多邊形，如果進行目標檢測任務，通過Edit->Create Rectangle來切換
• 在標註的過程中，儘量做到精細化的標註，儘量先將目標放大，再執行精細標註，標註的質量會對極大的影響檢測算法的精度
• 該工具默認會生成.json格式的標籤並和標註圖片存放在同一個目錄，.json格式的標籤文件適合用來構建COCO數據集格式，而.xml格式的標籤文件適合用來構建VOC數據集格式，

2.2 LabelImg

Github鏈接

2.2.1 工具安裝

方案1：

pip3 install labelImg
labelImg
labelImg [IMAGE_PATH] [PRE-DEFINED CLASS FILE]

方案2：

git clone https://github.com/tzutalin/labelImg.git
sudo apt-get install pyqt5-dev-tools
sudo pip3 install -r requirements/requirements-linux-python3.txt
make qt5py3
python3 labelImg.py
python3 labelImg.py [IMAGE_PATH] [PRE-DEFINED CLASS FILE]

2.2.2 工具使用簡介

  上圖展示了LabelImg工具的標註界面，工具的使用留給大家自己去熟悉，在這裏我只會強調幾個需要注意的地方，具體的內容如下所示：

• 該工具默認生成的標籤格式爲.xml，便於構造成VOC數據集的格式
• 該工具可以將以前使用的標籤保存起來，並將其設置爲默認標籤，這個單類目標檢測任務中能夠極大的節約你的時間

  總而言之，通過使用上面兩個工具對我們自己的數據集進行標註之後，我們最終會獲得一些.json或者.xml格式的標籤文件。當然你也可以使用其它的工具進行標註，具體的細節請看該鏈接。

3、標籤預處理

  通過標註工具我們已經獲得了我們需要的標籤文件，那麼我們爲什麼還需要標籤預處理操作呢？其實原因是這樣的，在計算機視覺界，你可能都會或多或少的聽說過PASCAL VOC和COCO數據集吧，由於這兩個數據集得到了大量的應用，學者們已經針對這兩個數據集寫了很好有用的代碼，比如如何快速的加載圖片和標籤文件等，那麼，對於一個新手來說，能夠複用大佬們的代碼是一個很明智的選擇，這樣你不僅不需要去寫冗長的加載和處理代碼，而且還可以極大的加快你的開發進程，這就是對生成的標籤進行預處理的目的。簡而言之，所謂的標籤預處理即是將自己的圖片和標籤文件調整爲PASCAL VOC和COCO數據集的格式，這樣我們只需要更換數據集，就可以很塊的完成數據的加載和解析工作。

3.1 PASCAL VOC數據集格式詳解

datasets
|__ VOC2007
    |_ JPEGImages
    |_ Annotations
    	|_ .xml
    |_ ImageSets
    	|_ Main
    		|_ train.txt/val.txt/trainval.txt/test.txt
    |_ SegmentationClass
|__ VOC2012
    |_ JPEGImages
    	|_ .jpg
    |_ Annotations
    |_ ImageSets
    |_ SegmentationClass
|__ ...

1、Annotations文件夾中存放使用標註工具獲得的.xml標籤文件；
2、ImageSets/Main文件夾中存放着訓練集\驗證集\訓練驗證集\測試集圖片的名稱，後續代碼訓練的圖片從這些文件中讀取；
3、JPEGImages文件夾中存放的是原始的訓練和驗證圖片，驗證圖片是隨機的從訓練集中獲得的；
4、SegmentationClass和SegmentationObject用來存放目標分割的信息，由於當前的任務僅僅是一個目標檢測的任務，因而不需要在這兩個文件夾下面存放任何文件；

3.2 構造新的PASCAL VOC數據集

# 切換到自己的想要存放數據集的目錄，xxx代表自己的具體目錄
cd xxx
#########################1、創建數據集包含的相關文件夾#########################
mkdir datasets
cd datasets
mkdir Annotations
mkdir ImageSets
mkdir JPEGImages
mkdir SegmentationClass
mkdir SegmentationObject
cd ImageSets
mkdir Main
#########################2、移動圖片和標籤到相應的文件夾#######################
cd xxx                          # xxx表示你的圖片路徑
mv * JPEGImages                 # 將圖片（.jpg/.png/.bmp等）移動到JPEGImages文件夾中
cd xxx                          # xxx表示.xml文件所在的路徑
mv * Annotations                # 將標籤文件（.xml）移動到Annotations文件夾中

#########################3、生成Main文件夾下面的.txt文件#######################
# coding=utf-8
import os
import random

# 用來劃分訓練集合驗證集的比例
trainval_percent = 0.95
train_percent = 0.85
# 設置輸入的.xml和輸出的.txt文件的路徑
xmlfilepath = './datasets/VOC2007/Annotations'
txtsavepath = '.datasets/VOC2007/ImageSets/Main'
total_xml = os.listdir(xmlfilepath)

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)

# 打開相應的文件，準備存儲內容
ftrainval = open('./datasets/VOC2007/ImageSets/Main/trainval.txt', 'w')
ftest = open('./datasets/VOC2007/ImageSets/Main/test.txt', 'w')
ftrain = open('./datasets/VOC2007/ImageSets/Main/train.txt', 'w')
fval = open('./datasets/VOC2007/ImageSets/Main/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()

  如果你是使用Labelme生成的.json格式的標籤文件，你想要使用VOC格式訓練網絡，那麼你需要使用下面這個小工具將.json格式轉換爲.xml格式

#########################4、對.json格式的標籤文件進行處理#######################
# coding=utf-8
import os
import numpy as np
import codecs
import json
from glob import glob
import cv2
import shutil
from sklearn.model_selection import train_test_split
#1.標籤路徑
labelme_path = "./xxx/"                 # 原始xxx標註數據路徑,需要更換成自己的數據集名稱
saved_path = "./datasets/VOC2007/"      # 保存路徑

#2.創建要求文件夾
if not os.path.exists(saved_path + "Annotations"):
    os.makedirs(saved_path + "Annotations")
if not os.path.exists(saved_path + "JPEGImages/"):
    os.makedirs(saved_path + "JPEGImages/")
if not os.path.exists(saved_path + "ImageSets/Main/"):
    os.makedirs(saved_path + "ImageSets/Main/")
    
#3.獲取待處理文件
files = glob(labelme_path + "*.json")
files = [i.split("/")[-1].split(".json")[0] for i in files]

#4.讀取標註信息並寫入 xml
for json_file_ in files:
    json_filename = labelme_path + json_file_ + ".json"
    json_file = json.load(open(json_filename,"r",encoding="utf-8"))
    height, width, channels = cv2.imread(labelme_path + json_file_ +".jpg").shape
    with codecs.open(saved_path + "Annotations/"+json_file_ + ".xml","w","utf-8") as xml:
        xml.write('<annotation>\n')
        xml.write('\t<folder>' + 'UAV_data' + '</folder>\n')
        xml.write('\t<filename>' + json_file_ + ".jpg" + '</filename>\n')
        xml.write('\t<source>\n')
        xml.write('\t\t<database>The Defect Detection</database>\n')
        xml.write('\t\t<annotation>Defect Detection</annotation>\n')
        xml.write('\t\t<image>flickr</image>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t</source>\n')
        xml.write('\t<owner>\n')
        xml.write('\t\t<flickrid>NULL</flickrid>\n')
        xml.write('\t\t<name>WZZ</name>\n')
        xml.write('\t</owner>\n')
        xml.write('\t<size>\n')
        xml.write('\t\t<width>'+ str(width) + '</width>\n')
        xml.write('\t\t<height>'+ str(height) + '</height>\n')
        xml.write('\t\t<depth>' + str(channels) + '</depth>\n')
        xml.write('\t</size>\n')
        xml.write('\t\t<segmented>0</segmented>\n')
        for multi in json_file["shapes"]:
            points = np.array(multi["points"])
            xmin = min(points[:,0])
            xmax = max(points[:,0])
            ymin = min(points[:,1])
            ymax = max(points[:,1])
            label = multi["label"]
            if xmax <= xmin:
                pass
            elif ymax <= ymin:
                pass
            else:
                xml.write('\t<object>\n')
                xml.write('\t\t<name>'+json_file["shapes"][0]["label"]+'</name>\n')
                xml.write('\t\t<pose>Unspecified</pose>\n')
                xml.write('\t\t<truncated>1</truncated>\n')
                xml.write('\t\t<difficult>0</difficult>\n')
                xml.write('\t\t<bndbox>\n')
                xml.write('\t\t\t<xmin>' + str(xmin) + '</xmin>\n')
                xml.write('\t\t\t<ymin>' + str(ymin) + '</ymin>\n')
                xml.write('\t\t\t<xmax>' + str(xmax) + '</xmax>\n')
                xml.write('\t\t\t<ymax>' + str(ymax) + '</ymax>\n')
                xml.write('\t\t</bndbox>\n')
                xml.write('\t</object>\n')
                print(json_filename,xmin,ymin,xmax,ymax,label)
        xml.write('</annotation>')
        
#5.複製圖片到 VOC2007/JPEGImages/下
image_files = glob(labelme_path + "*.jpg")
print("copy image files to VOC007/JPEGImages/")
for image in image_files:
    shutil.copy(image,saved_path +"JPEGImages/")
    
#6.split files for txt
txtsavepath = saved_path + "ImageSets/Main/"
ftrainval = open(txtsavepath+'/trainval.txt', 'w')
ftest = open(txtsavepath+'/test.txt', 'w')
ftrain = open(txtsavepath+'/train.txt', 'w')
fval = open(txtsavepath+'/val.txt', 'w')
total_files = glob("./VOC2007/Annotations/*.xml")
total_files = [i.split("/")[-1].split(".xml")[0] for i in total_files]
test_filepath = "./test"
for file in total_files:
    ftrainval.write(file + "\n")
#test
for file in os.listdir(test_filepath):
    ftest.write(file.split(".jpg")[0] + "\n")
#split
train_files,val_files = train_test_split(total_files,test_size=0.15,random_state=42)
#train
for file in train_files:
    ftrain.write(file + "\n")
#val
for file in val_files:
    fval.write(file + "\n")

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

3.3 COCO數據集格式詳解

COCO_ROOT
|__ annotations
    |_ instances_valminusminival2014.json
    |_ instances_minival2014.json
    |_ instances_train2014.json
    |_ instances_val2014.json
    |_ ...
|__ train2014
    |_ <im-1-name>.jpg
    |_ ...
    |_ <im-N-name>.jpg
|__ val2014
    |_ <im-1-name>.jpg
    |_ ...
    |_ <im-N-name>.jpg
|__ ...

1、COCO_ROOT表示自己的數據集所在的文件夾的名稱；
2、annotations文件夾用來存放標籤文件，instances_train2014.json和instances_val2014.json分別表示訓練集和驗證集的標籤；
3、train2014和val2014文件夾分別用來存儲訓練集和驗證集的圖片，命名格式爲im-N-name.jpg的格式；

3.4 構造新的COCO數據集

  本文更傾向於生成VOC格式的數據集，如果你需要生成COCO格式的數據，請參考博文1和博文2。

4、搭建SSD運行環境

# 克隆一份SSD代碼到本地
git clone https://github.com/lufficc/SSD.git 
cd SSD
# 創建conda虛擬環境
conda create -n ssd python=3.6
# 激活虛擬環境
conda activate ssd
# 安裝相應的python依賴包
pip install torch torchvision yacs tqdm opencv-python vizer
pip install tensorboardX

# 克隆一份cocoapi到本地
git clone https://github.com/cocodataset/cocoapi.git
# 切換到PythonAPI路徑並安裝
cd cocoapi/PythonAPI
python setup.py build_ext install
# 切換到ext文件夾
cd ext
# 編譯NMS等操作，用來加速網絡的訓練過程
python build.py build_ext develop

5、修改代碼訓練網絡

5.1 代碼架構詳解

configs-該文件夾下面存放的是不同的網絡配置文件；
	mobilenet_v2_ssd320_voc0712.yaml-表示mobilenet backbone的配置文件；
	vgg_ssd512_voc0712.yaml-表示vgg backbone的配置文件；
datasets-該文件夾存放數據文件，具體的結果看上文；
demo-該文件夾存放一些測試圖片；
ext-該文件夾存放CPU和GPU版本的NMS代碼；
figures-該文件夾存放網絡訓練過程中的一些可視化文件；
outputs-該文件夾存放訓練過程中的log文件和模型文件（.pth）；
demo.py-該文件用來測試新的圖片並輸出檢測結果；
train.py-該文件用來訓練網絡；
test.py-該文件用來測試網絡；
ssd-該文件夾中存放ssd算法的實現代碼；
	config-該文件夾存放着默認的配置文件；
	data-該文件夾中存放着基本的數據預處理源文件，包括voc和coco數據集；
	engine-該文件夾中存放着基本的訓練和推理引擎的實現源碼；
	layers-該文件夾中存放着separable_conv的實現源碼；
	modeling-該文件夾中存放着SSD網絡的Backbone和head網絡的源碼；
	solver-該文件中存放着SSD網絡中使用到的優化器的源碼；
	structures-該文件夾中存放着一些Box的help函數的源碼；
	utils-該文件夾中存放着一些訓練和推理SSD算法的小工具的源碼：

5.2 修改網絡配置參數

mv datasets SSD                          # 將構造好的VOC數據集移動到SSD文件夾中
cd configs                               # 切換到配置文件夾下面
vim vgg_ssd512_voc0712.yaml              # 使用vim打開配置文件

注：
1、vgg_ssd512_voc0712.yaml表示使用VGG作爲基準網絡，SSD網絡的輸入大小爲512x512 ，使用VOC2007格式的數據集來訓練新的網絡，大量的實驗結果表明該配置參數下能夠獲得最好的檢測效果。
2、vgg_ssd300_voc0712.yaml表示使用VGG作爲基準網絡，SSD網絡的輸入大小爲300x300 ，使用VOC2007格式的數據集來訓練新的網絡，由於輸入圖片的分辨率由512減小到300，整個網絡的訓練速度得到了提升，但是效果不如vgg_ssd512_voc0712.yaml架構。
3、mobilenet_v2_ssd320_voc0712.yaml表示使用Mobilenet_v2作爲基準網絡，SSD網絡的輸入大小爲300x300 ，使用VOC2007格式的數據集來訓練新的網絡，由於使用Mobilenet_v2作爲基準網絡，因此該架構下面的訓練速度最快，但是最終獲得精度也是最差的，這個架構適合於一些簡單的檢測任務，Mobilenet_v2不能能夠獲得較好的檢測效果，而且可以獲得接近實時的推理速度。

###############################修改前的網絡配置###############################
MODEL:
  NUM_CLASSES: 21
  BACKBONE:
    OUT_CHANNELS: (512, 1024, 512, 256, 256, 256, 256)
  PRIORS:
    FEATURE_MAPS: [64, 32, 16, 8, 4, 2, 1]
    STRIDES: [8, 16, 32, 64, 128, 256, 512]
    MIN_SIZES: [35.84, 76.8, 153.6, 230.4, 307.2, 384.0, 460.8]
    MAX_SIZES: [76.8, 153.6, 230.4, 307.2, 384.0, 460.8, 537.65]
    ASPECT_RATIOS: [[2], [2, 3], [2, 3], [2, 3], [2, 3], [2], [2]]
    BOXES_PER_LOCATION: [4, 6, 6, 6, 6, 4, 4]
INPUT:
  IMAGE_SIZE: 512
DATASETS:
  TRAIN: ("voc_2007_trainval", "voc_2012_trainval")
  TEST: ("voc_2007_test", )
SOLVER:
  MAX_ITER: 120000
  LR_STEPS: [80000, 100000]
  GAMMA: 0.1
  BATCH_SIZE: 24
  LR: 1e-3

OUTPUT_DIR: 'outputs/vgg_ssd512_voc0712'

###############################修改後網絡配置###############################
MODEL:
  NUM_CLASSES: 18
  BACKBONE:
    OUT_CHANNELS: (512, 1024, 512, 256, 256, 256, 256)
  PRIORS:
    FEATURE_MAPS: [64, 32, 16, 8, 4, 2, 1]
    STRIDES: [8, 16, 32, 64, 128, 256, 512]
    MIN_SIZES: [35.84, 76.8, 153.6, 230.4, 307.2, 384.0, 460.8]
    MAX_SIZES: [76.8, 153.6, 230.4, 307.2, 384.0, 460.8, 537.65]
    ASPECT_RATIOS: [[2], [2, 3], [2, 3], [2, 3], [2, 3], [2], [2]]
    BOXES_PER_LOCATION: [4, 6, 6, 6, 6, 4, 4]
INPUT:
  IMAGE_SIZE: 512
DATASETS:
  TRAIN: ("voc_2007_trainval",)
  TEST: ("voc_2007_test", )

SOLVER:
  MAX_ITER: 120000
  LR_STEPS: [80000, 100000]
  GAMMA: 0.1
  BATCH_SIZE: 12
  LR: 1e-4

OUTPUT_DIR: 'outputs/vgg_ssd512_voc0712'

注：
1、NUM_CLASSES參數表示檢測的類型，由於當前的缺陷數據集的類別數目爲17，另外加上一個__background__類，因此需要將該數值修改爲17+1=18;
2、BATCH_SIZE參數表示每次訓練的圖片的個數，由於訓練的圖片需要加載到內存中，該數值和你當前使用的顯卡的內容有着密切的關係，我根據自己的顯卡（Tesla T4）將BATCH_SIZE調整爲12；
3、LR參數表示網絡訓練時的學習率，由於我設置的迭代次數比較多，害怕網絡跳過最優值，因而我將LR調整爲1e-4；
4、TRAIN參數表示網絡訓練時使用的數據集，由於我們新建了一個類似於VOC2007的數據集，因而我刪除了voc_2012_trainval；
5、其它的這些配置參數，你可以根據自己的需要進行修改，但是我建議你不用修改。

5.3 修改VOC類別參數

cd .ssd/data/datasets                # 1、切換到datasets文件夾
vim voc.py                           # 2、打開voc文件       

###############################修改前的類別###############################
class_names = ('__background__',
               'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair',
               'cow', 'diningtable', 'dog', 'horse',
               'motorbike', 'person', 'pottedplant',
               'sheep', 'sofa', 'train', 'tvmonitor')

###############################修改後的類別###############################
 class_names = ('_background_', 'Buttercup', 'Colts Foot', 'Daffodil', 
 				'Daisy', 'Dandelion', 'Fritillary', 'Iris', 'Pansy', 
 				'Sunflower', 'Windflower', 'Snowdrop', 'LilyValley',
 				'Bluebell', 'Crocus', 'Tigerlily', 'Tulip', 'Cowslip')

5.4 下載模型

# 下載vgg_ssd300_voc0712.pth預訓練模型
wget https://github.com/lufficc/SSD/releases/download/1.2/vgg_ssd300_voc0712.pth
# 下載vgg_ssd512_voc0712.pth預訓練模型
wget https://github.com/lufficc/SSD/releases/download/1.2/vgg_ssd512_voc0712.pth
# 下載mobilenet_v2_ssd320_voc0712.pth預訓練模型
wget https://github.com/lufficc/SSD/releases/download/1.2/mobilenet_v2_ssd320_voc0712.pth
# 創建.torch文件夾用來存放模型
mkdir ~/.torch/
# 將預訓練好的模型移動到這個文件夾中
mv vgg_ssd300_voc0712.pth ~/.torch
mv vgg_ssd512_voc0712.pth ~/.torch
mv mobilenet_v2_ssd320_voc0712.pth ~/.torch

5.5 訓練模型

# 導入環境變量VOC_ROOT
export VOC_ROOT="./datasets"
# 使用vim打開train.py文件
vim train.py

###############################修改前網絡訓練參數###############################
parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training With PyTorch')
# 設置使用的配置文件
parser.add_argument( "--config-file", default="", metavar="FILE", help="path to config file",type=str)
# 設置是否進行local_rank
parser.add_argument("--local_rank", type=int, default=0)
# 設置保存log的步長
parser.add_argument('--log_step', default=10, type=int, help='Print logs every log_step')
# 設置保存checkpoint文件的步長
parser.add_argument('--save_step', default=2500, type=int, help='Save checkpoint every save_step')
# 設置執行網絡評估操作的步長
parser.add_argument('--eval_step', default=2500, type=int, help='Evaluate dataset every eval_step, disabled when eval_step < 0')
# 設置是否使用tensorboard進行loss和accuracy的可視化
parser.add_argument('--use_tensorboard', default=True, type=str2bool)
parser.add_argument("--skip-test", dest="skip_test", help="Do not test the final model", action="store_true")
# 設置是否打開使用命令行來改變參數
parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER)
args = parser.parse_args()

###############################修改後網絡訓練參數###############################
parser = argparse.ArgumentParser(description='Single Shot MultiBox Detector Training With PyTorch')
parser.add_argument("--config-file", default="configs/vgg_ssd512_voc0712.yaml", metavar="FILE", help="path to config file", type=str)
parser.add_argument("--local_rank", type=int, default=0)
parser.add_argument('--log_step', default=20, type=int, help='Print logs every log_step')
parser.add_argument('--save_step', default=2500, type=int, help='Save checkpoint every save_step')
parser.add_argument('--eval_step', default=2500, type=int, help='Evaluate dataset every eval_step, disabled when eval_step < 0')
parser.add_argument('--use_tensorboard', default=True, type=str2bool)
parser.add_argument("--skip-test", dest="skip_test", help="Do not test the final model", action="store_true")
parser.add_argument("opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER)
args = parser.parse_args()

注：
1、爲了便於後續的代碼運行，這裏直接將–config-file的默認參數設置爲vgg_ssd512_voc0712.yaml，即使用VGG作爲基準網路，輸入大小爲512x512，使用自己創建的VOC2007格式的數據集訓練新的模型。
2、對於其它的一些超參數，不建議你去改動。

# 如果你有一個GPU，執行該指令訓練模型
python train.py
# 如果你有4個GPU，執行該指令訓練模型
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS train.py

注：
1、下面展示是網絡訓練過程中輸出的配置參數。

2、下面展示的是網絡訓練過程中的Loss、Accuracy和運行時間等輸出。

5.6 Loss和Accuracy可視化

# 切換到相應的文件夾
cd ./SSD/outputs/vgg_ssd512_voc0712/
# 使用tensorboard打開網絡訓練過程中保存的log文件
tensorboard --logdir=tf_logs
# 在谷歌瀏覽器中輸入該網址查看logs
http://ubuntu:6006 

注：
1、Tensorboard展示的logs文件如下圖所示，cls_loss表示分類分支的loss曲線，橫軸表示的是網絡的迭代次數，縱軸表示cls_loss值，該數值越小表示網絡預測的結果和標籤的結果越接近；reg_loss表示迴歸分支的loss曲線，橫軸表示的是網絡的迭代次數，縱軸表示reg_loss值，該數值越小表示網絡預測的結果和標籤的結果越接近；total_loss是reg_loss和cls_loss之和，如果這3條曲線的整體趨勢是隨着迭代次數的增加呈現下降趨勢就表示我們的網絡正在不斷的收斂，不斷的尋找局部最優解。
2、LR表示網絡的學習率，該參數一般設置爲1e-3或者1e-4，隨着迭代次數的增加應該不斷的減小該數值的值以防跳躍最優解，這也就是所謂的衰減率。當在訓練的過程中出現NAN或者NULL時可以嘗試着調節該參數，或許可以解決該問題。

5.7 模型測試和效果分析

###############################使用test.txt下的圖片做測試############################
# 切換到改文件夾
cd SSD/outputs/vgg_ssd512_voc0712/
# 將訓練好的模型複製到model文件夾中
cp model_final.pth ../../model/
# 如果你有一個GPU，執行該指令訓練模型
python test.py --config-file configs/vgg_ssd300_voc0712.yaml
# 如果你有4個GPU，執行該指令訓練模型
export NGPUS=4
python -m torch.distributed.launch --nproc_per_node=$NGPUS test.py --config-file configs/vgg_ssd300_voc0712.yaml

###############################使用特定文件夾下的圖片做測試###########################
# 切換到改文件夾
cd SSD/outputs/vgg_ssd512_voc0712/
# 將訓練好的模型複製到model文件夾中
cp model_final.pth ../../model/
# 創建測試文件夾
mkdir test_imgs
# 將自己的測試圖片拷貝到該文件夾下面，這裏*.jpg只是用來進行演示
cp *jpg test_imgs
# 使用SSD對test_imgs文件夾下面的圖片進行預測
python demo.py --config-file configs/vgg_ssd300_voc0712.yaml --images_dir test_imgs

注：
1、下圖展示了運行python test.py之後的結果，這個輸出中主要包含測試圖片的數量和訓練好的模型在測試集上面的mAP指標值，該指標可以用來評估檢測模型的好壞，具體的細節請參考該博客。

2、下圖展示了運行python demo.py之後的結果，這個輸出中不僅包含着測試圖片的個數1768，同時包含着裝載每一張圖片所需要的時間，以及對每一張圖片進行預測所花費的時間，除此之外，你還可以看到當前的幀率值。

6、SSD性能改進

  如果直接使用現成的SSD訓練的結果可能並不能達到你的預期精度，那麼你可以從以下的幾個地方快速的對該模型進行改進，以下的幾個思路只是我自己的思路，希望大佬們不要吐槽。

• 思路1-該SSD網絡在訓練過程中已經包含了數據增強操作，具體的代碼在./SSD/ ssd/data/ transforms/init.py下面，但是對於多類別檢測任務而言，經常會存在類別不均衡問題，爲了解決該問題，最笨的方法就是使用數據增強工具將不同類別的數目調整爲相同數量，推薦使用Augmentor工具，具體的使用代碼如下所示：

# coding=utf-8
import Augmentor
# 1. 指定圖片所在目錄
p = Augmentor.Pipeline("./chahua_imgs/")
p.ground_truth("./chahua_labels")
# 2. 增強操作
# 放大 概率0.3，最小爲1.1倍，最大爲1.6倍；1不做變換
p.zoom(probability=0.5, min_factor=0.5, max_factor=1.6)
p.skew_tilt(probability=0.5, magnitude=1)
p.skew_corner(probability=0.5, magnitude=1)
p.random_distortion(probability=0.5, grid_width=5, grid_height=5, magnitude=1)
p.rotate_random_90(probability=0.5)
p.shear(probability=0.5, max_shear_left=10, max_shear_right=10)
p.flip_random(probability=0.5)

# 3. 指定增強後圖片數目總量
p.sample(1000)

• 思路2-如果你熟悉目標檢測算法，可能知道Focal_loss可以很好的抑制樣本不均衡問題，因而一個改進的思路就是利用Focal_loss替換掉原始的smooth_l1損失，具體的實現可以參考Focal_loss實現。
• 思路3-對於目標檢測問題而言，Loss函數的好壞直接或者間接會影響到算法的性能，針對這個問題當前已經有很好改進的Loss函數，比較具有代表性的就是IoU_Loss和IoU的變種GIoU和DIoU等，因而你可以嘗試着使用新的Loss來訓練SSD網絡，具體的實現可以參考DIoU實現。
• 思路4-NMS是目標檢測中必不可少的一步， 當前已經有很好NMS的變種算法，你或許僅僅需要修改幾行代碼就可以提升整個模型的性能，因而你可以嘗試着使用Soft-NMS等其它的變種去替換掉原始的NMS操作,具體的實現可以參考Soft-NMS實現。

7、Pytorch模型加速

  當你訓練完自己的SSD檢測模型之後，如果你使用MobileNet_v2的Backbone之後，仍然感覺到整個網絡的運行速度不能夠滿足你的需求，那麼你可能就需要做模型加速啦，模型加速的主要方式包括模型蒸餾、模型裁剪和模型量化，當前大量使用的技術是模型量化，如果你需要對Pytorch訓練好的模型做模型量化，建議你參考官網量化1、官網量化2。除了官方的量化工具外，你還可以使用這個Pytorch模型優化工具。具體的量化工作就交給大家自己去實現了，只有自己動手去做了才能真正的學懂知識，哈哈。

三、總結

  本文的主要目的是教你如何一步步使用現有SSD代碼在自己的數據集上面訓練出一個可用的檢測網絡。主要的步驟包括數據標註、標籤預處理、配置運行環境、修改模型參數、修改代碼參數、訓練模型、測試模型、優化模型和部署模型等。本文教給你的不僅僅訓練一個SSD網絡，當你理解了其中的道理，訓練一個Yolov3或者其它的檢測算法也就變成一件比較簡單的事情了。

參考資料

1、SSD代碼
2、Pytorch模型優化工具
3、官網量化1
4、官網量化2
5、Focal_loss實現
6、DIoU實現
7、Soft-NMS實現
8、Augmentor工具

注意事項

[1] 該博客是本人原創博客，如果您對該博客感興趣，想要轉載該博客，請與我聯繫（qq郵箱：[email protected]）,我會在第一時間回覆大家，謝謝大家的關注.
[2] 由於個人能力有限，該博客可能存在很多的問題，希望大家能夠提出改進意見。
[3] 如果您在閱讀本博客時遇到不理解的地方，希望您可以聯繫我，我會及時的回覆您，和您交流想法和意見，謝謝。
[4] 本文測試的圖片可以通過網盤鏈接進行下載。提取碼：hku5。
[5] 本人業餘時間承接各種本科畢設設計和各種小項目，包括圖像處理（數據挖掘、機器學習、深度學習等）、matlab仿真、python算法及仿真等，有需要的請加QQ：1575262785詳聊，備註“項目”！！！