ubuntu14.04中caffe下使用自己的数据集训练FCN

最近应项目需要训练一个FCN网络,原始FCN网络的类别数为21种,而我使用的数据集类别数为2(其中包含了背景),在网上看了一些关于FCN训练的帖子感觉还是没有一个系统清晰的指南,所以将自己的实现过程写下来分享给大家.

---

首先给出FCN源码下载地址:https://github.com/shelhamer/fcn.berkeleyvision.org

关于finetune的过程,主要有以下几个步骤:
1.整理自己的数据集,修改或重写数据输入层
2.修改.prototxt文件
3.修改solve.py文件
4.进行finetune

---

一. 数据集相关

原始的VOC数据集和SBDD数据集输入是通过VOCSegDataLayer和SBDDSegDataLayer两个class定义的,在voc_layers.py中,建议大家可以看看这两个python layers定义的方式,然后在根据自己的数据集决定是要修改还是重写.

VOC数据集的语义分割部分是将原图存放在VOCdevkit/VOC2012/JPEGImages,将每个原图对应的分割图存放在VOCdevkit/VOC2012/SegmentationClass,然后具体的train/val/test分组是存放在VOCdevkit/VOC2012/ImageSets/Segmentation中,每一个txt中就存放这对应集合的样本名称.在数据输入层VOCSegDataLayer中就是通过解析这些txt文件得到图像名称进行输入的.

这里由于我使用的数据集是按照不同train,val,test的文件夹存放的,并没有txt文件作为索引集合,所以我重写的一个数据输入层,用于从文件夹中读取样本,然后送入到网络中.这里只给出setup函数,其他的reshape,forward,backword函数跟VOC数据集是一样的.

class AerialImageDatasetLayer(caffe.Layer):
    def setup(self, bottom, top):
        #get the config from the .prototxt file
        params=eval(self.param_str)
        self.dataset_dir = params['dataset_dir']
        self.seed = params.get('seed',None)
        self.split = params['split']
        self.mean = params.get('mean',0)
        self.random = params.get('random',True)

        #check the number of top and bottom
        if len(top) is not 2:
            raise Exception('Need 2 output for this layer!')