1. Introduction

HugeCTR 能夠高效地利用 GPU 來進行推薦系統(tǒng)的訓練，為了使它還能直接被其他 DL 用戶，比如 TensorFlow 所直接使用，我們開發(fā)了 SparseOperationKit (SOK)，來將 HugeCTR 中的高級特性封裝為 TensorFlow 可直接調(diào)用的形式，從而幫助用戶在 TensorFlow 中直接使用 HugeCTR 中的高級特性來加速他們的推薦系統(tǒng)。

圖 1. SOK embedding 工作流程

SOK 以數(shù)據(jù)并行的方式接收輸入數(shù)據(jù)，然后在 SOK 內(nèi)部做黑盒式地模型轉(zhuǎn)換，最后將計算結(jié)果以數(shù)據(jù)并行的方式傳遞給初始 GPU。這種方式可以盡可能少地修改用戶已有的代碼，以更方便、快捷地在多個 GPU 上進行擴展。

SOK 不僅僅是加速了 TensorFlow 中的算子，而是根據(jù)業(yè)界中的實際需求提供了對應(yīng)的新解決方案，比如說 GPU HashTable。SOK 可以與 TensorFlow 1.15 和 TensorFlow 2.x 兼容使用；既可以使用 TensorFlow 自帶的通信工具，也可以使用 Horovod 等第三方插件來作為 embedding parameters 的通信工具。

2. TF2 Comparison/Performance

使用 MLPerf 的標準模型 DLRM 來對 SOK 的性能進行測試。

圖 2. SOK 性能測試數(shù)據(jù)

相比于 NVIDIA 的 DeepLearning Examples，使用 SOK 可以獲得更快的訓練速度以及更高的吞吐量。

3. API

SOK 提供了簡潔的、類 TensorFlow 的 API；使用 SOK 的方式非常簡單、直接；讓用戶通過修改幾行代碼就可以使用 SOK。

1. 定義模型結(jié)構(gòu)

左側(cè)是使用 TensorFlow 的 API 來搭建模型，右側(cè)是使用 SOK 的 API 來搭建相同的模型。使用 SOK 來搭建模型的時候，只需要將 TensorFlow 中的 Embedding Layer 替換為 SOK 對應(yīng)的 API 即可。

2. 使用 Horovod 來定義 training loop

同樣的，左側(cè)是使用 TensorFlow 來定義 training loop，右側(cè)是使用 SOK 時，training loop 的定義方式?？梢钥吹剑褂?SOK 時，只需要對 Embedding Variables 和 Dense Variables 進行分別處理即可。其中，Embedding Variables 部分由 SOK 管理，Dense Variables 由 TensorFlow 管理。

3. 使用 tf.distribute.MirroredStrategy 來定義 training loop

類似的，還可以使用 TensorFlow 自帶的通信工具來定義 training loop。

4. 開始訓練

在開始訓練過程時，使用 SOK 與使用 TensorFlow 時所用代碼完全一致。

4. 結(jié)語

SOK 將 HugeCTR 中的高級特性包裝為 TensorFlow 可以直接使用的模塊，通過修改少數(shù)幾行代碼即可在已有模型代碼中利用上 HugeCTR 的先進設(shè)計。

審核編輯 ：李倩