template<typename Dtype>
int feature_extraction_pipeline(int argc, char** argv);
int main(int argc, char** argv) {
return feature_extraction_pipeline<float>(argc, argv);
// return feature_extraction_pipeline<double>(argc, argv);
}
template<typename Dtype>
int feature_extraction_pipeline(int argc, char** argv) {
::google::InitGoogleLogging(argv[0]);
const int num_required_args = 5;
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_dataset_name1[,name2,...] num_mini_batches db_type"
" [CPU/GPU] [DEVICE_ID=0]\n"
"Note: you can extract multiple features in one pass by specifying"
" multiple feature blob names and dataset names separated by ','."
" The names cannot contain white space characters and the number of blobs"
" and datasets 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);
}
arg_pos = 0; // the name of the executable
std::string pretrained_binary_proto(argv[++arg_pos]);
std::string feature_extraction_proto(argv[++arg_pos]);
shared_ptr<Net<Dtype> > feature_extraction_net(
new Net<Dtype>(feature_extraction_proto, caffe::TEST));
feature_extraction_net->CopyTrainedLayersFrom(pretrained_binary_proto);
std::string extract_feature_blob_names(argv[++arg_pos]);
std::vector<std::string> blob_names;
boost::split(blob_names, extract_feature_blob_names, boost::is_any_of(","));
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]))
<< "Unknown feature blob name " << blob_names[i]
<< " in the network " << feature_extraction_proto;
}
int num_mini_batches = atoi(argv[++arg_pos]);
int index_exe = atoi(argv[5]);
LOG(ERROR)<< "Extacting Features";
FILE *fp;
if(index_exe == 0)
{
fp = fopen("/home/caffe/examples/_temp/lmdb1.txt", "w+");
}
else
{
fp = fopen("/home/caffe/examples/_temp/lmdb2.txt", "w+");
}
std::vector<Blob<Dtype>*> input_vec;
for (int batch_index = 0; batch_index < num_mini_batches; ++batch_index) {
feature_extraction_net->Forward(input_vec);
for (int i = 0; i < 1; ++i) {
const shared_ptr<Blob<Dtype> > 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;
const Dtype* feature_blob_data;
for (int n = 0; n < batch_size; ++n) {
feature_blob_data = feature_blob->cpu_data() +
feature_blob->offset(n);
for (int d = 0; d < dim_features; ++d) {
fprintf(fp, "%6f%c", feature_blob_data[d], ' ');
}
} // 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;
}
caffe中feature提取时做的修改
最新推荐文章于 2019-05-31 14:35:09 发布