TPU

TPU (张量处理单元) 是一种硬件，旨在加速张量计算，用于训练和推理。TPU 通常通过 Google 云服务访问，但也可以从 Google Colab 或 Kaggle 免费获得较小的 TPU。

本指南重点介绍如何从 Google Colab 在 TPU 上训练用于序列分类的 Keras 模型。请确保通过转到运行时 > 更改运行时类型并选择 TPU 来启用 TPU 运行时。

运行以下命令安装最新版本的 Transformers 和 Datasets。

!pip install --U transformers datasets

创建 tf.distribute.cluster_resolver.TPUClusterResolver 的实例，然后连接到远程集群并初始化 TPU。

import tensorflow as tf

resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

有多种分发策略可在多个 TPU 上运行模型。tpu.distribute.TPUStrategy 提供同步分布式训练。

strategy = tf.distribute.TPUStrategy(resolver)

加载并分词数据集 - 此示例使用 GLUE 基准测试中的 CoLA - 并将所有样本填充到最大长度，以便更轻松地作为数组加载并避免 XLA 编译问题。

from transformers import AutoTokenizer
from datasets import load_dataset
import numpy as np

dataset = load_dataset("glue", "cola")["train"]
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased")

train_data = tokenizer(
    dataset["sentence"],
    padding="max_length",
    truncation=True,
    max_length=128,
    return_tensors="np",
)
train_data = dict(train_data)
train_labels = np.array(dataset["label"])

模型必须在 Strategy.scope 内部创建，以便在每个 TPU 设备上复制模型层。

from transformers import TFAutoModelForSequenceClassification

with strategy.scope():
    model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)
    model.compile(optimizer="adam")

TPU 仅接受 tf.data.Dataset 输入，这与 Keras fit 方法不同，后者接受更广泛的输入。

BATCH_SIZE = 8 * strategy.num_replicas_in_sync

tf_dataset = tf.data.Dataset.from_tensor_slices((train_data, train_labels))
tf_dataset = tf_dataset.shuffle(len(tf_dataset))
tf_dataset = tf_dataset.batch(BATCH_SIZE, drop_remainder=True)

最后，调用 fit 开始训练。

model.fit(tf_dataset)

大型数据集

上面创建的数据集将每个样本填充到最大长度并将整个数据集加载到内存中。如果您正在处理更大的数据集，这可能不可行。在大型数据集上训练时，您可能需要创建 tf.TFRecord 或流式传输数据。

tf.TFRecord

tf.TFRecord 是用于存储训练数据的标准 tf.data 格式。对于非常大的训练作业，值得预处理您的数据并将其存储为 tf.TFRecord 格式，并在其上构建 tf.data 管道。请参阅下表，以帮助您决定 tf.TFRecord 是否对您有帮助。

优点	缺点
适用于所有 TPU 实例	与云存储相关的成本
支持海量数据集和巨大的吞吐量	某些数据类型（图像）可能占用大量存储空间
适用于在整个 TPU Pod 上进行训练
预处理是提前完成的，最大限度地提高了训练速度

在将数据集写入 tf.TFRecord 之前对其进行预处理和分词，以避免每次加载数据时都进行写入。

训练时增强是一个例外，因为在写入 tf.TFRecord 后应用的增强会导致每个 epoch 都进行相同的增强。相反，在加载数据的 tf.data 管道中应用增强。

在实践中，您可能无法将整个数据集加载到内存中。一次加载数据集的一部分并将其转换为 TFRecord，然后重复此操作，直到整个数据集都为 TFRecord 格式。然后，您可以使用所有文件的列表来创建 TFRecordDataset。下面的示例演示了单个文件的简单性。

tokenized_data = tokenizer(
    dataset["sentence"],
    padding="max_length",
    truncation=True,
    max_length=128,
    return_tensors="np",
)
labels = dataset["label"]

with tf.io.TFRecordWriter("dataset.tfrecords") as file_writer:
    for i in range(len(labels)):
        features = {
            "input_ids": tf.train.Feature(
                int64_list=tf.train.Int64List(value=tokenized_data["input_ids"][i])
            ),
            "attention_mask": tf.train.Feature(
                int64_list=tf.train.Int64List(value=tokenized_data["attention_mask"][i])
            ),
            "labels": tf.train.Feature(
                int64_list=tf.train.Int64List(value=[labels[i]])
            ),
        }
        features = tf.train.Features(feature=features)
        example = tf.train.Example(features=features)
        record_bytes = example.SerializeToString()
        file_writer.write(record_bytes)

使用保存的文件名构建 TFRecordDataset 以加载它。

def decode_fn(sample):
    features = {
        "input_ids": tf.io.FixedLenFeature((128,), dtype=tf.int64),
        "attention_mask": tf.io.FixedLenFeature((128,), dtype=tf.int64),
        "labels": tf.io.FixedLenFeature((1,), dtype=tf.int64),
    }
    return tf.io.parse_example(sample, features)

# TFRecordDataset can handle gs:// paths
tf_dataset = tf.data.TFRecordDataset(["gs://matt-tf-tpu-tutorial-datasets/cola/dataset.tfrecords"])
tf_dataset = tf_dataset.map(decode_fn)
tf_dataset = tf_dataset.shuffle(len(dataset)).batch(BATCH_SIZE, drop_remainder=True)
tf_dataset = tf_dataset.apply(
    tf.data.experimental.assert_cardinality(len(labels) // BATCH_SIZE)
)

现在可以将数据集传递给 fit 方法。

model.fit(tf_dataset)

从原始数据流式传输

数据可以以其原始格式存储，并在加载数据时在 tf.data 管道中进行预处理。对于许多具有复杂分词方案的模型，此方法不受支持，但某些模型（如 BERT）受支持，因为它们的分词可以编译。请参阅下表，以帮助您决定此方法是否对您有帮助。

优点	缺点
适用于原始格式的高度压缩大数据（图像、音频）	需要编写完整的预处理管道
如果原始数据在公共云存储桶中可用，则很方便	复杂的即时预处理可能会降低吞吐量
如果数据存储在 Google Cloud 中，则适用于所有 TPU 实例	如果数据尚未在云存储中，则必须将其放入云存储中
	不太适合文本数据，因为编写分词管道很困难（文本请使用 `TFRecord`）

下面的示例演示了图像模型的流式数据。

加载图像数据集并获取底层图像文件路径和标签的列表。

from datasets import load_dataset

image_dataset = load_dataset("beans", split="train")
filenames = image_dataset["image_file_path"]
labels = image_dataset["labels"]

将数据集中的本地文件名转换为 Google Cloud Storage 中的 gs:// 路径。

# strip everything but the category directory and filenames
base_filenames = ['/'.join(filename.split('/')[-2:]) for filename in filenames]
# prepend the Google Cloud base path to everything instead
gs_paths = ["gs://matt-tf-tpu-tutorial-datasets/beans/"+filename for filename in base_filenames]

# create tf_dataset
tf_dataset = tf.data.Dataset.from_tensor_slices(
    {"filename": gs_paths, "labels": labels}
)
tf_dataset = tf_dataset.shuffle(len(tf_dataset))

像 AutoImageProcessor 这样的 Transformers 预处理类与框架无关，无法由 tf.data 编译到管道中。为了解决这个问题，从 AutoImageProcessor 获取归一化值（mean 和 std），并在 tf.data 管道中使用它们。

from transformers import AutoImageProcessor

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
image_size = (processor.size["height"], processor.size["width"])
image_mean = processor.image_mean
image_std = processor.image_std

使用这些归一化值创建一个函数来加载和预处理图像。

BATCH_SIZE = 8 * strategy.num_replicas_in_sync

def decode_fn(sample):
    image_data = tf.io.read_file(sample["filename"])
    image = tf.io.decode_jpeg(image_data, channels=3)
    image = tf.image.resize(image, image_size)
    array = tf.cast(image, tf.float32)
    array /= 255.0
    array = (array - image_mean) / image_std
    array = tf.transpose(array, perm=[2, 0, 1])
    return {"pixel_values": array, "labels": sample["labels"]}

tf_dataset = tf_dataset.map(decode_fn)
tf_dataset = tf_dataset.batch(BATCH_SIZE, drop_remainder=True)
print(tf_dataset.element_spec)

现在可以将数据集传递给 fit 方法。

from transformers import TFAutoModelForImageClassification

with strategy.scope():
    model = TFAutoModelForImageClassification.from_pretrained(image_model_checkpoint)
    model.compile(optimizer="adam")

model.fit(tf_dataset)

使用 prepare_tf_dataset 流式传输

prepare_tf_dataset() 创建一个 tf.data 管道，该管道从 tf.data.Dataset 加载样本。该管道使用 tf.numpy_function 或 from_generator（TensorFlow 无法编译），以访问底层的 tf.data.Dataset。它也无法在 Colab TPU 或 TPU 节点上工作，因为该管道从本地磁盘流式传输数据。请参阅下表，以帮助您决定此方法是否对您有帮助。

优点	缺点
代码简单	仅适用于 TPU VM
TPU/GPU 上相同的方法	数据必须作为 Hugging Face Dataset 提供
数据集不必适合内存	数据必须适合本地存储
支持可变填充	在大型 TPU pod 切片上，数据加载可能成为瓶颈

prepare_tf_dataset() 仅适用于 TPU VM。由于数据集存储在磁盘上，这意味着它可以处理大于可用内存的数据，因此将分词器输出作为列添加到数据集中。通过使用 tf.data 管道包装数据集，使用 prepare_tf_dataset() 从数据集流式传输数据。

def tokenize_function(examples):
    return tokenizer(
        examples["sentence"], padding="max_length", truncation=True, max_length=128
    )
# add the tokenizer output to the dataset as new columns
dataset = dataset.map(tokenize_function)

# prepare_tf_dataset() chooses columns that match the models input names
tf_dataset = model.prepare_tf_dataset(
    dataset, batch_size=BATCH_SIZE, shuffle=True, tokenizer=tokenizer
)

现在可以将数据集传递给 fit 方法。

from transformers import AutoTokenizer, TFAutoModelForSequenceClassification

with strategy.scope():
    model = TFAutoModelForSequenceClassification.from_pretrained(model_checkpoint)
    model.compile(optimizer="adam")

model.fit(tf_dataset)

TPU 类型

TPU 有两种类型：TPU 节点和 TPU VM。

TPU 节点间接访问远程 TPU。它需要一个单独的 VM 来初始化您的网络和数据管道，然后将其转发到远程节点。Google Colab TPU 是 TPU 节点的一个示例。您不能使用本地数据，因为 TPU 位于远程位置，并且数据必须存储在 Google Cloud Storage 中，数据管道可以在其中访问它。

TPU VM 直接连接到 TPU 所在的机器，并且它们通常更易于使用，尤其是在涉及您的数据管道时。

如果可能，我们建议避免使用 TPU 节点，因为它比 TPU VM 更难调试。TPU 节点将来也可能不受支持，并成为一种过时的访问方法。

单个 TPU（v2-8、v3-8、v4-8）运行 8 个副本。TPU 可以存在于 pod 中，这些 pod 可以同时运行数百甚至数千个副本。当您仅使用 pod 的一部分时，它被称为 pod 切片。在 Google Colab 上，您通常会获得单个 v2-8 TPU。

XLA

XLA 是一个用于高性能执行的线性代数编译器，默认情况下使用它来提高 TPU 的性能。

在 TPU 上执行代码之前，最好先在 CPU 或 GPU 上尝试，因为它更容易调试。您可以训练几个步骤，以确保模型和数据管道按预期工作。在 compile 方法中设置 jit_compile=True 以启用 XLA 编译（但请记住在 TPU 上运行之前删除此行代码）。

以下部分概述了使您的代码与 XLA 兼容的三个规则。Transformers 默认对模型和损失函数强制执行前两个规则，但如果您要编写自己的模型和损失函数，请不要忘记它们。

数据依赖条件语句

任何 if 语句都不能依赖于 tf.Tensor 内的值。以下代码无法由 XLA 编译。

if tf.reduce_sum(tensor) > 10:
    tensor = tensor / 2.0

要使用 XLA 编译，请使用 tf.cond 或删除条件语句，并使用指示变量代替，如下所示。

sum_over_10 = tf.cast(tf.reduce_sum(tensor) > 10, tf.float32)
tensor = tensor / (1.0 + sum_over_10)

数据依赖形状

tf.Tensor 的形状不能依赖于它们的值。例如，tf.unique 无法编译，因为它返回一个张量，其中包含输入中每个唯一值的实例。此输出的形状取决于输入 tf.Tensor 的重复程度。

这是标签掩码期间的一个问题，其中标签设置为负值以指示在计算损失时应忽略它们。以下代码无法由 XLA 编译，因为 masked_outputs 和 masked_labels 的形状取决于掩码的位置数量。

label_mask = labels >= 0
masked_outputs = outputs[label_mask]
masked_labels = labels[label_mask]
loss = compute_loss(masked_outputs, masked_labels)
mean_loss = torch.mean(loss)

要使用 XLA 编译，请通过计算每个位置的损失并在计算均值时将分子和分母中掩码的位置归零来避免数据依赖形状。将 tf.bool 转换为 tf.float32 作为指示变量，使您的代码与 XLA 兼容。

label_mask = tf.cast(labels >= 0, tf.float32)
loss = compute_loss(outputs, labels)
loss = loss * label_mask
mean_loss = tf.reduce_sum(loss) / tf.reduce_sum(label_mask)

重新编译不同的输入形状

如果输入形状是可变的，XLA 会重新编译您的模型，这会产生巨大的性能问题。这在文本模型中尤其常见，因为输入文本在分词后具有可变的长度。

过度的填充也会严重降低训练速度，因为它需要更多的计算和内存来处理。

为了避免不同的形状，请使用填充将所有输入填充到相同的长度，并使用 attention_mask。尝试将批量的样本填充为 32 或 64 个令牌的倍数。使用参数 padding="max_length"、padding="longest" 或 pad_to_multiple_of 来帮助填充。这通常会少量增加令牌的数量，但会显着减少唯一输入形状的数量，因为每个输入形状都是 32 或 64 的倍数。更少的唯一输入形状需要更少的重新编译。

< > 在 GitHub 上更新