TensorNet是什么？

TensorNet是一個構建在TensorFlow之上針對廣告推薦等大規模稀疏場景優化的分布式訓練框架。TensorNet的目標是讓所有使用TensorFlow的開發者可以快速的、方便的訓練出稀疏參數超過百億的超大模型。

訓練帶有大規模稀疏特征模型的主要挑戰

在廣告、搜索、推薦等場景下的深度模型都會存在大量的高維離散稀疏特征，訓練帶有高維稀疏特征的模型主要有兩個問題：

1. 訓練樣本規模大。比如對于360廣告場景會有超過100TB的訓練數據。
2. 模型參數多。比如對于360廣告場景會有超過100億的參數。

使用單機模式訓練模型速度慢，耗時長，嚴重制約了模型的迭代速度，使用分布式訓練已經成為業界標準。

使用TensorFlow訓練稀疏特征模型的主要問題

TensorFlow是最受開發者歡迎的深度學習訓練框架，但是TensorFlow對訓練帶有大規模稀疏特征的模型不太友好，主要問題有：

1. TensorFlow支持的特征維度有限。一般的，TensorFlow需要對每一個特征定義一個矩陣，這個矩陣受限于內存，往往不能太大。
2. TensorFlow2.x對參數服務器的支持較少，使用同步訓練的模式會比較慢。

TensorNet——基于TensorFlow的專為大規模稀疏特征模型優化的分布式訓練框架

TensorNet在復用TensorFlow的所有功能的基礎之上，專門定制使其支持大規模稀疏特征模型的訓練。TensorNet的主要提升包括：

1. 使TensorFlow支持的稀疏特征的維度接近于無限。
2. 使TensorFlow2.2 keras支持基于參數服務器的異步訓練模式，極大的提升了訓練速度。在360真實業務場景下我們將原來的離線訓練時間由3.5小時提升到了25分鐘。
3. 配合TensorNet通過split graph的方法可以對在線推理的性能進行優化。在360真實場景測試中我們發現有近 35% 的性能提升。

TensorNet分布式訓練架構

TensorNet支持異步和同步模式訓練。異步模式在僅有CPU的集群中速度提升十分顯著，同步模式在網卡速度超過100GbE的GPU集群中表現突出。

TensorNet異步訓練架構

在僅有CPU的集群中使用參數服務器的異步訓練模式是訓練模型速度最快的方法，TensorNet異步訓練架構與TensorFlow的異步訓練架構有很大的區別：

1. TensorNet將sparse參數和與dense參數分別使用不同的parameter server管理。
2. TensorNet不設單獨的parameter server節點。在每個worker中都會維護一個sparse paramter server和dense parameter server。這省去了開發人員管理ps節點和worker節點的不少麻煩。
3. TensorNet對sparse參數使用分布式哈希表按照哈希值均勻分布不同的節點上。這相較于TensorFlow需要讓開發者根據自身情況將tensor分布在不同的ps節點上的方法更加靈活，這不僅減小了節點通信熱點的概率，還減輕了開發者的工作量。
4. TensorNet將模型的所有dense參數合并后使用分布式數組切分到不同的機器上，每次pull和push參數的時候只有一次網絡請求。相較于TensorFlow對每個tensor都有一次網絡請求的方法極大的減少了網絡請求的次數從而提升了模型訓練的速度。

TensorNet同步訓練架構

TensorNet同步訓練架構基本與TensorFlow的MultiWorkerMirroredStrategy架構一致，主要區別如下：

1. TensorNet使用單獨的sparse parameter server節點保存所有sparse參數。通過parameter server可以解決TensorFlow支持的sparse特征維度不能太大的問題。
2. TensorNet對sparse參數做了特殊的定制化的同步。TensorNet在訓練時由于每個batch內的sparse參數的`IndexedSlices`指向的內容與TensorFlow默認的不同，我們對此做了定制化的同步。

TensorNet核心優化

TensorNet最核心的優化是將模型的embedding tensor優化到了最小。如下圖所示，對于最簡單的wide&deep模型，如果在一個廣告系統中有3億用戶，那么就需要定義一個維度為3億的embedding矩陣，在訓練模型時需要在這個3億維的矩陣上做embedding_lookup得到當前batch內的用戶的embedding信息，近而在embedding之上做更加復雜的操作。

TensorFlow中的實現在高維稀疏場景下，embedding矩陣太大，占用內存多。很顯然當特征較多的時候單機無法存儲整個模型。

TensorNet使用一個較小的，可以容納特征在一個batch內所有數據的embedding矩陣代替TensorFlow默認實現中需要定義的較大的embedding矩陣。

如下圖所示，在batch_size設置為1024的場景下，對于用戶id特征，在TensorNet中只需要定義一個維度為1024的embedding矩陣，TensorNet的主要處理步驟如下：

1. 定義模型時定義userid的embedding矩陣的維度為一個batch內所有用戶id個數的最大值。
2. 訓練模型時得到當前batch內的所有用戶id。
3. 將用戶id排序，并按照先后順序為每個userid分配索引，索引從0開始，對應為下圖中的virtual sparse feature。
4. 使用userid從parameter server中獲取相應的embedding向量，然后按照其對應的索引放置到embedding矩陣中。
5. 使用轉換后的virtual sparse feature作為模型的輸入。

從上述可見，TensorNet由于極大的減小了模型所需要的embedding矩陣，從而可以極大的減小內存的開銷，以及通過parameter server的方式使得稀疏特征的維度可以支持到接近無限維，從而可以極大的提升模型的刻畫能力。

TensorNet Inference優化

由于TensorNet只更改了模型的第一層，從而模型的inference也變得極其簡單。

在使用TensorNet構造模型的時候，可以將模型切分為兩部分，如下圖所示，embedding_lookup_graph只在離線訓練時使用，在線inference時只需要將sparse embedding導出成字典供inference_graph作為輸入即可，具體的請參考以下系列文章：

1. 為inference準備——模型切分: https://github.com/Qihoo360/tensornet/blob/master/doc/tutorial/03-split-to-sub-graph.ipynb

2. 使用XLA方式進行在線預估: https://github.com/Qihoo360/tensornet/blob/master/doc/tutorial/04-deploy-tf-graph-online.ipynb

3. sparse embedding字典導出: https://github.com/Qihoo360/tensornet/blob/master/doc/tutorial/05-export-sparse-feature-embedding.ipynb

TensorNet中split graph inference方案

在360內部場景中我們測試發現通過split graph配合XLA AOT的方法性能提升近35%。

TensorNet開源及使用

TensorNet已經成功落地應用到了360廣告ctr預估相關的場景中，并取得了顯著的效果，我們已將代碼、文檔及我們在360廣告的應用經驗全部整理到了項目中，歡迎關注。

tensorNet主頁：https://github.com/Qihoo360/TensorNet

tensornet快速上手：https://github.com/Qihoo360/tensornet/blob/master/doc/tutorial/01-begin-with-wide-deep.ipynb

更多文檔請看：https://github.com/Qihoo360/TensorNet/README.md

聯系方式：張彥升(zhangyansheng@360.cn),姚磊(yaolei@360.cn)

微信交流群：

https://weixin.qq.com/g/AQYAAB0H54Ncc-QyDTtKD2l-4V19WXGdpbUfFa6QqwS8NphK5kZ5tWAD0_c342jf (二維碼自動識別)