Hugging Face 的 TensorFlow 哲學

引言

儘管 PyTorch 和 JAX 的競爭日益激烈，TensorFlow 仍然是使用最廣泛的深度學習框架。它與其他兩個庫在某些非常重要的方面也有所不同。特別是，它與其高階 API Keras 以及資料載入庫 tf.data 緊密整合。

PyTorch 工程師（想象一下我在這裡盯著開放式辦公室的黑暗角落）傾向於將此視為一個需要克服的問題；他們的目標是弄清楚如何讓 TensorFlow 不礙事，以便他們可以使用他們習慣的低階訓練和資料載入程式碼。這完全是處理 TensorFlow 的錯誤方式！Keras 是一個很棒的高階 API。如果你在任何比幾個模組更大的專案中將其推開，當你意識到你需要它時，你最終會自己重現它的大部分功能。

作為經驗豐富、備受推崇且極具吸引力的 TensorFlow 工程師，我們希望利用尖端模型令人難以置信的功能和靈活性，但我們希望使用我們熟悉的工具和 API 來處理它們。這篇部落格文章將探討我們在 Hugging Face 所做的選擇，以實現這一目標，以及作為 TensorFlow 程式設計師，您應該從框架中獲得什麼。

插曲：30 秒瞭解 🤗

經驗豐富的使用者可以隨意略讀或跳過此部分，但如果這是您第一次接觸 Hugging Face 和 transformers，我應該首先向您概述該庫的核心思想：您只需按名稱請求一個預訓練模型，即可用一行程式碼獲取它。最簡單的方法是直接使用 TFAutoModel 類

from transformers import TFAutoModel

model = TFAutoModel.from_pretrained("bert-base-cased")

這一行程式碼將例項化模型架構並載入權重，為您提供原始著名 BERT 模型的精確複製品。然而，這個模型本身不會做太多事情——它缺乏輸出頭或損失函式。實際上，它是一個神經網路的“主幹”，在最後一個隱藏層之後就停止了。那麼如何給它加上輸出頭呢？很簡單，只需使用不同的 AutoModel 類。這裡我們載入 Vision Transformer (ViT) 模型並新增一個影像分類頭

from transformers import TFAutoModelForImageClassification

model_name = "google/vit-base-patch16-224"
model = TFAutoModelForImageClassification.from_pretrained(model_name)

現在我們的 model 有一個輸出頭，並且可選地，一個適合其新任務的損失函式。如果新的輸出頭與原始模型不同，則其權重將隨機初始化。所有其他權重將從原始模型載入。但我們為什麼要這樣做？我們為什麼要使用現有模型的主幹，而不是從頭開始構建我們需要的模型呢？

事實證明，在大量資料上預訓練的大模型幾乎是所有機器學習問題的更好起點，而不是簡單地隨機初始化權重的標準方法。這被稱為**遷移學習**，如果你仔細想想，這是有道理的——很好地解決文字任務需要一些語言知識，而很好地解決視覺任務需要一些影像和空間知識。沒有遷移學習時，機器學習之所以如此資料飢渴，僅僅是因為這種基本的領域知識必須為每個問題從頭開始重新學習，這需要大量的訓練樣本。然而，透過使用遷移學習，一個問題可以透過一千個訓練樣本來解決，而這在沒有遷移學習的情況下可能需要一百萬個樣本，並且通常具有更高的最終準確性。有關此主題的更多資訊，請檢視 Hugging Face 課程的相關部分！

然而，在使用遷移學習時，非常重要的一點是，您必須以與訓練期間處理輸入相同的方式處理模型輸入。這確保了當我們將模型的知識遷移到新問題時，模型需要重新學習的內容儘可能少。在 transformers 中，這種預處理通常由**分詞器**處理。分詞器可以像模型一樣載入，使用 AutoTokenizer 類。請務必載入與您要使用的模型匹配的分詞器！

from transformers import TFAutoModel, AutoTokenizer

# Make sure to always load a matching tokenizer and model!
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
model = TFAutoModel.from_pretrained("bert-base-cased")

# Let's load some data and tokenize it
test_strings = ["This is a sentence!", "This is another one!"]
tokenized_inputs = tokenizer(test_strings, return_tensors="np", padding=True)

# Now our data is tokenized, we can pass it to our model, or use it in fit()!
outputs = model(tokenized_inputs)

這當然只是該庫的冰山一角——如果您想了解更多，可以檢視我們的筆記本或我們的程式碼示例。在 keras.io 也有其他幾個該庫的實際應用示例！

至此，您已經瞭解了 transformers 中的一些基本概念和類。我上面寫的所有內容都是與框架無關的（除了 TFAutoModel 中的“TF”），但是當您真正想訓練和部署模型時，不同框架之間的差異就開始顯現了。這就引出了本文的重點：作為一名 TensorFlow 工程師，您對 transformers 應該有何期待？

哲學 #1：所有 TensorFlow 模型都應是 Keras 模型物件，所有 TensorFlow 層都應是 Keras 層物件。

對於一個 TensorFlow 庫來說，這幾乎是不言而喻的，但無論如何都值得強調。從使用者的角度來看，這種選擇最重要的影響是您可以直接在我們的模型上呼叫 Keras 方法，例如 fit()、compile() 和 predict()。

例如，假設您的資料已準備好並分詞，那麼從序列分類模型獲取 TensorFlow 預測就像這樣簡單

model = TFAutoModelForSequenceClassification.from_pretrained(my_model)
model.predict(my_data)

如果您想訓練該模型，它只是

model.fit(my_data, my_labels)

然而，這種便利並不意味著您僅限於我們開箱即用的任務。Keras 模型可以作為其他模型中的層來組合，因此如果您有一個涉及拼接五種不同模型的巨大銀河大腦想法，那麼沒有什麼能阻止您，除了您有限的 GPU 記憶體。也許您想將預訓練語言模型與預訓練視覺轉換器合併以建立混合模型，例如Deepmind 最近的 Flamingo，或者您想建立下一個像 Dall-E Mini Craiyon 那樣流行的文字到影像的轟動作品？這是一個使用 Keras 子類化的混合模型示例

class HybridVisionLanguageModel(tf.keras.Model):
  def __init__(self):
    super().__init__()
    self.language = TFAutoModel.from_pretrained("gpt2")
    self.vision = TFAutoModel.from_pretrained("google/vit-base-patch16-224")

  def call(self, inputs):
    # I have a truly wonderful idea for this
    # which this code box is too short to contain

哲學 #2：預設提供損失函式，但可以輕鬆更改。

在 Keras 中，訓練模型的標準方法是建立模型，然後使用最佳化器和損失函式對其進行 compile()，最後進行 fit()。使用 transformers 載入模型非常容易，但設定損失函式可能很棘手——即使對於標準語言模型訓練，您的損失函式也可能出奇地不明顯，並且一些混合模型具有極其複雜的損失。

我們的解決方案很簡單：如果您在沒有損失引數的情況下 compile()，我們將為您提供您可能想要的損失函式。具體來說，我們將為您提供一個與您的基礎模型和輸出型別都匹配的損失函式——如果您在沒有損失的情況下 compile() 一個基於 BERT 的掩碼語言模型，我們將為您提供一個掩碼語言建模損失，該損失可以正確處理填充和掩碼，並且只會計算損壞的令牌上的損失，完全匹配原始 BERT 訓練過程。如果出於某種原因，您真的、真的不希望您的模型在編譯時沒有任何損失，那麼只需在編譯時指定 loss=None 即可。

model = TFAutoModelForQuestionAnswering.from_pretrained("bert-base-cased")
model.compile(optimizer="adam")  # No loss argument!
model.fit(my_data, my_labels)

但同樣重要的是，一旦您想做更復雜的事情，我們希望讓您自由發揮。如果您為 compile() 指定一個損失引數，那麼模型將使用它而不是預設損失。當然，如果您建立自己的子類化模型，例如上面提到的 HybridVisionLanguageModel，那麼您可以透過編寫 call() 和 train_step() 方法來完全控制模型功能的各個方面。

哲學 實現細節 #3：標籤是靈活的

過去，一個困惑的來源是標籤究竟應該傳遞給模型的哪個地方。將標籤傳遞給 Keras 模型的標準方式是作為單獨的引數，或者作為 (輸入，標籤) 元組的一部分

model.fit(inputs, labels)

過去，我們要求使用者在使用預設損失時在輸入字典中傳遞標籤。原因是計算該特定模型損失的程式碼包含在 call() 正向傳播方法中。這可行，但對於 Keras 模型來說肯定不標準，並導致了一些問題，包括與標準 Keras 指標不相容，更不用說一些使用者的困惑。謝天謝地，這不再是必需的。我們現在建議以 Keras 的正常方式傳遞標籤，儘管舊方法出於向後相容性原因仍然有效。總的來說，許多以前很麻煩的事情現在應該對我們的 TensorFlow 模型“開箱即用”——試試看吧！

哲學 #4：你不應該自己編寫資料管道，尤其是對於常見任務

除了預訓練模型的巨大開放儲存庫 transformers 之外，還有一個巨大的開放資料集儲存庫 🤗 datasets，包括文字、視覺、音訊等。這些資料集可以輕鬆轉換為 TensorFlow 張量和 Numpy 陣列，從而方便地將其用作訓練資料。這裡有一個快速示例，展示了我們如何對資料集進行分詞並將其轉換為 Numpy。一如既往，請確保您的分詞器與您要訓練的模型匹配，否則事情會變得非常奇怪！

from datasets import load_dataset
from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
from tensorflow.keras.optimizers import Adam

dataset = load_dataset("glue", "cola")  # Simple text classification dataset
dataset = dataset["train"]  # Just take the training split for now

# Load our tokenizer and tokenize our data
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
tokenized_data = tokenizer(dataset["text"], return_tensors="np", padding=True)
labels = np.array(dataset["label"]) # Label is already an array of 0 and 1

# Load and compile our model
model = TFAutoModelForSequenceClassification.from_pretrained("bert-base-cased")
# Lower learning rates are often better for fine-tuning transformers
model.compile(optimizer=Adam(3e-5))

model.fit(tokenized_data, labels)

這種方法在有效時非常好，但對於更大的資料集，您可能會發現它開始成為一個問題。為什麼？因為分詞後的陣列和標籤必須完全載入到記憶體中，並且由於 Numpy 不處理“不規則”陣列，所以每個分詞後的樣本都必須填充到整個資料集中最長樣本的長度。這將使您的陣列變得更大，而且所有這些填充令牌也會減慢訓練速度！

作為一名 TensorFlow 工程師，此時您通常會轉向 tf.data 來構建一個數據管道，該管道將從儲存中流式傳輸資料，而不是將其全部載入到記憶體中。但這很麻煩，所以我們幫您搞定了。首先，讓我們使用 map() 方法將分詞器列新增到資料集中。請記住，我們的資料集預設是磁碟支援的——在您將它們轉換為陣列之前，它們不會載入到記憶體中！

def tokenize_dataset(data):
    # Keys of the returned dictionary will be added to the dataset as columns
    return tokenizer(data["text"])

dataset = dataset.map(tokenize_dataset)

現在我們的資料集有了我們想要的列，但我們如何用它來訓練呢？很簡單——用一個 tf.data.Dataset 封裝它，所有問題就解決了——資料是動態載入的，並且填充只應用於批次而不是整個資料集，這意味著我們需要的填充令牌要少得多

tf_dataset = model.prepare_tf_dataset(
    dataset,
    batch_size=16,
    shuffle=True
)

model.fit(tf_dataset)

為什麼 prepare_tf_dataset() 是模型的一個方法？很簡單：因為您的模型知道哪些列是有效輸入，並自動過濾掉資料集中不是有效輸入名稱的列！如果您希望對正在建立的 tf.data.Dataset 有更精確的控制，您可以使用更低階的 Dataset.to_tf_dataset() 代替。

哲學 #5：XLA 很棒！

XLA 是 TensorFlow 和 JAX 共享的即時編譯器。它將線性代數程式碼轉換為更最佳化的版本，執行速度更快，記憶體使用更少。它真的很酷，我們儘可能地支援它。它對於在 TPU 上執行模型非常重要，但它也為 GPU 甚至 CPU 提供了速度提升！要使用它，只需使用 jit_compile=True 引數編譯您的模型（這適用於所有 Keras 模型，而不僅僅是 Hugging Face 的模型）

model.compile(optimizer="adam", jit_compile=True)

我們最近在這個領域做了一些重大改進。最重要的是，我們更新了 generate() 程式碼以使用 XLA——這是一個迭代地從語言模型生成文字輸出的函式。這帶來了巨大的效能提升——我們舊的 TF 程式碼比 PyTorch 慢得多，但新程式碼比它快得多，並且與 JAX 的速度相似！有關更多資訊，請參閱我們關於 XLA 生成的部落格文章。

然而，XLA 除了生成之外還有其他用途！我們還進行了一些修復，以確保您可以使用 XLA 訓練模型，因此我們的 TF 模型在語言模型訓練等任務中達到了類似 JAX 的速度。

但重要的是要明確 XLA 的主要限制：XLA 期望輸入形狀是靜態的。這意味著如果您的任務涉及可變序列長度，您將需要為您傳遞給模型的每個不同輸入形狀執行一次新的 XLA 編譯，這會真正抵消效能優勢！您可以在我們的 TensorFlow 筆記本和上面提到的 XLA 生成部落格文章中看到我們如何處理此問題的一些示例。

哲學 #6：部署與訓練同樣重要

TensorFlow 擁有豐富的生態系統，尤其是在模型部署方面，這是其他更注重研究的框架所缺乏的。我們正在積極努力讓您能夠使用這些工具來部署您的整個模型以進行推理。我們特別關注支援 TF Serving 和 TFX。如果您對此感興趣，請檢視我們關於使用 TF Serving 部署模型的部落格文章！

然而，部署 NLP 模型的一個主要障礙是輸入仍然需要分詞，這意味著僅僅部署您的模型是不夠的。對 tokenizers 的依賴在許多部署場景中可能會令人煩惱，因此我們正在努力使分詞能夠嵌入到您的模型本身中，從而允許您僅部署單個模型工件來處理從輸入字串到輸出預測的整個管道。目前，我們只支援最常見的模型，如 BERT，但這是一個活躍的工作領域！如果您想嘗試，可以使用以下程式碼片段

# This is a new feature, so make sure to update to the latest version of transformers!
# You will also need to pip install tensorflow_text

import tensorflow as tf
from transformers import TFAutoModel, TFBertTokenizer


class EndToEndModel(tf.keras.Model):
    def __init__(self, checkpoint):
        super().__init__()
        self.tokenizer = TFBertTokenizer.from_pretrained(checkpoint)
        self.model = TFAutoModel.from_pretrained(checkpoint)

    def call(self, inputs):
        tokenized = self.tokenizer(inputs)
        return self.model(**tokenized)

model = EndToEndModel(checkpoint="bert-base-cased")

test_inputs = [
    "This is a test sentence!",
    "This is another one!",
]
model.predict(test_inputs)  # Pass strings straight to model!

結論：我們是一個開源專案，這意味著社群就是一切

做了一個很酷的模型？分享它！一旦您建立了賬戶並設定了憑據，就這麼簡單

model_name = "google/vit-base-patch16-224"
model = TFAutoModelForImageClassification.from_pretrained(model_name)

model.fit(my_data, my_labels)

model.push_to_hub("my-new-model")

您還可以使用 PushToHubCallback 在較長的訓練執行期間定期上傳檢查點！無論哪種方式，您都將獲得一個模型頁面和自動生成的模型卡，最重要的是，任何人都可以使用與載入任何現有模型完全相同的 API 來使用您的模型進行預測，或作為進一步訓練的起點

model_name = "your-username/my-new-model"
model = TFAutoModelForImageClassification.from_pretrained(model_name)

我認為，大型知名基礎模型與單個使用者微調的模型之間沒有區別，這體現了 Hugging Face 的核心信念——使用者能夠構建偉大事物。機器學習從來就不應該僅僅是少數幾家封閉公司提供的結果；它應該是一系列開放的工具、工件、實踐和知識的集合，不斷地被擴充套件、測試、批判和構建——一個集市，而不是一座大教堂。如果您碰巧有了新的想法、新的方法，或者您訓練了一個新的模型並取得了巨大成果，請告訴大家！

同樣，您是否缺少什麼？有 Bug 嗎？有什麼讓人不舒服的地方？有什麼應該很直觀但不是的？告訴我們吧！如果您願意（象徵性地）拿起鐵鍬開始修復，那就更好了，但即使您沒有時間或技能親自改進程式碼庫，也不要害羞地表達出來。通常，核心維護者可能會錯過問題，因為使用者沒有提出，所以不要認為我們一定知道某些事情！如果它困擾您，請在論壇上提問，或者如果您確定這是一個 Bug 或缺少重要功能，請提交問題。

當然，這些很多都是小細節，但用一個（相當笨拙的）詞來說，偉大的軟體是由成千上萬的小提交組成的。正是透過使用者和維護者的持續集體努力，開源軟體才得以改進。機器學習將在 2020 年代成為一個主要的社會問題，而開源軟體和社群的實力將決定它是否會成為一種開放和民主的力量，接受批評和重新評估，或者它是否被大型黑箱模型所主導，而這些模型的所有者不允許外人，甚至那些模型對其做出決策的人，看到他們寶貴的專有權重。所以不要害羞——如果有什麼不對勁，如果您對如何做得更好有想法，如果您想貢獻但不知道從何開始，那麼請告訴我們！

（如果你能製作一個表情包來嘲諷 PyTorch 團隊，那在你酷炫的新功能合併後，就更好了。）