NVcaffe源碼閱讀——Net&Solver

NVcaffe源碼閱讀——Net&Solver

Solver

在caffe中，爲了區分多GPU數據並行時，負責更新網絡參數的solver被稱作root solver。其他僅僅計算梯度用的solver爲WorkerSolver，是Solver類的子類。nvcaffe將WorkerSolver類同Solver本身合併，將root solver和worker solver的概念融入到統一的框架下去，添加了is_root()等函數，也添加了Reduce()等用於並行處理的函數。

Step函數

nvcaffe中solver的Init()、InitTrainNet()和InitTestNet()沒有明顯改動。

1. 整理可學習參數的梯度空間

程序中使用了很多even(T val)函數。該函數返回一個val臨近的偶數：

inline T even(T val) { return val & 1 ? val + 1 : val; }

對於Step()函數，nvcaffe首先調用了net.cpp所定義的InitializeLearnableDiffSpace()函數處理參數的梯度數值存儲空間。和caffe使用ClearParamDiffs()函數對各個參數blob的diff空間進行寫0覆蓋相比，nvcaffe使用的InitializeLearnableDiffSpace()函數增加了內存對齊的操作（個人推斷）。該函數是專門針對GPU顯存的操作。在net.cpp的InitializeLearnableDiffSpace()函數中，nvcaffe使用了定義在gpu_memory.hpp中的Workspace結構體，將所有的參數梯度顯存放在連續的空間之中。該函數首先計算每個待學習參數所佔用的空間，計算方法是even(元素個數)*數據類型大小。個人推測之所以使用even()函數，是爲了在使用FLOAT16類型時能夠讓顯存對齊。之後使用這些計算好的大小來分配顯存空間，最終仍然使用的是try_allocate()函數（詳見”Blob的重新構建”的”Syncedmem”一節），並在此之後用0初始化空間。雖然各個梯度空間被像碎牛肉一樣地壓在了一起，但是各個空間的初始位置使用learnable_params_ptrs_這個指針數組保存了起來。

相比之下，caffe將梯度和參數本身以blob爲單位在一起存放。nvcaffe則是將learnable_params_ptrs_的各個元素賦給每個blob的diff指針，但空間是連續的。

2. 多GPU之間的模型分發

如果使用了多GPU並行，則需要將net分發到各個設備上。兩個框架的不同點在於callback_->on_start函數上，該函數定義在parallel.cpp的P2PSync空間裏。caffe使用了cudaMemcpyAsync函數來同步數據，nvcaffe使用的是ncclBcast（如果編譯時開啓了使用NCCL）並輔以一些同步用的函數。

3. 異步並行更新權重

nvcaffe在step()函數中另一個顯著的不同點在於權重的更新。在caffe中，權重的更新在每次前向後向傳播後進行，使用ApplyUpdate()函數，而nvcaffe則是在循環迭代之前開啓了一個新線程，專門負責權重的更新，調用solver.cpp中的Reduce()函數，以及進一步的net.cpp中的ReduceAndUpdate()函數：

    reduce_thread_.reset(new boost::thread(&Solver::Reduce, this,
      Caffe::current_device(), mode, random_seed, solver_count, root_solver));
    while (iter_ < stop_iter) {  
    ...
    //start iteration

3.1 異步模型

step()函數的異步更新權重功能基本依賴於nvcaffe對net.cpp的改進。net.cpp維護了一個異步隊列，該隊列存儲的元素是各個層的id號碼：

BlockingQueue<int> reduction_queue_;

這個異步隊列存儲了經過網絡反向計算了梯度之後了的、需要被更新權重的層的id。一方面，step()線程調用net.cpp的BackwardFromToAu()函數計算反向過程。每當計算完一個層的數據，且該層的參數需要被更新時，便將該層的id放入reduction_queue_中。每當計算完一個batch的數據後，放入一個END_OF_ITERATION標識符。另一方面，Reduce()線程則輪詢reduction_queue_中的元素。一旦發現隊列中有待處理參數信息時便調用實例化的solver中的ApplyUpdate()函數（例如，sgd_solver.cpp中的實現）。由於reduction_queue_本身是加鎖的隊列，因此能夠爲這兩個線程提供異步的一致性。

step()函數最後調用了Net::Finalize()函數，以確保step()函數線程在Reduce()線程之後結束。

3.2 ReduceAndUpdate()函數與buckets

當使用多GPU進行並行訓練時，用戶可以使用NetParameter中的reduce_buckets選項來優化權重更新過程。
該過程的處理是在ReduceAndUpdate()函數當中。簡單來講，reduce_buckets用來設置當reduction_queue_累計了多少個待處理參數個數的時候才調用一次權重更新函數，類似於批處理的概念。在caffe.proto的註釋中，作者建議reduce_buckets的默認參數對於大部分網絡來說是比較好的設置。

multi-gpu TestAll&Test

nvcaffe也支持訓練中多GPU的測試。在Test(const int test_net_id, const int iters, bool use_multi_gpu)函數中，如果使用了多GPU並行，則程序多執行一步同步各GPU測試結果的操作。TestAll(const int iters, bool use_multi_gpu)調用Test函數。在solver的構造函數中，use_multi_gpu的取值取決於Caffe::solver_count() > 1的結果，即是否使用了多GPU。