XLA

加速线性代数（XLA）是一种线性代数编译器，可优化不同硬件和框架上的模型运行时。

本指南将重点介绍如何使用 XLA 加速 TensorFlow 模型。

TensorFlow

XLA 可以在不更改任何源代码的情况下加速 TensorFlow 模型。它已经与 TensorFlow 库一起打包，并通过 tf.function 等任何图形创建函数中的 jit_compile 触发。

如果您正在使用 Keras 方法，例如 fit 和 predict，可以通过将 jit_compile=True 传递给 compile 来启用 XLA。

model.compile(jit_compile=True)

XLA 可用于加速任意 tf.function。

具有 TensorFlow 实现的模型，例如 GPT2、T5、OPT 和 Whisper，都与 XLA 兼容。加速取决于模型，但通常来说，Transformers 中的 TensorFlow 模型可以获得大约 100 倍的加速。

函数

TensorFlow 模型中的典型前向传播如下所示。要使用 XLA 运行前向传播，请使用 tf.function 包装模型并将 jit_compile=True 设置为 true。

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")]
)
# 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_fn = tf.function(model, jit_compile=True)
+ _ = xla_fn(random_inputs)

模型的默认 call 函数用于编译 XLA 图。但是，如果您想用 XLA 编译任何其他模型函数，请使用 tf.function 包装它们。

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

文本生成

您还可以使用 XLA 编译其他模型函数。例如，通过使用 tf.function 包装 generate() 来为文本生成启用 XLA。

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"]

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"

追踪

首次执行启用 XLA 的函数时，它会尝试在一个称为*追踪*的过程中推断计算图。这是一个耗时的步骤，但对该函数的任何后续调用都会快得多，因为它不必再次追踪计算图。

为了确保函数只被追踪一次，输入必须与构建图时的形状相同。这对于图像等固定输入形状通常不是问题，但对于文本等可变形状输入可能是一个问题。

解决这个问题的一种方法是填充您的文本，使其始终具有相同的形状。在 tokenizer 中配置填充选项，例如 pad_to_multiple_of。

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)

# Call 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 的更多信息。

一个演示与 XLA 兼容的编码器-解码器和仅解码器文本生成模型的笔记本。
使用 TensorFlow 和 XLA 加速文本生成博客文章比较了与 XLA 兼容的模型的基准，并友好的介绍了 TensorFlow 中的 XLA。
Hugging Face 如何通过 XLA 改进文本生成性能博客文章讨论了在 Transformers 中向 TensorFlow 模型添加 XLA 的设计理念。
图和 tf.function 简介指南。
使用 tf.function 获得更好的性能指南。
XLA 文档。

< > 在 GitHub 上更新