TFX 最新博文：如何使用 TensorFlow 生態(tài)系統(tǒng)實現(xiàn)快速高效的 BERT 部署？

本文作者： skura

2020-03-14 07:45

導(dǎo)語：為什么TFX是BERT部署的完美搭檔？

Transformer 模型，特別是 BERT 模型，徹底改變了 NLP，并在諸如情感分析、實體抽取或問答等任務(wù)上開辟了新的天地。BERT 模型允許數(shù)據(jù)科學(xué)家站在巨人的肩膀上。當(dāng)企業(yè)用這些模型在大型語料庫中進(jìn)預(yù)訓(xùn)練時，數(shù)據(jù)科學(xué)家可以將遷移學(xué)習(xí)應(yīng)用到這些多用途的 transformer 模型中，并針對其領(lǐng)域特定的問題取得突破性的結(jié)果。

在 SAP 的 Concur 實驗室，研究人員希望使用 BERT 解決差旅費領(lǐng)域的新問題。他們想簡化 BERT 的推論。但不幸的是，他們一直沒找到完美的解決方案。

通過與 TensorFlow 團(tuán)隊合作，并使用他們的最新成果，Concur 實驗室最終實現(xiàn)了一致、簡單、快速的 BERT 模型推斷，可以在幾毫秒內(nèi)實現(xiàn)從原始文本的分類預(yù)測。

那么，各種 TensorFlow 庫和組件如何幫助研究人員達(dá)到這個里程碑？作為 SAP Concur 實驗室的高級機(jī)器學(xué)習(xí)工程師，Hannes Hapke 寫了一篇博文分享了他們的做法。TFX 團(tuán)隊的研究人員也參與了這篇博文的編輯。

這篇博客文章將向你概述如何使用 TensorFlow 生態(tài)系統(tǒng)來實現(xiàn)可伸縮、快速、高效的 BERT 部署。全文編譯如下。

本文討論的方法將允許開發(fā)人員使用 TensorFlow Extended（TFX）v0.21 或更高版本來訓(xùn)練 TensorFlow 模型。但是，對 tf.text 操作的支持尚未包含在 TensorFlow Service（v2.1）的當(dāng)前版本中，而是包含在nightly docker版本中，并將包含在 v2.2 版本中。

如果你想查看完整的示例，請查看 Colab notebook。它展示了整個 TensorFlow 擴(kuò)展（TFX）管道以生成可部署的 BERT 模型，并將預(yù)處理步驟作為模型圖的一部分。

BERT 部署的現(xiàn)狀

transformer 模型的最新進(jìn)展令人吃驚。但不幸的是，將這些模型用在生產(chǎn)中卻并不簡單或者完美。理想情況下，我們希望將原始文本發(fā)送到服務(wù)器，但是 BERT 模型需要對輸入文本進(jìn)行預(yù)處理，然后才能從實際模型中獲得預(yù)測。

有一些解決方案通過在客戶端對文本進(jìn)行預(yù)處理來解決這個問題，還有的解決方案是，在服務(wù)器端實現(xiàn)了一個中間步驟來處理輸入數(shù)據(jù)。這兩個方法都不太正確，因為它們需要額外的部署協(xié)調(diào)，降低了推理的效率。

TFX 最新博文：如何使用 TensorFlow 生態(tài)系統(tǒng)實現(xiàn)快速高效的 BERT 部署？

圖 1：當(dāng)前的 BERT 部署

什么是最佳部署？

在部署模型時，方法越簡單越好。我們希望部署 transformer 模型，并將預(yù)處理作為模型圖的一部分。由于預(yù)處理被集成到模型圖中，我們可以簡單地將單個模型部署到模型服務(wù)器，刪除其他部署依賴項，并充分利用模型服務(wù)器的優(yōu)點。例如，我們可以批處理預(yù)測請求，以優(yōu)化推理硬件的使用。

利用 TensorFlow 生態(tài)系統(tǒng)部署 BERT

TensorFlow 對我們來說是一個非常高效的框架，因為它不僅是一個機(jī)器學(xué)習(xí)框架，而且還提供了一個具有廣泛的支持包和工具的生態(tài)系統(tǒng)。TensorFlow Serving 對我們來說非常有用。它提供了簡單、一致、可伸縮的模型部署。

我們密切關(guān)注的另一個系統(tǒng)項目是 TensorFlow Transform。它為我們提供了一個將模型預(yù)處理步驟構(gòu)建為圖的機(jī)會，然后，我們可以將這些圖與實際的深度學(xué)習(xí)模型一起導(dǎo)出。TensorFlow Transform 要求所有預(yù)處理步驟都表示為 TensorFlow 操作。正因為如此，TensorFlow 文本的最新發(fā)展對我們來說也非常有幫助。RaggedTensors 的實現(xiàn)不僅開啟了新的應(yīng)用，而且，其庫還提供了實現(xiàn)自然語言預(yù)處理步驟所需的功能。

在 TensorFlowWorld 2019 上展示了 TensorFlow Text 的一個新功能，那就是完全實現(xiàn)一個 BERT 標(biāo)記器。因此，我們可以用幾行 TensorFlow 代碼來表示預(yù)處理步驟。我們還通過使用另一個 TensorFlow 工具：TensorFlow Extended（TFX），實現(xiàn)了模型管道和部署一致的目標(biāo)。TFX 允許我們以可重復(fù)的方式表達(dá)整個 ML 管道，因此有助于我們部署一致的機(jī)器學(xué)習(xí)模型。

TFX 最新博文：如何使用 TensorFlow 生態(tài)系統(tǒng)實現(xiàn)快速高效的 BERT 部署？

圖 2：附 tf.Text 的 TFX 管道

用 TensorFlow ops 編寫預(yù)處理步驟

理想的模型部署接受原始文本作為模型的輸入，并提供模型預(yù)測作為結(jié)果。簡化 BERT 部署的關(guān)鍵是將預(yù)處理步驟表示為 TensorFlow ops。BERT 模型要求將原始輸入文本標(biāo)記為標(biāo)記 id、輸入掩碼的伴隨數(shù)據(jù)結(jié)構(gòu)和生成的輸入類型 id。借助 TensorFlow Text，我們現(xiàn)在可以用更少的代碼來實現(xiàn)這一點。在接下來的文章中，我們將討論從原始文本到特定于 BERT 的數(shù)據(jù)結(jié)構(gòu)轉(zhuǎn)換的細(xì)節(jié)，包括添加特定于 BERT 的標(biāo)記。

vocab_file_path = load_bert_layer().resolved_object.vocab_file.asset_path
bert_tokenizer = text.BertTokenizer(vocab_lookup_table=vocab_file_path,
token_out_type=tf.int64,
lower_case=do_lower_case)
...
input_word_ids = tokenize_text(text)
input_mask = tf.cast(input_word_ids > 0, tf.int64)
input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

zeros_dims = tf.stack(tf.shape(input_mask))
input_type_ids = tf.fill(zeros_dims, 0)
input_type_ids = tf.cast(input_type_ids, tf.int64)

TFX 最新博文：如何使用 TensorFlow 生態(tài)系統(tǒng)實現(xiàn)快速高效的 BERT 部署？

圖 3：BERT 標(biāo)記器

利用 TensorFlow Transform 和上述代碼，可以將預(yù)處理圖與訓(xùn)練好的 TensorFlow 模型一起導(dǎo)出。隨著 TensorFlow Serving 的最新更新，我們部署的 BERT 模型現(xiàn)在可以接受原始文本作為輸入。在這里，沒有其他部署依賴項！

使用 TensorFlow Transform 為我們提供了一些實際的好處。一方面，我們可以在組織上劃分?jǐn)?shù)據(jù)預(yù)處理和模型架構(gòu)工作之間的職責(zé)。另一方面，我們可以方便地調(diào)試、測試和生成預(yù)處理輸出的統(tǒng)計信息。transform 組件輸出轉(zhuǎn)換后的訓(xùn)練集，作為 TFRecords，這些訓(xùn)練集易于檢查。

在 Transform 輸出的「調(diào)試」過程中，我們發(fā)現(xiàn)了一些小錯誤，這些錯誤不會導(dǎo)致模型的訓(xùn)練失敗，但可能會影響其性能。例如，[SEP] 標(biāo)記中的偏移量。從技術(shù)上來講，這里不需要 TensorFlow Transform。由于每個示例的預(yù)處理獨立于整個語料庫進(jìn)行，我們可以很容易地將其直接構(gòu)建到模型圖中。但是，我們發(fā)現(xiàn)用這種方法構(gòu)建和調(diào)試管道會更容易。

TFX 最新博文：如何使用 TensorFlow 生態(tài)系統(tǒng)實現(xiàn)快速高效的 BERT 部署？

圖 4：BERT 層

理想的部署？

簡化開發(fā)

通過使用各種 TensorFlow 工具，我們能夠以簡單而簡潔的方式部署 BERT 模型。將預(yù)處理步驟集成到模型圖中可以減少訓(xùn)練數(shù)據(jù)和推理數(shù)據(jù)之間的傾斜風(fēng)險。部署的模型不需要額外的客戶機(jī)或服務(wù)器依賴項，這進(jìn)一步降低了模型錯誤的風(fēng)險。我們可以將我們的 BERT 模型與 TensorFlow Serving 一致地部署，同時還可以利用模型優(yōu)化——如批處理推斷。

推理性能

我們最初的性能測試看起來很有希望。在我們的演示 BERT 模型圖上的推論，包含預(yù)處理步驟和每個預(yù)測模型。這些模型平均約 15.5 ms，它們是在單個 V100 GPU、最多 128 個標(biāo)記、gRPC 請求上測量的，由 GPU 的非優(yōu)化 TensorFlow 服務(wù)構(gòu)建的無基礎(chǔ) BERT 模型。這和以前在客戶端使用 BERT 標(biāo)記化的部署以及使用 TensorFlow Serving 托管的分類模型的平均推斷時間大致相同。當(dāng)然，根據(jù)你的機(jī)器及其型號的不同，你可能會看到不同的結(jié)果。

了解更多信息

如果你想深入了解代碼，請查看 Colab notebook，其中有一個使用預(yù)先訓(xùn)練的 BERT 模型實現(xiàn)情感分類模型的示例。如果你想試用我們的演示部署，請查看 Concur 實驗室的演示頁面，上面展示了我們的情感分類項目。

如果你對 TensorFlow Extended（TFX）和 TensorFlow Transform 的內(nèi)部工作感興趣，請進(jìn)入 TFX 用戶指南并查看即將出版的 O'Reilly 出版物「Building Machine Learning Pipelines, Automating Model Life Cycles With TensorFlow」。

要了解有關(guān) TFX 的更多信息，請訪問 TFX 網(wǎng)站，加入 TFX 討論組，瀏覽 TFX 博客中的其他文章，或者在 YouTube 上觀看我們的 TFX 播放列表，并訂閱 TensorFlow 頻道。