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!')
        if len(bottom) is not 0:
            raise Exception('Need no input for this layer!')

        #get the image names in dataset_dir/split/images
        self.splitimages_dir = os.path.join(self.dataset_dir, self.split, 'images')
        self.splitlabels_dir = os.path.join(self.dataset_dir, self.split, 'gt')
        self.splitimages_names = list(os.listdir(self.splitimages_dir))
        self.indx = 0

        #decide whather need to random the images order
        if 'train' not in self.split:
            self.random = False

        if self.random is True:
            random.seed(self.seed)
            self.indx = random.randint(0, len(self.splitimages_names)-1)

二. 修改.prototxt

首先再次说明一下我的类别数是2,也就是二分类问题,使用的是voc-fcn32s(所以所提到的相关文件都是在voc-fcn32s文件夹中的)这种只进行了一次upsample的结构,基于vgg16-fcn.caffemodel进行finetune.

主要修改的prototxt文件有train.prototxt,val.prototxt和deploy.prototxt三个文件:
1.train.prototxt
有两个改的方面,一个是数据输入层的参数,另一个是后面score_fr和upscore层:
将data层改为:

layer {
  name: "data"
  type: "Python"
  top: "data"
  top: "label"
  python_param {
    module: "AerialImageDataset"
    layer: "AerialImageDatasetLayer"
    param_str: "{\'dataset_dir\': \'./data/dataset\', \'seed\': 1337, \'split\': \'train\', \'mean\': (104.00699, 116.66877, 122.67892)}"
  }
}

其中./data/dataset要修改为你存放数据集的路径.
将score_fr和upscore层改为:

layer {
  name: "score_fr"
  type: "Convolution"
  bottom: "fc7"
  top: "score_fr"
  param {
    lr_mult: 1
    decay_mult: 1
  }
  param {
    lr_mult: 2
    decay_mult: 0
  }
  convolution_param {
    num_output: 2
    pad: 0
    kernel_size: 1
  }
}
layer {
  name: "upscore"
  type: "Deconvolution"
  bottom: "score_fr"
  top: "upscore"
  param {
    lr_mult: 0
  }
  convolution_param {
    num_output: 2
    bias_term: false
    kernel_size: 64
    stride: 32
  }
}

其实主要是修改了两个层的num_output参数,这个参数就是你的类别数.

2.val.prototxt
val结构声明中类似于train.prototxt修改方法,也是修改那两个部分,只不过要将data层中的split参数改为val
3.deploy.prototxt
只需要类似train.prototxt修改score_fr和upscore层.

三. 修改solve.py

solve.py作用是初始化device,初始化网络,copy源网络的权重参数以及进行训练和验证,相当于是训练网络的main函数.
先直观的给出修改后的solve.py:

#!/usr/bin/env python
import caffe
import surgery, score

import numpy as np
import os
import sys

try:
    import setproctitle
    setproctitle.setproctitle(os.path.basename(os.getcwd()))
except:
    pass

#weights = '../voc-fcn8s/fcn8s-heavy-pascal.caffemodel'
weights = '../ilsvrc-nets/vgg16-fcn.caffemodel'
vgg_proto = '../ilsvrc-nets/VGG_ILSVRC_16_layers_deploy.prototxt'

# init
#caffe.set_device(int(sys.argv[1]))
caffe.set_device(0)
caffe.set_mode_gpu()

solver = caffe.SGDSolver('solver.prototxt')
#solver.net.copy_from(weights)
vgg_net=caffe.Net(vgg_proto,weights,caffe.TRAIN)
surgery.transplant(solver.net,vgg_net)
del vgg_net

# surgeries
interp_layers = [k for k in solver.net.params.keys() if 'up' in k]
surgery.interp(solver.net, interp_layers)

# scoring
# val = np.loadtxt('../data/segvalid11.txt', dtype=str)
val = os.listdir('/home/irsa/fcn.berkeleyvision.org/voc-fcn8s/data/dataset/val/images')
for _ in range(25):
    solver.step(4000)
    score.seg_tests(solver, False, val, layer='score')

可以看出有以下几个方面的修改:
(1) 添加了#!/usr/bin/env python,这个的作用是可以直接在终端中输入./solve.py来运行py程序(同时也是方便后续.sh文件的调用).
(2) 修改了weights的定义,添加了vgg_proto变量.weights当然是源模型文件,vgg_proto则是weights对应模型的模型定义文件,这里为什么要定义这个参数呢,原因是vgg网络中fc6和fc7层类型都是全连层,而我们定义的FCN网络中fc7和fc6层为1x1的conv层,其实本质上两者的参数都是一样的,但由于层类型不同使用solver.net.copy_from(weights)方式不能将VGG中fc6和fc7的训练好的参数赋给FCN,这就会导致我们在训练的时候loss一直很大不会减小.
(3)使用GPU加速,所以添加了caffe.set_device(0)和caffe.set_mode_gpu()
(4)跟上面(2)中提到的相照应,修改了模型参数的赋值方式,注释了solver.net.copy_from(weights),添加了

vgg_net=caffe.Net(vgg_proto,weights,caffe.TRAIN)
surgery.transplant(solver.net,vgg_net)
del vgg_net

注:VGG_ILSVRC_16_layers_deploy.prototxt下载地址:http://pan.baidu.com/s/1geLL6Sz

(5)修改了val数据集样本名称索引的读入方式(这也是有由于我使用的样本集存放方式不同导致的),将val = np.loadtxt('../data/segvalid11.txt', dtype=str)修改为

val = os.listdir('/home/irsa/fcn.berkeleyvision.org/voc-fcn8s/data/dataset/val/images')

至此,所有的修改完成了

四. 进行finetune

在终端中cd 到solve.py存放的目录,然后运行:

./solve.py

参考文献:

http://blog.csdn.net/wangkun1340378/article/details/56834642
http://www.cnblogs.com/xuanxufeng/p/6243342.html
https://blog.csdn.net/jiongnima/article/details/78549326

---

版权为NBJ所有 请勿转载

---