1. 程式人生 > >為 Caffe 新增新的 DataLayer

為 Caffe 新增新的 DataLayer

目標

復現 DeepID 用 Caffe 實現人臉識別時,網路的訓練的框架往往是這樣的:

structure-orig

就是說 Image List 中的資料是按對整理好的,類內 (intra class) 類間 (inter class) 資料交替排列。這樣就可以直接利用 ImageDataLayer 獲得一個個均勻的 Batch。現在只要對 Loss Layer 簡單做一下修改,網路已經可以正常訓練了,相當簡單。

但簡單的代價也相當明顯:

  • 資料類內與類間的組合已經固定,想用新的組合訓練,就得重新整理一個 ImageList。
  • 左圖中類內資料與類間資料交替排列的方式不利於 LossLayer 計算的優化。

新的框架

所以如果能把資料整理這一步整合到 DataLayer 中,將會得到更大程度的簡化。而且可以額外附加一些功能,比如隨機抽取類內類間資料;類內資料抽取區別較大的組合,類間資料則抽取區別較小的組合,如果計算速度允許的話。

於是新的訓練框架應該是下面這樣的,DataLayer 的名字都取好了,叫 TripletDataLayer!

structure

TripletDataLayer 會在每次迭代過程中抽取一個 Batch 的資料,Batch 中的資料等分為 Idenities A, Idenities P, Idenities N 三個部分, 分別代表 FaceNet 1 中提出的 Anchor, Positive, Negative。Anchor 與 Positive 是類內組合,Anchor 與 Negative 是類間組合,後面都用這種方式表示。可以看到由於 Batch 內的資料是有序的,對應的 Loss Layer 就不需要 label 資訊了(SoftMax Loss 除外)。

A-N-P

TripletDataLayer 的實現

為 Caffe 新增新資料層的步驟 ( 以 TripletDataLayer 為例 ) 如下:

  1. 在 src/caffe/proto/caffe.proto 中定義相關的引數
  2. 在 src/caffe/layers/ 目錄下新增 triplet_data_layer.cpp
  3. 在 include/caffe/layers/ 目錄下新增 triplet_data_layer.hpp
  4. 在 src/caffe/test 目錄下新增 test_triplet_data_layer.cpp (可選)

第 2, 3 步就是實現 TripletDataLayer 類,最為重要。然後把需要定製的引數新增到 caffe.proto 檔案中。準確地完成這兩步的話,新增 DataLayer 層的工作就算完成了。如果不能確定,最好在 src/caffe/test 目錄下寫一個測試用例。

分析 ImageDataLayer

我的實現原則是 對 Caffe 原始碼做最少的修改。縱觀現有的 DataLayer 最佳參考的是 ImageDataLayer 。理所當然地應該分析一下它的實現,下面是我畫的一個非常粗糙的 ImageDataLayer 類圖:

classes

可以看到,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;
}

然後在 TripletDataLayer 類中就可以讀取這些引數了 ( caffe.pb.cc 與 caffe.pb.h 會在 Caffe 編譯過程中自動編譯出來 )。比如我們想讀取配置檔案中的 batch_size 引數,可以這樣做:

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

待續。。。