为 Caffe 添加新的 DataLayer

目标

复现 DeepID 用 Caffe 实现人脸识别时，网络的训练的框架往往是这样的：



就是说 Image List 中的数据是按对整理好的，类内 (intra class) 类间 (inter class) 数据交替排列。这样就可以直接利用 ImageDataLayer 获得一个个均匀的 Batch。现在只要对 Loss Layer 简单做一下修改，网络已经可以正常训练了，相当简单。

但简单的代价也相当明显：

数据类内与类间的组合已经固定，想用新的组合训练，就得重新整理一个 ImageList。

左图中类内数据与类间数据交替排列的方式不利于 LossLayer 计算的优化。

新的框架

所以如果能把数据整理这一步集成到 DataLayer 中，将会得到更大程度的简化。而且可以额外附加一些功能，比如随机抽取类内类间数据；类内数据抽取区别较大的组合，类间数据则抽取区别较小的组合，如果计算速度允许的话。

于是新的训练框架应该是下面这样的，DataLayer 的名字都取好了，叫 TripletDataLayer!



TripletDataLayer 会在每次迭代过程中抽取一个 Batch 的数据，Batch 中的数据等分为 Idenities A, Idenities P, Idenities N 三个部分, 分别代表 FaceNet 1 中提出的 Anchor, Positive, Negative。Anchor 与 Positive 是类内组合，Anchor 与 Negative 是类间组合，后面都用这种方式表示。可以看到由于 Batch 内的数据是有序的，对应的 Loss Layer 就不需要 label 信息了（SoftMax Loss 除外）。

TripletDataLayer 的实现

为 Caffe 添加新数据层的步骤 ( 以 TripletDataLayer 为例 ) 如下：

在 src/caffe/proto/caffe.proto 中定义相关的参数

在 src/caffe/layers/ 目录下添加 triplet_data_layer.cpp

在 include/caffe/layers/ 目录下添加 triplet_data_layer.hpp

在 src/caffe/test 目录下添加 test_triplet_data_layer.cpp (可选)

第 2, 3 步就是实现 TripletDataLayer 类，最为重要。然后把需要定制的参数添加到 caffe.proto 文件中。准确地完成这两步的话，添加 DataLayer 层的工作就算完成了。如果不能确定，最好在 src/caffe/test 目录下写一个测试用例。

分析 ImageDataLayer

我的实现原则是 对 Caffe 源码做最少的修改。纵观现有的 DataLayer 最佳参考的是 ImageDataLayer 。理所当然地应该分析一下它的实现，下面是我画的一个非常粗糙的 ImageDataLayer 类图：



可以看到，ImageDataLayer 与普通的 Layer 不同，它不需要实现 .cu 文件。因为 GPU 的操作已经在高层的类中实现了，同样数据预取操作 (经典的生产者消费者问题，每次预取 3 个 Batch) 也已在高层类中实现。所以我们只要重写 load_batch() 函数即可。

在 caffe.proto 中定义外部参数

为参数定义一个 message 结构

在 ImageDataParameter 的基础上修改。虽然结构名的选择是自由的，但最好按 Caffe 的风格来写，即结构名与类名对应。有关 protobuf 的语法请自行 Google。

message TripletDataParameter
{
// 图片列表的路径，可以是文件名或文件路径
required string source = 1;

// 是否根据图片的特征筛选数据
optional bool use_feature = 15 [default = false];

// 图片特征文件的扩展名，特征文件与对应图片位于同一目录下
optional string feature_extension = 14  [default = ".feat"];

// 如果每个人的图片数量不等时，让每个 Batch 的数据分布与训练集相同
optional bool batch_follow_distribution = 16 [default = true];

// 存放图片的目录（所有图片都存放在同一目录时使用）
optional string root_folder = 12 [default = ""];

// Batch size，Layer 初始化时把它改为 3 的倍数
optional uint32 batch_size = 4 [default = 1];

// 启动时随机跳过几个数据，与 ImageDataLayer 略有不同
optional uint32 rand_skip = 7 [default = 0];

// 下面的参数的意义与 ImageDataLayer 完全相同
optional bool shuffle = 8 [default = false];
optional uint32 new_height = 9 [default = 0];
optional uint32 new_width = 10 [default = 0];
optional bool is_color = 11 [default = true];
optional float scale = 2 [default = 1];
optional string mean_file = 3;
optional uint32 crop_size = 5 [default = 0];
optional bool mirror = 6 [default = false];
}

在 LayerParameter 中添加参数的定义

同样，为了统一性参数名照样与类名相关。应注意标识号不能与前面的字段重复。

message LayerParameter {
...
optional TripletDataParameter triplet_data_param = 512<
4000
/span>;
}

然后在 TripletDataLayer 类中就可以读取这些参数了 ( caffe.pb.cc 与 caffe.pb.h 会在 Caffe 编译过程中自动编译出来 )。比如我们想读取配置文件中的 batch_size 参数，可以这样做：

int batch_size_ = this->layer_param_.triplet_data_param().batch_size();

待续。。。