SSD學習系列（三）基于自己資料集訓練SSD模型

LMDB資料形式

将自己的資料集轉換成LMDB格式後就可以基于它開始訓練了。其對應檔案目錄如下所示。

/work/xxx/caffe_ssd/data/VOCdevkit/VOC201909/lmdb/

VOC201909_test_lmdb  VOC201909_trainval_lmdb
           

編輯訓練腳本 ssd_pascal.py

sudo vi examples/ssd/ssd_pascal.py
           

resize_width和height代表網絡size，可以設為500x500來提高點識别準确度，但是訓練時相應的要求更大的顯存。

如果是3個GPU，則上圖所示的 gpus = "0,1,2"。其它依此類推。 batch size如果顯存足夠大可以适當增加些來提高訓練速度。

test image數目可以根據資料集情況進行修改。 我這裡隻有12張測試圖檔，每次就測一張 。

訓練遇到的問題

在真正開始訓練前，可以考慮編輯solver.prototxt中的幾個參數，如下紅框所示。

訓練指令如下

sudo python xxx/caffe_ssd/examples/ssd/ssd_pascal.py
           

問題1：不識别caffe 

Traceback (most recent call last):  
File "examples/ssd/ssd_pascal.py", line 2, in <module>    
import caffeModuleNotFoundError: 
No module named 'caffe'
           

解決辦法依然是在ssd_pascal.py中插入兩行：

#在import caffe之前插入這兩行
import sys
sys.path.append('/work/tieshan/caffe_ssd/python')
           

問題2：顯存不夠

I0918 17:37:18.623419  5384 blocking_queue.cpp:50] Data layer prefetch queue empty
F0918 17:37:20.669662  5384 syncedmem.cpp:56] Check failed: 
error == cudaSuccess (2 vs. 0)  out of memory
           

我看了一下，如果網絡size為300x300所需顯存要9個多GB，其batch size等于32，剛好可以在1080ti上訓練。如果是500x500的size，則要改batch size為8才能跑的起來。不過使用500x500或512x512，loss值很容易為nan值，有待進一步check。