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 团队，那在你酷炫的新功能合并后，就更好了。）