Hugging Face 的 logo

Transformers 文档

TensorFlow 模型的 XLA 集成

Hugging Face's logo
加入 Hugging Face 社区

并获取增强文档体验

开始使用

TensorFlow 模型的 XLA 集成

加速线性代数，简称 XLA，是一种用于加速 TensorFlow 模型运行时的编译器。来自 官方文档

XLA（加速线性代数）是一种针对线性代数的领域特定编译器，可以加速 TensorFlow 模型，并且可能无需更改源代码。

在 TensorFlow 中使用 XLA 非常简单 - 它打包在 tensorflow 库中，并且可以通过任何创建图形的函数（例如 tf.function）中的 jit_compile 参数触发。当使用 Keras 方法（如 fit() 和 predict()）时，只需将 jit_compile 参数传递给 model.compile() 即可启用 XLA。但是，XLA 不仅限于这些方法 - 它还可以用于加速任何任意 tf.function。

🤗 Transformers 中的几个 TensorFlow 方法已被重写为与 XLA 兼容，包括针对诸如 GPT2、T5 和 OPT 等模型的文本生成，以及针对诸如 Whisper 等模型的语音处理。

虽然加速的具体幅度很大程度上取决于模型，但在 🤗 Transformers 中的 TensorFlow 文本生成模型中，我们观察到大约 100 倍的加速。本文档将解释如何为这些模型使用 XLA 以获得最大的性能。如果您有兴趣了解有关基准测试和我们 XLA 集成背后的设计理念的更多信息，我们还将提供指向其他资源的链接。

使用 XLA 运行 TF 函数

让我们考虑 TensorFlow 中的以下模型

import tensorflow as tf

model = tf.keras.Sequential(
    [tf.keras.layers.Dense(10, input_shape=(10,), activation="relu"), tf.keras.layers.Dense(5, activation="softmax")]
)

上述模型接受维度为 (10, ) 的输入。我们可以使用该模型运行前向传递，如下所示

# Generate random inputs for the model.
batch_size = 16
input_vector_dim = 10
random_inputs = tf.random.normal((batch_size, input_vector_dim))

# Run a forward pass.
_ = model(random_inputs)

为了使用 XLA 编译的函数运行前向传递，我们需要执行以下操作

xla_fn = tf.function(model, jit_compile=True)
_ = xla_fn(random_inputs)

模型的默认 call() 函数用于编译 XLA 图。但是，如果您想将任何其他模型函数编译到 XLA 中，也可以使用以下方法

my_xla_fn = tf.function(model.my_xla_fn, jit_compile=True)

从 🤗 Transformers 运行使用 XLA 的 TF 文本生成模型

要在 🤗 Transformers 中启用 XLA 加速生成，您需要安装最新版本的 transformers。您可以通过运行以下命令来安装它：

pip install transformers --upgrade

然后您可以运行以下代码：

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

# Will error if the minimal version of Transformers is not installed.
from transformers.utils import check_min_version

check_min_version("4.21.0")


tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")
input_string = ["TensorFlow is"]

# One line to create an XLA generation function
xla_generate = tf.function(model.generate, jit_compile=True)

tokenized_input = tokenizer(input_string, return_tensors="tf")
generated_tokens = xla_generate(**tokenized_input, num_beams=2)

decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")
# Generated -- TensorFlow is an open-source, open-source, distributed-source application # framework for the

如您所见，在 generate() 上启用 XLA 只需一行代码。其余代码保持不变。但是，上述代码片段中有一些特定于 XLA 的注意事项。您需要了解这些才能实现 XLA 可以带来的加速效果。我们在下一节中讨论这些内容。

需要注意的事项

当您第一次执行启用 XLA 的函数（如上面的 xla_generate()）时，它将在内部尝试推断计算图，这是一个耗时的过程。此过程称为“跟踪”。

您可能会注意到生成时间不快。鉴于函数的输入与最初构建计算图的输入具有相同的形状，因此 xla_generate()（或任何其他启用 XLA 的函数）的后续调用无需推断计算图。虽然这对具有固定输入形状的模态（例如图像）来说不是问题，但如果您正在处理可变输入形状的模态（例如文本），则必须注意这一点。

为了确保 xla_generate() 始终使用相同的输入形状进行操作，您可以在调用分词器时指定 padding 参数。

import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")
input_string = ["TensorFlow is"]

xla_generate = tf.function(model.generate, jit_compile=True)

# Here, we call the tokenizer with padding options.
tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")

generated_tokens = xla_generate(**tokenized_input, num_beams=2)
decoded_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
print(f"Generated -- {decoded_text}")

这样，您可以确保 xla_generate() 的输入将始终接收与其跟踪的形状相同的输入，从而导致生成时间加快。您可以使用以下代码进行验证

import time
import tensorflow as tf
from transformers import AutoTokenizer, TFAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2", padding_side="left", pad_token="</s>")
model = TFAutoModelForCausalLM.from_pretrained("openai-community/gpt2")

xla_generate = tf.function(model.generate, jit_compile=True)

for input_string in ["TensorFlow is", "TensorFlow is a", "TFLite is a"]:
    tokenized_input = tokenizer(input_string, pad_to_multiple_of=8, padding=True, return_tensors="tf")
    start = time.time_ns()
    generated_tokens = xla_generate(**tokenized_input, num_beams=2)
    end = time.time_ns()
    print(f"Execution time -- {(end - start) / 1e6:.1f} ms\n")

在 Tesla T4 GPU 上，您可以期望输出如下所示：

Execution time -- 30819.6 ms

Execution time -- 79.0 ms

Execution time -- 78.9 ms

第一次调用 xla_generate() 由于跟踪而很耗时，但后续调用速度要快几个数量级。请记住，在任何时候更改生成选项都会触发重新跟踪，从而导致生成时间变慢。

本文档没有涵盖 🤗 Transformers 提供的所有文本生成选项。我们鼓励您阅读文档以了解高级用例。

其他资源

在这里，如果您想更深入地了解 🤗 Transformers 和一般的 XLA，我们为您提供了一些其他资源。

< > 在 GitHub 上更新

←调试 使用 `torch.compile()` 优化推理→