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__ 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___top_ 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

然后就开始运行了。
Fast-rcnn 训练(3)- 训练自己的数据集_第1张图片
我这里默认是迭代40000次。
训练后生成的模型文件默认会保存在$FRCN_ROOT/output/default/train文件夹下面,并且每10000次迭代都会生成一个
Fast-rcnn 训练(3)- 训练自己的数据集_第2张图片


最后如果想看一下训练的模型的效果的话,可以使用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了

你可能感兴趣的:(code)