快速入手一個簡單的分類網絡

本系列文章由 @yhl_leo 出品，轉載請註明出處。
文章連接： http://blog.csdn.net/yhl_leo/article/details/53727411

---

在之前的一篇博客中，我整理了如何根據CIFAR10的數據組織方式，製做本身的數據集，而後略微調整tensorflow 提供的demo進行訓練，得到了一些關注，如今從新公佈一個簡單的方法，不須要製做像CIFAR10那樣的數據集，也不用lmdb數據格式，直接使用原始數據，利用caffe訓練簡單的分類網絡。

發佈於GitHub: yhlleo/CreateSimpleNetworks.

在caffe的layer中，已有image_data_layer，對於image+label類型的訓練數據，數據讀取過程很簡單：

LOG(INFO) << "Opening file " << source;
  std::ifstream infile(source.c_str());
  string line;
  size_t pos;
  int label;
  while (std::getline(infile, line)) {
    pos = line.find_last_of(' ');
    label = atoi(line.substr(pos + 1).c_str());
    lines_.push_back(std::make_pair(line.substr(0, pos), label));
  }

  CHECK(!lines_.empty()) << "File is empty";

  if (this->layer_param_.image_data_param().shuffle()) {
    // randomly shuffle data
    LOG(INFO) << "Shuffling data";
    const unsigned int prefetch_rng_seed = caffe_rng_rand();
    prefetch_rng_.reset(new Caffe::RNG(prefetch_rng_seed));
    ShuffleImages();
  }
  LOG(INFO) << "A total of " << lines_.size() << " images.";

即，須要製做訓練文件列表格式爲：

...
/path/img1.jpg 0
/path/img2.jpg 1
...

完成訓練文件列表後，簡單搭建起一個小型網絡：

指定好train.prototxt， solver.prototxt和deploy.prototxt文件，就能夠訓練。

啓動訓練：

## train.py ##
from __future__ import division
import numpy as np
import sys
caffe_root = '/path/caffe/' 
sys.path.insert(0, caffe_root)
import caffe

# init
caffe.set_mode_gpu()
caffe.set_device(0)

solver = caffe.SGDSolver('/path/Models/solver.prototxt')
solver.step(60000)

批量測試：

import numpy as np
import os, cv2
import time
import caffe

# Make sure that caffe is on the python path:
caffe_root = '/path/caffe/'  
import sys
sys.path.insert(0, caffe_root + 'python')

caffe.set_mode_gpu()
caffe.set_device(0)

def findImages(dir,topdown=True):
    im_list = []
    if not os.path.exists(dir):
        print "Path for {} not exist!".format(dir)
        raise
    else:
        for root, dirs, files in os.walk(dir, topdown):
            for fl in files:
                im_list.append(os.path.join(root, fl))
    return im_list

data_root = '/path/test/test1'
test_lst = findImages(data_root)
savefolder = '/path/test/'
name = 'test1.txt'
OutDir = open(savefolder+name, 'w');

net = caffe.Net('/path/Models/xh_deploy.prototxt', \
    '/path/train/net_iter_60000.caffemodel', caffe.TEST)
time_consum = []

for idx in range(len(test_lst)):
    im = cv2.imread(test_lst[idx], cv2.IMREAD_UNCHANGED)
    sp = im.shape

    in_ = np.array(im, dtype=np.float32)
    in_ = in_[:,:,::-1]
    in_ = in_.transpose((2,0,1))
    net.blobs['data'].reshape(1, *in_.shape)
    net.blobs['data'].data[...] = in_

    start =time.clock()
    net.forward()   
    end = time.clock()
    time_consum.append(end-start)

    fuse = net.blobs['prob'].data[0]
    fname = test_lst[idx].split('/')[-1]
    OutDir.write("%s %.3f %.3f %.3f\n"%(fname, fuse[0], fuse[1], fuse[2]))
print sum(time_consum)/len(time_consum)
OutDir.close()

測試結果（數據集分爲兩類），所以四列分別對應着：文件名，label爲0的機率，label爲1的機率和其它類別的機率：

1-1.jpg 1.000 0.000 0.000
1-2.jpg 1.000 0.000 0.000
1-3.jpg 1.000 0.000 0.000
1-4.jpg 1.000 0.000 0.000
1-5.jpg 1.000 0.000 0.000
1-6.jpg 1.000 0.000 0.000
1-7.jpg 1.000 0.000 0.000
1-8.jpg 1.000 0.000 0.000
1001-1.jpg 0.594 0.405 0.001
1002-1.jpg 0.009 0.990 0.000
1002-10.jpg 1.000 0.000 0.000
...