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修改）
#include // for snprintf
#include
#include
#include
#include
#include "boost/algorithm/string.hpp"
#include "google/protobuf/text_format.h"
#include "leveldb/db.h"
#include "leveldb/write_batch.h"
#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/net.hpp"
#include "caffe/proto/caffe.pb.h"
#include "caffe/util/io.hpp"
#include "caffe/vision_layers.hpp"
using namespace caffe; // NOLINT(build/namespaces)
template
int feature_extraction_pipeline(int argc, char** argv);
int main(int argc, char** argv) {
return feature_extraction_pipeline(argc, argv);
// return feature_extraction_pipeline(argc, argv);
}
template
class writeDb
{
public:
void open(string dbName)
{
db.open(dbName.c_str());
}
void write(const Dtype &data)
{
db<
int feature_extraction_pipeline(int argc, char** argv) {
::google::InitGoogleLogging(argv[0]);
const int num_required_args = 6;
if (argc < num_required_args) {
LOG(ERROR)<<
"This program takes in a trained network and an input data layer, and then"
" extract features of the input data produced by the net.\n"
"Usage: extract_features pretrained_net_param"
" feature_extraction_proto_file extract_feature_blob_name1[,name2,...]"
" save_feature_leveldb_name1[,name2,...] img_path [CPU/GPU]"
" [DEVICE_ID=0]\n"
"Note: you can extract multiple features in one pass by specifying"
" multiple feature blob names and leveldb names seperated by ','."
" The names cannot contain white space characters and the number of blobs"
" and leveldbs must be equal.";
return 1;
}
int arg_pos = num_required_args;
arg_pos = num_required_args;
if (argc > arg_pos && strcmp(argv[arg_pos], "GPU") == 0) {
LOG(ERROR)<< "Using GPU";
uint device_id = 0;
if (argc > arg_pos + 1) {
device_id = atoi(argv[arg_pos + 1]);
CHECK_GE(device_id, 0);
}
LOG(ERROR) << "Using Device_id=" << device_id;
Caffe::SetDevice(device_id);
Caffe::set_mode(Caffe::GPU);
} else {
LOG(ERROR) << "Using CPU";
Caffe::set_mode(Caffe::CPU);
}
Caffe::set_phase(Caffe::TEST);
arg_pos = 0; // the name of the executable
string pretrained_binary_proto(argv[++arg_pos]);//网络模型参数文件
string feature_extraction_proto(argv[++arg_pos]);
shared_ptr > feature_extraction_net(
new Net(feature_extraction_proto));
feature_extraction_net->CopyTrainedLayersFrom(pretrained_binary_proto);//将网络参数load进内存
string extract_feature_blob_names(argv[++arg_pos]);
vector blob_names;//要抽取特征的layer的名字，可以是多个
boost::split(blob_names, extract_feature_blob_names, boost::is_any_of(","));
string save_feature_leveldb_names(argv[++arg_pos]);
vector leveldb_names;// 这里我改写成一个levedb为一个文件，数据格式不使用真正的levedb，而是自定义
boost::split(leveldb_names, save_feature_leveldb_names,
boost::is_any_of(","));
CHECK_EQ(blob_names.size(), leveldb_names.size()) <<
" the number of blob names and leveldb names must be equal";
size_t num_features = blob_names.size();
for (size_t i = 0; i < num_features; i++) {
CHECK(feature_extraction_net->has_blob(blob_names[i])) //检测blob的名字在网络中是否存在
<< "Unknown feature blob name " << blob_names[i]
<< " in the network " << feature_extraction_proto;
}
vector > > feature_dbs;
for (size_t i = 0; i < num_features; ++i) //打开db，准备写入数据
{
LOG(INFO)<< "Opening db " << leveldb_names[i];
writeDb* db = new writeDb();
db->open(leveldb_names[i]);
feature_dbs.push_back(shared_ptr >(db));
}
LOG(ERROR)<< "Extacting Features";
const shared_ptr > layer = feature_extraction_net->layer_by_name("data");//获取第一层
MyImageDataLayer* my_layer = (MyImageDataLayer*)layer.get();
my_layer->setImgPath(argv[++arg_pos],1);//"/media/G/imageset/clothing/针织衫/针织衫_1.jpg"
//设置图片路径
vector*> input_vec;
vector image_indices(num_features, 0);
int num_mini_batches = 1;//atoi(argv[++arg_pos]);//共多少次迭代。 每次迭代的数量在prototxt用batchsize指定
for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index) //共num_mini_batches次迭代
{
feature_extraction_net->Forward(input_vec);//一次正向传播
for (int i = 0; i < num_features; ++i) //多层特征
{
const shared_ptr > feature_blob = feature_extraction_net
->blob_by_name(blob_names[i]);
int batch_size = feature_blob->num();
int dim_features = feature_blob->count() / batch_size;
Dtype* feature_blob_data;
for (int n = 0; n < batch_size; ++n)
{
feature_blob_data = feature_blob->mutable_cpu_data() +
feature_blob->offset(n);
feature_dbs[i]->write("3 ");
for (int d = 0; d < dim_features; ++d)
{
feature_dbs[i]->write((Dtype)(d+1));
feature_dbs[i]->write(":");
feature_dbs[i]->write(feature_blob_data[d]);
feature_dbs[i]->write(" ");
}
feature_dbs[i]->write("\n");
} // for (int n = 0; n < batch_size; ++n)
} // for (int i = 0; i < num_features; ++i)
} // for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index)
LOG(ERROR)<< "Successfully extracted the features!";
return 0;
}
my_data_layer.cpp（参考image_data_layer修改）
#include // NOLINT(readability/streams)
#include // NOLINT(readability/streams)
#include
#include
#include
#include "caffe/layer.hpp"
#include "caffe/util/io.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/rng.hpp"
#include "caffe/vision_layers.hpp"
namespace caffe {
template
MyImageDataLayer::~MyImageDataLayer() {
}
template
void MyImageDataLayer::setImgPath(string path,int label)
{
lines_.clear();
lines_.push_back(std::make_pair(path, label));
}
template
void MyImageDataLayer::SetUp(const vector*>& bottom,
vector*>* top) {
Layer::SetUp(bottom, top);
const int new_height = this->layer_param_.image_data_param().new_height();
const int new_width = this->layer_param_.image_data_param().new_width();
CHECK((new_height == 0 && new_width == 0) ||
(new_height > 0 && new_width > 0)) << "Current implementation requires "
"new_height and new_width to be set at the same time.";
/*
* 因为下面需要随便拿一张图片来初始化blob。
* 因此需要硬盘上有一张图片。
* 1 从prototxt读取一张图片的路径，
* 2 其实也可以在这里将用于初始化的图片路径写死
*/
/*1*/
/*
const string& source = this->layer_param_.image_data_param().source();
LOG(INFO) << "Opening file " << source;
std::ifstream infile(source.c_str());
string filename;
int label;
while (infile >> filename >> label) {
lines_.push_back(std::make_pair(filename, label));
}
*/
/*2*/
lines_.push_back(std::make_pair("/home/linger/init.jpg",1));
//上面1和2代码可以任意用一段
lines_id_ = 0;
// Read a data point, and use it to initialize the top blob. (随便)读取一张图片，来初始化blob
Datum datum;
CHECK(ReadImageToDatum(lines_[lines_id_].first, lines_[lines_id_].second,
new_height, new_width, &datum));
// image
const int crop_size = this->layer_param_.image_data_param().crop_size();
const int batch_size = 1;//this->layer_param_.image_data_param().batch_size();
const string& mean_file = this->layer_param_.image_data_param().mean_file();
if (crop_size > 0) {
(*top)[0]->Reshape(batch_size, datum.channels(), crop_size, crop_size);
prefetch_data_.Reshape(batch_size, datum.channels(), crop_size, crop_size);
} else {
(*top)[0]->Reshape(batch_size, datum.channels(), datum.height(),
datum.width());
prefetch_data_.Reshape(batch_size, datum.channels(), datum.height(),
datum.width());
}
LOG(INFO) << "output data size: " << (*top)[0]->num() << ","
<< (*top)[0]->channels() << "," << (*top)[0]->height() << ","
<< (*top)[0]->width();
// label
(*top)[1]->Reshape(batch_size, 1, 1, 1);
prefetch_label_.Reshape(batch_size, 1, 1, 1);
// datum size
datum_channels_ = datum.channels();
datum_height_ = datum.height();
datum_width_ = datum.width();
datum_size_ = datum.channels() * datum.height() * datum.width();
CHECK_GT(datum_height_, crop_size);
CHECK_GT(datum_width_, crop_size);
// check if we want to have mean
if (this->layer_param_.image_data_param().has_mean_file()) {
BlobProto blob_proto;
LOG(INFO) << "Loading mean file from" << mean_file;
ReadProtoFromBinaryFile(mean_file.c_str(), &blob_proto);
data_mean_.FromProto(blob_proto);
CHECK_EQ(data_mean_.num(), 1);
CHECK_EQ(data_mean_.channels(), datum_channels_);
CHECK_EQ(data_mean_.height(), datum_height_);
CHECK_EQ(data_mean_.width(), datum_width_);
} else {
// Simply initialize an all-empty mean.
data_mean_.Reshape(1, datum_channels_, datum_height_, datum_width_);
}
// Now, start the prefetch thread. Before calling prefetch, we make two
// cpu_data calls so that the prefetch thread does not accidentally make
// simultaneous cudaMalloc calls when the main thread is running. In some
// GPUs this seems to cause failures if we do not so.
prefetch_data_.mutable_cpu_data();
prefetch_label_.mutable_cpu_data();
data_mean_.cpu_data();
}
//--------------------------------下面是读取一张图片数据-----------------------------------------------
template
void MyImageDataLayer::fetchData() {
Datum datum;
CHECK(prefetch_data_.count());
Dtype* top_data = prefetch_data_.mutable_cpu_data();
Dtype* top_label = prefetch_label_.mutable_cpu_data();
ImageDataParameter image_data_param = this->layer_param_.image_data_param();
const Dtype scale = image_data_param.scale();//image_data_layer相关参数
const int batch_size = 1;//image_data_param.batch_size(); 这里我们只需要一张图片
const int crop_size = image_data_param.crop_size();
const bool mirror = image_data_param.mirror();
const int new_height = image_data_param.new_height();
const int new_width = image_data_param.new_width();
if (mirror && crop_size == 0) {
LOG(FATAL) << "Current implementation requires mirror and crop_size to be "
<< "set at the same time.";
}
// datum scales
const int channels = datum_channels_;
const int height = datum_height_;
const int width = datum_width_;
const int size = datum_size_;
const int lines_size = lines_.size();
const Dtype* mean = data_mean_.cpu_data();
for (int item_id = 0; item_id < batch_size; ++item_id) {//读取一图片
// get a blob
CHECK_GT(lines_size, lines_id_);
if (!ReadImageToDatum(lines_[lines_id_].first,
lines_[lines_id_].second,
new_height, new_width, &datum)) {
continue;
}
const string& data = datum.data();
if (crop_size) {
CHECK(data.size()) << "Image cropping only support uint8 data";
int h_off, w_off;
// We only do random crop when we do training.
h_off = (height - crop_size) / 2;
w_off = (width - crop_size) / 2;
// Normal copy 正常读取，把裁剪后的图片数据读给top_data
for (int c = 0; c < channels; ++c) {
for (int h = 0; h < crop_size; ++h) {
for (int w = 0; w < crop_size; ++w) {
int top_index = ((item_id * channels + c) * crop_size + h)
* crop_size + w;
int data_index = (c * height + h + h_off) * width + w + w_off;
Dtype datum_element =
static_cast(static_cast(data[data_index]));
top_data[top_index] = (datum_element - mean[data_index]) * scale;
}
}
}
} else {
// Just copy the whole data 正常读取，把图片数据读给top_data
if (data.size()) {
for (int j = 0; j < size; ++j) {
Dtype datum_element =
static_cast(static_cast(data[j]));
top_data[item_id * size + j] = (datum_element - mean[j]) * scale;
}
} else {
for (int j = 0; j < size; ++j) {
top_data[item_id * size + j] =
(datum.float_data(j) - mean[j]) * scale;
}
}
}
top_label[item_id] = datum.label();//读取该图片的标签
}
}
template
Dtype MyImageDataLayer::Forward_cpu(const vector*>& bottom,
vector*>* top) {
//更新input
fetchData();
// Copy the data
caffe_copy(prefetch_data_.count(), prefetch_data_.cpu_data(),
(*top)[0]->mutable_cpu_data());
caffe_copy(prefetch_label_.count(), prefetch_label_.cpu_data(),
(*top)[1]->mutable_cpu_data());
return Dtype(0.);
}
#ifdef CPU_ONLY
STUB_GPU_FORWARD(ImageDataLayer, Forward);
#endif
INSTANTIATE_CLASS(MyImageDataLayer);
} // namespace caffe
在data_layers.hpp添加一下代码，参考ImageDataLayer写的。
template
class MyImageDataLayer : public Layer {
public:
explicit MyImageDataLayer(const LayerParameter& param)
: Layer(param) {}
virtual ~MyImageDataLayer();
virtual void SetUp(const vector*>& bottom,
vector*>* top);
virtual inline LayerParameter_LayerType type() const {
return LayerParameter_LayerType_MY_IMAGE_DATA;
}
virtual inline int ExactNumBottomBlobs() const { return 0; }
virtual inline int ExactNumTopBlobs() const { return 2; }
void fetchData();
void setImgPath(string path,int label);
protected:
virtual Dtype Forward_cpu(const vector*>& bottom,
vector*>* top);
virtual void Backward_cpu(const vector*>& top,
const vector& propagate_down, vector*>* bottom) {}
vector > lines_;
int lines_id_;
int datum_channels_;
int datum_height_;
int datum_width_;
int datum_size_;
Blob prefetch_data_;
Blob prefetch_label_;
Blob data_mean_;
Caffe::Phase phase_;
}; 修改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