TPU

TPU (Tensor Processing Unit) 是一种旨在加速张量计算以进行训练和推理的硬件。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))

Transformers 预处理类，例如 AutoImageProcessor 是与框架无关的，无法由 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 数据集提供
数据集不必完全加载到内存中	数据必须适合本地存储
支持可变填充	在大规模 TPU Pod 切片上，数据加载可能是瓶颈

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

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 slice。在 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 上更新