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 上更新