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。