Fast-rcnn 训练（3）- 训练自己的数据集

训练文件的入口就是tools文件夹下的train_net.py

作者给的训练方式是

./tools/train_net.py --gpu 0 --solver models/VGG16/solver.prototxt \
--weights data/imagenet_models/VGG16.v2.caffemodel

这里训练的话默认是训练VOC的数据，所以需要先把VOC的数据放在指定的文件夹下。其实这里还有一个隐藏的参数

--imdb voc_2007_trainval

imdb后面跟的就是你要训练的数据集的名字。

所以后面要定义自己的数据集。

–weights后面跟的就是你要训练的模型。

这里如果要训练一个空模型的话，需要先从data/scripts文件夹下面运行fetch_imagenet_models.sh，下载空的模型。

首先看一下/lib/datasets/文件夹里面，有一个__init__.py的文件

最后有一段代码：

if _which(MATLAB) is None:
msg = ("MATLAB command '{}' not found. "
"Please add '{}' to your PATH.").format(MATLAB, MATLAB)
raise EnvironmentError(msg)

就是为了先检查有没有matlab的，没有的话会报错，这里因为是训练自己的数据集，后面matlab的函数是不需要的，所以把这段代码删掉。

剩下factory.py，pascol_voc.py和imdb.py三个重要的文件

factory.py

def _selective_search_IJCV_top_k(split, year, top_k):
"""Return an imdb that uses the top k proposals from the selective search
IJCV code.
"""
imdb = datasets.pascal_voc(split, year)
imdb.roidb_handler = imdb.selective_search_IJCV_roidb
imdb.config['top_k'] = top_k
return imdb

# Set up voc_<year>_<split> using selective search "fast" mode
for year in ['2007', '2012']:
for split in ['train', 'val', 'trainval', 'test']:
name = 'voc_{}_{}'.format(year, split)
__sets[name] = (lambda split=split, year=year:
datasets.pascal_voc(split, year))

# Set up voc_<year>_<split>_top_<k> using selective search "quality" mode
# but only returning the first k boxes
for top_k in np.arange(1000, 11000, 1000):
for year in ['2007', '2012']:
for split in ['train', 'val', 'trainval', 'test']:
name = 'voc_{}_{}_top_{:d}'.format(year, split, top_k)
__sets[name] = (lambda split=split, year=year, top_k=top_k:
_selective_search_IJCV_top_k(split, year, top_k))

代码中前面的这些，如果是用来训练自己的数据集，那么就都可以删掉不用了。只要模仿上面的格式写一个自己的数据集的类，添加到__sets集合中就行了。

所以在这里，我先提前照着pascal_voc.py,写了一个自己数据集的类，叫Rjmgc.py（后面会给出代码）

首先，我的数据集位置放在了$FRCN_ROOT/Rjmgc_data这个目录下

所以factory.py里只需要添加这些话就好了

from datasets.Rjmgc import Rjmgc
Rjmgc_devkit_path = '/home/hh/fast-rcnn/Rjmgc_data'
__sets['Rjmgc_train'] = (lambda imageset = 'train',devkit = Rjmgc_devkit_path:Rjmgc(imageset,devkit))

先说一下我的训练数据的存放格式

首先，数据集的地址是/fast-rcnn/Rjmgc_data

训练用的图片放在了/fast-rcnn/Rjmgc_data/data/Images

训练图片的文件列表“train.txt”放在了/fast-rcnn/Rjmgc_data/data/ImageSets

标记xml文件放在了/fast-rcnn/Rjmgc_data/data/Annotations

然后需要一个记录了selective_search的mat文件，放在了/fast-rcnn/Rjmgc_data下

Rjmgc.py

首先是init方法，因为只需要car和bus，就只保留了两个类

def __init__(self,image_set,devkit_path):
datasets.imdb.__init__(self,image_set)
self._image_set = image_set
self._devkit_path = devkit_path
self._data_path = os.path.join(self._devkit_path,'data')
self._classes = ('__background__','car','bus')
self._class_to_ind = dict(zip(self.classes, xrange(self.num_classes)))
self._image_ext = '.jpg'
self._image_index = self._load_image_set_index()
self._roidb_handler = self.selective_search_roidb

self.config = {'cleanup': True,
'use_salt': True,
'top_k': 2000}

assert os.path.exists(self._devkit_path), \
'Devkit path does not exist: {}'.format(self._devkit_path)
assert os.path.exists(self._data_path), \
'Path does not exist: {}'.format(self._data_path)

修改函数 _load_selective_search_roidb为

def _load_selective_search_roidb(self, gt_roidb):
filename = os.path.abspath(os.path.join(self._devkit_path,
"train.mat"))
assert os.path.exists(filename), \
'Selective search data not found at: {}'.format(filename)
raw_data = sio.loadmat(filename)['all_boxes'].ravel()

box_list = []
for i in xrange(raw_data.shape[0]):
box_list.append(raw_data[i][:, (1, 0, 3, 2)])

return self.create_roidb_from_box_list(box_list, gt_roidb)

其实就是根据自己的mat文件的存放位置，修改了一下读取数据的格式。

后面的两个函数其实不需要了

selective_search_IJCV_roidb

_load_selective_search_IJCV_roidb

最后把main函数也统一的修改了

if __name__ == '__main__':
d = datasets.Rjmgc('train', '')
res = d.roidb
from IPython import embed; embed()

imdb.py

不需要修改

修改模型文件的配置，这里主要参考http://www.cnblogs.com/louyihang-loves-baiyan/p/4906690.html

比如这里需要使用caffenet网络进行训练

所以要到$FRCN_ROOT/models/CaffeNet文件夹里，找到train.prototxt文件进行修改

首先在data层把num_classes 从原来的21类( 20类+背景) ，改成 3类 (car+bus+背景)

接在在cls_score层把num_output 从原来的21 改成 3

在bbox_pred层把num_output 从原来的84 改成12， 为检测类别个数乘以4，比如这里是3类那就是3*4=12

这样前期的操作就算是完成了，就可以开始训练了

在我这里，为了方便，写了一个脚本文件

cd /home/hh/fast-rcnn/data
rm -r cache
cd /home/hh/fast-rcnn/tools
./train_net.py --solver /home/hh/fast-rcnn/models/CaffeNet/solver.prototxt --weights /home/hh/fast-rcnn/data/imagenet_models/Rjmgc_empty.caffemodel --imdb Rjmgc_train

然后就开始运行了。



我这里默认是迭代40000次。

训练后生成的模型文件默认会保存在$FRCN_ROOT/output/default/train文件夹下面，并且每10000次迭代都会生成一个



最后如果想看一下训练的模型的效果的话，可以使用demo.py查看效果

代码中

caffemodel = os.path.join(cfg.ROOT_DIR, 'data', 'fast_rcnn_models',
NETS[args.demo_net][1])

指定了测试用的模型，所以只需要把模型指定为刚训练的模型即可

caffemodel = '/home/hh/fast-rcnn/output/default/train/caffenet_fast_rcnn_iter_40000.caffemodel'

后面修改一下要测试的图片名，就可以直接运行demo了