用自己的資料集訓練tf-ssd模型
資料集製作
因為老闆接的豐田的一個專案,工廠那邊要求能識別出雨天打傘的行人、交通錐形桶、躺在地上的人等,PASCAL VOC的資料集類別裡沒這些,是滿足不了他們要求了,所以要去製作資料集訓練網路。我們去現場拍了些視訊,然後用我之前寫的一個指令碼解析,得到原始圖片,拿給實驗室的學弟學妹們標註。標註工具也是github上一個開源專案,作者搞了個GUI,生成帶標籤的.xml檔案,這裡也想感謝他。https://github.com/tzutalin/labelImg
標註完了之後得到大量的.xml檔案,然後我們需要將這些資料做成VOC的一樣格式。先來看看VOC的訓練集是什麼樣的:
因為只是做檢測,生成Annotations、ImageSets、JPEGImages這三個資料夾就可以了。後面的兩個segmentation還有一個test,這三個夾子不用管。Annotations是用來放.xml檔案的,JPEGImages放原始的jpg圖片,像這樣:
ImageSets這個夾子開啟長這樣:
我們只要生成Main資料夾就可以了,這個資料夾是用來存放資料對應的.txt檔案的。開啟Main:
其他的都不重要,生成這三個畫圈的檔案加一個test.txt就行了。利用程式生成test.txt, train.txt, trainval.txt, val.txt, 程式碼如下:
import os import random xmlfilepath=r'/home/ogai/ngy/nissd/mydataset/Annotations' saveBasePath=r"/home/ogai/ngy/nissd/" trainval_percent=0.8 train_percent=0.7 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) print("train and val size",tv) print("traub suze",tr) ftrainval = open(os.path.join(saveBasePath,'mydataset/ImageSets/Main/trainval.txt'), 'w') ftest = open(os.path.join(saveBasePath,'mydataset/ImageSets/Main/test.txt'), 'w') ftrain = open(os.path.join(saveBasePath,'mydataset/ImageSets/Main/train.txt'), 'w') fval = open(os.path.join(saveBasePath,'mydataset/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()
到此為止資料集製作完畢。
生成tfrecords
我們需要先修改一下/datasets/pascalvoc_common.py中的類別標籤定義,我的做法:
""" VOC_LABELS = { 'none': (0, 'Background'), 'aeroplane': (1, 'Vehicle'), 'bicycle': (2, 'Vehicle'), 'bird': (3, 'Animal'), 'boat': (4, 'Vehicle'), 'bottle': (5, 'Indoor'), 'bus': (6, 'Vehicle'), 'car': (7, 'Vehicle'), 'cat': (8, 'Animal'), 'chair': (9, 'Indoor'), 'cow': (10, 'Animal'), 'diningtable': (11, 'Indoor'), 'dog': (12, 'Animal'), 'horse': (13, 'Animal'), 'motorbike': (14, 'Vehicle'), 'Person': (15, 'Person'), 'pottedplant': (16, 'Indoor'), 'sheep': (17, 'Animal'), 'sofa': (18, 'Indoor'), 'train': (19, 'Vehicle'), 'tvmonitor': (20, 'Indoor'), } """ VOC_LABELS = { 'none': (0, 'Background'), 'cone': (1, 'Cone'), 'umbrellaman': (2, 'Umbrella Man'), # 類別我就不一一列舉了 }
總之你的資料有多少類就換成多少類。
寫個shell指令碼,利用tf_convert_data.py來生成tfrecords:#!/bin/bash
DATASET_DIR=./mydataset/
OUTPUT_DIR=./tfrecords
python tf_convert_data.py \
--dataset_name=pascalvoc \
--dataset_dir=${DATASET_DIR} \
--output_name=voc_2007_train \
--output_dir=${OUTPUT_DIR}
但是呢,我這裡報錯了:UnicodeDecodeError: 'utf-8' codec can't decode byte 0xff in position 0: invalid start byte
解決方法,修改pascalvoc_to_tfrecords.py的第83行,將'r'改成'rb'
image_data = tf.gfile.FastGFile(filename, 'rb').read()
Fine-tuning訓練
在預訓練好的ssd模型(vgg300)上訓練自己的資料,然後寫個shell指令碼,利用train_ssd_network.py來訓練,注意路徑:
#!/bin/bash
DATASET_DIR=./tfrecords
TRAIN_DIR=./logs/
CHECKPOINT_PATH=./checkpoints/ssd_300_vgg.ckpt
python train_ssd_network.py \
--train_dir=${TRAIN_DIR} \
--dataset_dir=${DATASET_DIR} \
--dataset_name=pascalvoc_2007 \
--dataset_split_name=train \
--model_name=ssd_300_vgg \
--checkpoint_path=${CHECKPOINT_PATH} \
--save_summaries_secs=60 \
--save_interval_secs=600 \
--weight_decay=0.0005 \
--optimizer=adam \
--learning_rate=0.001 \
--batch_size=16
TRAIN_DIR是自己訓練產生的模型的checkpoints目錄,待會demo要從這裡載入checkpoints。
我訓練了一天多一點,機器是1070ti(6g),然而loss總是在1~5之間振盪,無法收斂,我也不知道為什麼。後來我直接關閉了訓練,用最新的checkpoints載入測試,奇怪的是竟然發現結果還行......
Demo
用jupyter notebook開啟/notebooks/ssd_notebook.ipynb,然後需要修改一下checkpoints載入路徑,替換成自己訓練的結果。
改下類別標籤:
網上找了點圖片,測測結果:
Evaluation
寫個shell指令碼,利用eval_ssd_network.py來evaluate:
#!/bin/bash
EVAL_DIR=./logs/
DATASET_DIR=./tfrecords/
CHECKPOINT_PATH=./logs/model.ckpt-252448
python eval_ssd_network.py \
--eval_dir=${EVAL_DIR} \
--dataset_dir=${DATASET_DIR} \
--dataset_name=pascalvoc_2007 \
--dataset_split_name=train \
--model_name=ssd_300_vgg \
--checkpoint_path=${CHECKPOINT_PATH} \
--batch_size=1
EVAL_DIR是eval結果儲存的路徑,DATASET_DIR是用來eval的資料集的路徑。
但是報錯了:TypeError: Can not convert a tuple into a Tensor or Operation
解決方法,在eval_ssd_network.py裡定義一個flatten函式:
def flatten(x):
result = []
for el in x:
if isinstance(el, tuple):
result.extend(flatten(el))
else:
result.append(el)
return result
然後將原本第318和338行的
eval_op = list(names_to_updates.values()),
改為
eval_op = flatten(list(names_to_updates.values())),