caffe自定义层：imagedata层实现多标签

骑猪看日落 2022-02-23 09:14 299阅读 0赞

简介注意一下caffe的版本，我用的是最新版本的

我们都知道ImageDataLayer是直接读取原图进行分类，它的label是单label，文件格式如下

train.txt示例  
001.jpg 1  
002.jpg 2  
003.jpg 3

layer \{  
name: "demo"  
type: "ImageData"  
top: "data"  
top: "label"  
include \{  
phase: TRAIN  
\}  
transform\_param \{  
scale: 0.00390625  
mean\_value: 128  
\}  
image\_data\_param \{  
source: "your path/train.txt"  
root\_folder: "your image data path"  
new\_height: xxx  
new\_width: xxx  
batch\_size: 32  
shuffle: true \#每个epoch都会进行shuffle  
\}  
\}  
修改代码

由于ImageataLayer的限制，我们只能在train.txt中放置单label，现在我们来修改ImageataLayer的代码来实现多label  
主要涉及三个文件

caffe/src/caffe/proto/caffe.proto  
caffe/include/caffe/layers/image\_data\_layer.hpp  
caffe/src/caffe/layers/image\_data\_layer.cpp  
定位到caffe/src/caffe/proto/caffe.proto中message ImageDataParameter

// 添加一个参数  
// Specify the label dim. default 1.  
// 有几种label，比如性别、年龄两种label，在网络结构里就把此参数设置为2  
optional uint32 label\_dim = IDNumber \[default = 1\];   
// IDNumber是和其它参数不冲突的ID数字  
定位到caffe/include/caffe/layers/image\_data\_layer.hpp

// 修改vector<std::pair<std::string, int> > lines\_;  
// string对应那个train.txt中的图片名称，in对应label，我们把int改为int\*,实现多label  
vector<std::pair<std::string, int \*> > lines\_;  
定位到caffe/src/caffe/layers/image\_data\_layer.cpp

// DataLayerSetUp函数  
// 原本的加载图片名称和label的代码  
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));  
\}  
// 修改为这样  
std::ifstream infile(source.c\_str());  
string filename;  
// 获取label的种类  
int label\_dim = this->layer\_param\_.image\_data\_param().label\_dim();  
// 注意这里默认每个label直接以空格隔开，每个图片名称及其label占一行，如果你的格式不同，可自行修改读取方式  
while (infile >> filename) \{  
int\* labels = new int\[label\_dim\];  
for(int i = 0;i < label\_dim;++i)\{  
infile >> labels\[i\];  
\}  
lines\_.push\_back(std::make\_pair(filename, labels));  
\}  
// 原本的输出label  
vector<int> label\_shape(1, batch\_size);  
top\[1\]->Reshape(label\_shape);  
for (int i = 0; i < this->prefetch\_.size(); ++i) \{  
this->prefetch\_\[i\]->label\_.Reshape(label\_shape);  
\}  
// 修改为这样  
 vector<int> label\_shape(2);  
label\_shape\[0\] = batch\_size;  
label\_shape\[1\] = label\_dim;  
top\[1\]->Reshape(label\_shape); // label的输出shape batch\_size\*label\_dim  
for (int i = 0; i < this->PREFETCH\_COUNT; ++i) \{  
this->prefetch\_\[i\].label\_.Reshape(label\_shape);  
\}  
// 注意：caffe最新版本prefetch\_的结构由之前的Batch<Dtype> prefetch\_\[PREFETCH\_COUNT\];  
// 改为 vector<shared\_ptr<Batch<Dtype> > > prefetch\_; 由对象数组改为了存放shared指针的vector。  
// 所以此处的this->PREFETCH\_COUNT改为this->prefetch\_.size();  
// 此处的this->prefetch\_\[i\].label\_.Reshape(label\_shape);  
// 改为this->prefetch\_\[i\]->label\_.Reshape(label\_shape);把.改成指针的->  
// load\_batch函数  
// 在函数一开始先获取下label\_dim参数  
int label\_dim = this->layer\_param\_.image\_data\_param().label\_dim();  
// 原本的预取label  
prefetch\_label\[item\_id\] = lines\_\[lines\_id\_\].second;  
// 修改为这样  
for(int i = 0;i < label\_dim;++i)\{  
// lines\_\[lines\_id\_\].second就是最开始改为的int\*,多label  
prefetch\_label\[item\_id \* label\_dim + i\] = lines\_\[lines\_id\_\].second\[i\];  
\}  
  
最后进行make

示例

train.txt  
001.jpg 1 3 2  
002.jpg 2 4 7  
003.jpg 3 0 9  
  
\# trainval.prototxt  
layer \{  
name: "data"  
type: "ImageData"  
top: "data"  
top: "label"  
include \{  
phase: TRAIN  
\}  
transform\_param \{  
mirror: true  
mean\_value: 128  
mean\_value: 128  
mean\_value: 128  
\}  
image\_data\_param \{  
mirror: true  
source: "your path/train.txt"  
root\_folder: "your image data path"  
new\_height: xxx  
new\_width: xxx  
batch\_size: 32  
shuffle: true \#每个epoch都会进行shuffle  
label\_dim: 3  
\}  
\}  
layer \{  
name: "slice"  
type: "Slice"  
bottom: "label"  
top: "label\_1"  
top: "label\_2"  
top: "label\_3"  
slice\_param \{  
axis: 1  
slice\_point:1  
slice\_point:2  
\}  
\}  
......  
\# 后续代码参考我另一篇caffe实现多label输入的代码  
总结：https://blog.csdn.net/u013010889/article/details/54614067