caffe导读（1） caffe源码修改：抽取任意一张图片的特征

caffe源码修改：抽取任意一张图片的特征

目前caffe不是很完善，输入的图片数据需要在prototxt指定路径。但是我们往往有这么一个需求：训练后得到一个模型文件，我们想拿这个模型文件来对一张图片抽取特征或者预测分类等。如果非得在prototxt指定路径，就很不方便。因此，这样的工具才是我们需要的：给一个可执行文件通过命令行来传递图片路径，然后caffe读入图片数据，进行一次正向传播。

因此我做了这么一个工具，用来抽取任意一张图片的特征。

这工具的使用方法如下：

extract_one_feature.bin ./model/caffe_reference_imagenet_model ./examples/_temp/imagenet_val.prototxt fc7 ./examples/_temp/features /media/G/imageset/clothing/针织衫/针织衫_426.jpg CPU

参数1：./model/caffe_reference_imagenet_model是训练后的模型文件

参数2：./examples/_temp/imagenet_val.prototxt 网络配置文件

参数3：fc7是blob的名字

参数4：./examples/_temp/features 将该图片的特征保存在该文件

参数5：图片路径

参数6：GPU或者CPU模式

（其实我还想到更好的工具，如果该可执行文件是监听模式的，就是通过一定的方式，给该进程传递 图片路径，进程接到任务就执行。

这样子的话，就不需要每次抽一张图片都要申请内存空间。(*^__^*) 嘻嘻……）

下面给出初步修改方法，大家可以根据自己需求再修改。

extract_one_feature.cpp（该文件参考过源码中extract_features.cpp修改）

[cpp]  view plain copy

1. #include <stdio.h>  // for snprintf  
2. #include <string>  
3. #include <vector>  
4. #include <iostream>  
5. #include <fstream>  
6.   
7. #include "boost/algorithm/string.hpp"  
8. #include "google/protobuf/text_format.h"  
9. #include "leveldb/db.h"  
10. #include "leveldb/write_batch.h"  
11.   
12. #include "caffe/blob.hpp"  
13. #include "caffe/common.hpp"  
14. #include "caffe/net.hpp"  
15. #include "caffe/proto/caffe.pb.h"  
16. #include "caffe/util/io.hpp"  
17. #include "caffe/vision_layers.hpp"  
18.   
19. using namespace caffe;  // NOLINT(build/namespaces)  
20.   
21. template<typename Dtype>  
22. int feature_extraction_pipeline(int argc, char** argv);  
23.   
24. int main(int argc, char** argv) {  
25.   return feature_extraction_pipeline<float>(argc, argv);  
26. //  return feature_extraction_pipeline<double>(argc, argv);  
27. }  
28.   
29. template<typename Dtype>  
30. class writeDb  
31. {  
32. public:  
33.     void open(string dbName)  
34.     {  
35.         db.open(dbName.c_str());  
36.     }  
37.     void write(const Dtype &data)  
38.     {  
39.         db<<data;  
40.     }  
41.     void write(const string &str)  
42.     {  
43.         db<<str;  
44.     }  
45.     virtual ~writeDb()  
46.     {  
47.         db.close();  
48.     }  
49. private:  
50.     std::ofstream db;  
51. };  
52.   
53. template<typename Dtype>  
54. int feature_extraction_pipeline(int argc, char** argv) {  
55.   ::google::InitGoogleLogging(argv[0]);  
56.   const int num_required_args = 6;  
57.   if (argc < num_required_args) {  
58.     LOG(ERROR)<<  
59.     "This program takes in a trained network and an input data layer, and then"  
60.     " extract features of the input data produced by the net.\n"  
61.     "Usage: extract_features  pretrained_net_param"  
62.     "  feature_extraction_proto_file  extract_feature_blob_name1[,name2,...]"  
63.     "  save_feature_leveldb_name1[,name2,...]  img_path  [CPU/GPU]"  
64.     "  [DEVICE_ID=0]\n"  
65.     "Note: you can extract multiple features in one pass by specifying"  
66.     " multiple feature blob names and leveldb names seperated by ','."  
67.     " The names cannot contain white space characters and the number of blobs"  
68.     " and leveldbs must be equal.";  
69.     return 1;  
70.   }  
71.   int arg_pos = num_required_args;  
72.   
73.   arg_pos = num_required_args;  
74.   if (argc > arg_pos && strcmp(argv[arg_pos], "GPU") == 0) {  
75.     LOG(ERROR)<< "Using GPU";  
76.     uint device_id = 0;  
77.     if (argc > arg_pos + 1) {  
78.       device_id = atoi(argv[arg_pos + 1]);  
79.       CHECK_GE(device_id, 0);  
80.     }  
81.     LOG(ERROR) << "Using Device_id=" << device_id;  
82.     Caffe::SetDevice(device_id);  
83.     Caffe::set_mode(Caffe::GPU);  
84.   } else {  
85.     LOG(ERROR) << "Using CPU";  
86.     Caffe::set_mode(Caffe::CPU);  
87.   }  
88.   Caffe::set_phase(Caffe::TEST);  
89.   
90.   arg_pos = 0;  // the name of the executable  
91.   string pretrained_binary_proto(argv[++arg_pos]);//网络模型参数文件  
92.   
93.   string feature_extraction_proto(argv[++arg_pos]);  
94.   
95.   shared_ptr<Net<Dtype> > feature_extraction_net(  
96.       new Net<Dtype>(feature_extraction_proto));  
97.   
98.   feature_extraction_net->CopyTrainedLayersFrom(pretrained_binary_proto);//将网络参数load进内存  
99.   
100.   
101.   string extract_feature_blob_names(argv[++arg_pos]);  
102.   vector<string> blob_names;//要抽取特征的layer的名字，可以是多个  
103.   boost::split(blob_names, extract_feature_blob_names, boost::is_any_of(","));  
104.   
105.   string save_feature_leveldb_names(argv[++arg_pos]);  
106.   vector<string> leveldb_names;// 这里我改写成一个levedb为一个文件，数据格式不使用真正的levedb，而是自定义  
107.   boost::split(leveldb_names, save_feature_leveldb_names,  
108.                boost::is_any_of(","));  
109.   CHECK_EQ(blob_names.size(), leveldb_names.size()) <<  
110.       " the number of blob names and leveldb names must be equal";  
111.   size_t num_features = blob_names.size();  
112.   
113.   for (size_t i = 0; i < num_features; i++) {  
114.     CHECK(feature_extraction_net->has_blob(blob_names[i]))  //检测blob的名字在网络中是否存在  
115.         << "Unknown feature blob name " << blob_names[i]  
116.         << " in the network " << feature_extraction_proto;  
117.   }  
118.   
119.   
120.   vector<shared_ptr<writeDb<Dtype> > > feature_dbs;  
121.   for (size_t i = 0; i < num_features; ++i) //打开db，准备写入数据  
122.   {  
123.     LOG(INFO)<< "Opening db " << leveldb_names[i];  
124.     writeDb<Dtype>* db = new writeDb<Dtype>();  
125.     db->open(leveldb_names[i]);  
126.     feature_dbs.push_back(shared_ptr<writeDb<Dtype> >(db));  
127.   }  
128.   
129.   
130.   
131.   LOG(ERROR)<< "Extacting Features";  
132.   
133.   const shared_ptr<Layer<Dtype> > layer = feature_extraction_net->layer_by_name("data");//获取第一层  
134.   MyImageDataLayer<Dtype>* my_layer = (MyImageDataLayer<Dtype>*)layer.get();  
135.   my_layer->setImgPath(argv[++arg_pos],1);//"/media/G/imageset/clothing/针织衫/针织衫_1.jpg"  
136.   //设置图片路径  
137.   
138.   vector<Blob<float>*> input_vec;  
139.   vector<int> image_indices(num_features, 0);  
140.   int num_mini_batches = 1;//atoi(argv[++arg_pos]);//共多少次迭代。  每次迭代的数量在prototxt用batchsize指定  
141.   for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index) //共num_mini_batches次迭代  
142.   {  
143.     feature_extraction_net->Forward(input_vec);//一次正向传播  
144.     for (int i = 0; i < num_features; ++i) //多层特征  
145.     {  
146.       const shared_ptr<Blob<Dtype> > feature_blob = feature_extraction_net  
147.           ->blob_by_name(blob_names[i]);  
148.       int batch_size = feature_blob->num();  
149.       int dim_features = feature_blob->count() / batch_size;  
150.   
151.       Dtype* feature_blob_data;  
152.   
153.       for (int n = 0; n < batch_size; ++n)  
154.       {  
155.         feature_blob_data = feature_blob->mutable_cpu_data() +  
156.             feature_blob->offset(n);  
157.         feature_dbs[i]->write("3 ");  
158.         for (int d = 0; d < dim_features; ++d)  
159.         {  
160.           feature_dbs[i]->write((Dtype)(d+1));  
161.           feature_dbs[i]->write(":");  
162.           feature_dbs[i]->write(feature_blob_data[d]);  
163.           feature_dbs[i]->write(" ");  
164.         }  
165.         feature_dbs[i]->write("\n");  
166.   
167.       }  // for (int n = 0; n < batch_size; ++n)  
168.     }  // for (int i = 0; i < num_features; ++i)  
169.   }  // for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index)  
170.   
171.   
172.   LOG(ERROR)<< "Successfully extracted the features!";  
173.   return 0;  
174. }  

my_data_layer.cpp（参考image_data_layer修改）

[cpp]  view plain copy

1. #include <fstream>  // NOLINT(readability/streams)  
2. #include <iostream>  // NOLINT(readability/streams)  
3. #include <string>  
4. #include <utility>  
5. #include <vector>  
6.   
7. #include "caffe/layer.hpp"  
8. #include "caffe/util/io.hpp"  
9. #include "caffe/util/math_functions.hpp"  
10. #include "caffe/util/rng.hpp"  
11. #include "caffe/vision_layers.hpp"  
12.   
13. namespace caffe {  
14.   
15.   
16. template <typename Dtype>  
17. MyImageDataLayer<Dtype>::~MyImageDataLayer<Dtype>() {  
18. }  
19.   
20.   
21. template <typename Dtype>  
22. void MyImageDataLayer<Dtype>::setImgPath(string path,int label)  
23. {  
24.     lines_.clear();  
25.     lines_.push_back(std::make_pair(path, label));  
26. }  
27.   
28.   
29. template <typename Dtype>  
30. void MyImageDataLayer<Dtype>::SetUp(const vector<Blob<Dtype>*>& bottom,  
31.       vector<Blob<Dtype>*>* top) {  
32.   Layer<Dtype>::SetUp(bottom, top);  
33.   const int new_height  = this->layer_param_.image_data_param().new_height();  
34.   const int new_width  = this->layer_param_.image_data_param().new_width();  
35.   CHECK((new_height == 0 && new_width == 0) ||  
36.       (new_height > 0 && new_width > 0)) << "Current implementation requires "  
37.       "new_height and new_width to be set at the same time.";  
38.   
39.   /* 
40.    * 因为下面需要随便拿一张图片来初始化blob。 
41.    * 因此需要硬盘上有一张图片。 
42.    * 1 从prototxt读取一张图片的路径， 
43.    * 2 其实也可以在这里将用于初始化的图片路径写死 
44.   */  
45.   
46.   /*1*/  
47.   /* 
48.   const string& source = this->layer_param_.image_data_param().source(); 
49.   LOG(INFO) << "Opening file " << source; 
50.   std::ifstream infile(source.c_str()); 
51.   string filename; 
52.   int label; 
53.   while (infile >> filename >> label) { 
54.     lines_.push_back(std::make_pair(filename, label)); 
55.   } 
56.   */  
57.   
58.   /*2*/  
59.   lines_.push_back(std::make_pair("/home/linger/init.jpg",1));  
60.   
61.   //上面1和2代码可以任意用一段  
62.   
63.   lines_id_ = 0;  
64.   // Read a data point, and use it to initialize the top blob. (随便)读取一张图片，来初始化blob  
65.   Datum datum;  
66.   CHECK(ReadImageToDatum(lines_[lines_id_].first, lines_[lines_id_].second,  
67.                          new_height, new_width, &datum));  
68.   // image  
69.   const int crop_size = this->layer_param_.image_data_param().crop_size();  
70.   const int batch_size = 1;//this->layer_param_.image_data_param().batch_size();  
71.   const string& mean_file = this->layer_param_.image_data_param().mean_file();  
72.   if (crop_size > 0) {  
73.     (*top)[0]->Reshape(batch_size, datum.channels(), crop_size, crop_size);  
74.     prefetch_data_.Reshape(batch_size, datum.channels(), crop_size, crop_size);  
75.   } else {  
76.     (*top)[0]->Reshape(batch_size, datum.channels(), datum.height(),  
77.                        datum.width());  
78.     prefetch_data_.Reshape(batch_size, datum.channels(), datum.height(),  
79.         datum.width());  
80.   }  
81.   LOG(INFO) << "output data size: " << (*top)[0]->num() << ","  
82.       << (*top)[0]->channels() << "," << (*top)[0]->height() << ","  
83.       << (*top)[0]->width();  
84.   // label  
85.   (*top)[1]->Reshape(batch_size, 1, 1, 1);  
86.   prefetch_label_.Reshape(batch_size, 1, 1, 1);  
87.   // datum size  
88.   datum_channels_ = datum.channels();  
89.   datum_height_ = datum.height();  
90.   datum_width_ = datum.width();  
91.   datum_size_ = datum.channels() * datum.height() * datum.width();  
92.   CHECK_GT(datum_height_, crop_size);  
93.   CHECK_GT(datum_width_, crop_size);  
94.   // check if we want to have mean  
95.   if (this->layer_param_.image_data_param().has_mean_file()) {  
96.     BlobProto blob_proto;  
97.     LOG(INFO) << "Loading mean file from" << mean_file;  
98.     ReadProtoFromBinaryFile(mean_file.c_str(), &blob_proto);  
99.     data_mean_.FromProto(blob_proto);  
100.     CHECK_EQ(data_mean_.num(), 1);  
101.     CHECK_EQ(data_mean_.channels(), datum_channels_);  
102.     CHECK_EQ(data_mean_.height(), datum_height_);  
103.     CHECK_EQ(data_mean_.width(), datum_width_);  
104.   } else {  
105.     // Simply initialize an all-empty mean.  
106.     data_mean_.Reshape(1, datum_channels_, datum_height_, datum_width_);  
107.   }  
108.   // Now, start the prefetch thread. Before calling prefetch, we make two  
109.   // cpu_data calls so that the prefetch thread does not accidentally make  
110.   // simultaneous cudaMalloc calls when the main thread is running. In some  
111.   // GPUs this seems to cause failures if we do not so.  
112.   prefetch_data_.mutable_cpu_data();  
113.   prefetch_label_.mutable_cpu_data();  
114.   data_mean_.cpu_data();  
115.   
116.   
117. }  
118.   
119. //--------------------------------下面是读取一张图片数据-----------------------------------------------  
120. template <typename Dtype>  
121. void MyImageDataLayer<Dtype>::fetchData() {  
122.       Datum datum;  
123.       CHECK(prefetch_data_.count());  
124.       Dtype* top_data = prefetch_data_.mutable_cpu_data();  
125.       Dtype* top_label = prefetch_label_.mutable_cpu_data();  
126.       ImageDataParameter image_data_param = this->layer_param_.image_data_param();  
127.       const Dtype scale = image_data_param.scale();//image_data_layer相关参数  
128.       const int batch_size = 1;//image_data_param.batch_size(); 这里我们只需要一张图片  
129.   
130.       const int crop_size = image_data_param.crop_size();  
131.       const bool mirror = image_data_param.mirror();  
132.       const int new_height = image_data_param.new_height();  
133.       const int new_width = image_data_param.new_width();  
134.   
135.       if (mirror && crop_size == 0) {  
136.         LOG(FATAL) << "Current implementation requires mirror and crop_size to be "  
137.             << "set at the same time.";  
138.       }  
139.       // datum scales  
140.       const int channels = datum_channels_;  
141.       const int height = datum_height_;  
142.       const int width = datum_width_;  
143.       const int size = datum_size_;  
144.       const int lines_size = lines_.size();  
145.       const Dtype* mean = data_mean_.cpu_data();  
146.   
147.       for (int item_id = 0; item_id < batch_size; ++item_id) {//读取一图片  
148.         // get a blob  
149.         CHECK_GT(lines_size, lines_id_);  
150.         if (!ReadImageToDatum(lines_[lines_id_].first,  
151.               lines_[lines_id_].second,  
152.               new_height, new_width, &datum)) {  
153.           continue;  
154.         }  
155.         const string& data = datum.data();  
156.         if (crop_size) {  
157.           CHECK(data.size()) << "Image cropping only support uint8 data";  
158.           int h_off, w_off;  
159.           // We only do random crop when we do training.  
160.             h_off = (height - crop_size) / 2;  
161.             w_off = (width - crop_size) / 2;  
162.   
163.             // Normal copy 正常读取，把裁剪后的图片数据读给top_data  
164.             for (int c = 0; c < channels; ++c) {  
165.               for (int h = 0; h < crop_size; ++h) {  
166.                 for (int w = 0; w < crop_size; ++w) {  
167.                   int top_index = ((item_id * channels + c) * crop_size + h)  
168.                                   * crop_size + w;  
169.                   int data_index = (c * height + h + h_off) * width + w + w_off;  
170.                   Dtype datum_element =  
171.                       static_cast<Dtype>(static_cast<uint8_t>(data[data_index]));  
172.                   top_data[top_index] = (datum_element - mean[data_index]) * scale;  
173.                 }  
174.               }  
175.             }  
176.   
177.         } else {  
178.           // Just copy the whole data 正常读取，把图片数据读给top_data  
179.           if (data.size()) {  
180.             for (int j = 0; j < size; ++j) {  
181.               Dtype datum_element =  
182.                   static_cast<Dtype>(static_cast<uint8_t>(data[j]));  
183.               top_data[item_id * size + j] = (datum_element - mean[j]) * scale;  
184.             }  
185.           } else {  
186.             for (int j = 0; j < size; ++j) {  
187.               top_data[item_id * size + j] =  
188.                   (datum.float_data(j) - mean[j]) * scale;  
189.             }  
190.           }  
191.         }  
192.         top_label[item_id] = datum.label();//读取该图片的标签  
193.   
194.       }  
195. }  
196.   
197. template <typename Dtype>  
198. Dtype MyImageDataLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,  
199.       vector<Blob<Dtype>*>* top) {  
200.   
201.   //更新input  
202.     fetchData();  
203.   
204.   // Copy the data  
205.   caffe_copy(prefetch_data_.count(), prefetch_data_.cpu_data(),  
206.              (*top)[0]->mutable_cpu_data());  
207.   caffe_copy(prefetch_label_.count(), prefetch_label_.cpu_data(),  
208.              (*top)[1]->mutable_cpu_data());  
209.   
210.   return Dtype(0.);  
211. }  
212.   
213. #ifdef CPU_ONLY  
214. STUB_GPU_FORWARD(ImageDataLayer, Forward);  
215. #endif  
216.   
217. INSTANTIATE_CLASS(MyImageDataLayer);  
218.   
219. }  // namespace caffe  

在data_layers.hpp添加一下代码，参考ImageDataLayer写的。

[cpp]  view plain copy

1. template <typename Dtype>  
2. class MyImageDataLayer : public Layer<Dtype>  {  
3.  public:  
4.   explicit MyImageDataLayer(const LayerParameter& param)  
5.       : Layer<Dtype>(param) {}  
6.   virtual ~MyImageDataLayer();  
7.   virtual void SetUp(const vector<Blob<Dtype>*>& bottom,  
8.       vector<Blob<Dtype>*>* top);  
9.   
10.   virtual inline LayerParameter_LayerType type() const {  
11.     return LayerParameter_LayerType_MY_IMAGE_DATA;  
12.   }  
13.   virtual inline int ExactNumBottomBlobs() const { return 0; }  
14.   virtual inline int ExactNumTopBlobs() const { return 2; }  
15.   void fetchData();  
16.   void setImgPath(string path,int label);  
17.  protected:  
18.   virtual Dtype Forward_cpu(const vector<Blob<Dtype>*>& bottom,  
19.       vector<Blob<Dtype>*>* top);  
20.   
21.   virtual void Backward_cpu(const vector<Blob<Dtype>*>& top,  
22.       const vector<bool>& propagate_down, vector<Blob<Dtype>*>* bottom) {}  
23.   
24.   
25.   vector<std::pair<std::string, int> > lines_;  
26.   int lines_id_;  
27.   int datum_channels_;  
28.   int datum_height_;  
29.   int datum_width_;  
30.   int datum_size_;  
31.   Blob<Dtype> prefetch_data_;  
32.   Blob<Dtype> prefetch_label_;  
33.   Blob<Dtype> data_mean_;  
34.   Caffe::Phase phase_;  
35. };  

修改caffe.proto,在适当的位置添加下面信息，也是参考image_data写的。

MY_IMAGE_DATA = 36;

optional MyImageDataParameter my_image_data_param = 36;

// Message that stores parameters used by MyImageDataLayer
message MyImageDataParameter {
  // Specify the data source.
  optional string source = 1;
  // For data pre-processing, we can do simple scaling and subtracting the
  // data mean, if provided. Note that the mean subtraction is always carried
  // out before scaling.
  optional float scale = 2 [default = 1];
  optional string mean_file = 3;
  // Specify the batch size.
  optional uint32 batch_size = 4;
  // Specify if we would like to randomly crop an image.
  optional uint32 crop_size = 5 [default = 0];
  // Specify if we want to randomly mirror data.
  optional bool mirror = 6 [default = false];
  // The rand_skip variable is for the data layer to skip a few data points
  // to avoid all asynchronous sgd clients to start at the same point. The skip
  // point would be set as rand_skip * rand(0,1). Note that rand_skip should not
  // be larger than the number of keys in the leveldb.
  optional uint32 rand_skip = 7 [default = 0];
  // Whether or not ImageLayer should shuffle the list of files at every epoch.
  optional bool shuffle = 8 [default = false];
  // It will also resize images if new_height or new_width are not zero.
  optional uint32 new_height = 9 [default = 0];
  optional uint32 new_width = 10 [default = 0];
}

以上每行位置不在一起，可以参考读一个image_data对应的位置。

本文作者：linger

本文链接：http://blog.csdn.net/lingerlanlan/article/details/39400375