作者: Matthew Watson
创建日期 2022/04/18
上次修改 2023/07/15
描述: 使用 KerasHub 从零开始训练 Transformer 模型。
KerasHub 旨在简化构建最先进的文本处理模型。在本指南中,我们将展示库组件如何简化从零开始预训练和微调 Transformer 模型的过程。
本指南分为三个部分:
以下指南使用 Keras 3,它可以在 tensorflow
、jax
或 torch
中的任何一个后端上工作。我们在下面选择 jax
后端,这会给我们一个特别快的训练步骤,但你也可以随意选择其他后端。
!pip install -q --upgrade keras-hub
!pip install -q --upgrade keras # Upgrade to Keras 3.
import os
os.environ["KERAS_BACKEND"] = "jax" # or "tensorflow" or "torch"
import keras_hub
import tensorflow as tf
import keras
接下来,我们可以下载两个数据集。
最后,我们将下载一个 WordPiece 词汇表,以便在本指南后面进行子词分词。
# Download pretraining data.
keras.utils.get_file(
origin="https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-103-raw-v1.zip",
extract=True,
)
wiki_dir = os.path.expanduser("~/.keras/datasets/wikitext-103-raw/")
# Download finetuning data.
keras.utils.get_file(
origin="https://dl.fbaipublicfiles.com/glue/data/SST-2.zip",
extract=True,
)
sst_dir = os.path.expanduser("~/.keras/datasets/SST-2/")
# Download vocabulary data.
vocab_file = keras.utils.get_file(
origin="https://storage.googleapis.com/tensorflow/keras-nlp/examples/bert/bert_vocab_uncased.txt",
)
接下来,我们定义一些将在训练期间使用的超参数。
# Preprocessing params.
PRETRAINING_BATCH_SIZE = 128
FINETUNING_BATCH_SIZE = 32
SEQ_LENGTH = 128
MASK_RATE = 0.25
PREDICTIONS_PER_SEQ = 32
# Model params.
NUM_LAYERS = 3
MODEL_DIM = 256
INTERMEDIATE_DIM = 512
NUM_HEADS = 4
DROPOUT = 0.1
NORM_EPSILON = 1e-5
# Training params.
PRETRAINING_LEARNING_RATE = 5e-4
PRETRAINING_EPOCHS = 8
FINETUNING_LEARNING_RATE = 5e-5
FINETUNING_EPOCHS = 3
我们使用 tf.data 加载数据,这使得我们可以定义用于分词和预处理文本的输入管道。
# Load SST-2.
sst_train_ds = tf.data.experimental.CsvDataset(
sst_dir + "train.tsv", [tf.string, tf.int32], header=True, field_delim="\t"
).batch(FINETUNING_BATCH_SIZE)
sst_val_ds = tf.data.experimental.CsvDataset(
sst_dir + "dev.tsv", [tf.string, tf.int32], header=True, field_delim="\t"
).batch(FINETUNING_BATCH_SIZE)
# Load wikitext-103 and filter out short lines.
wiki_train_ds = (
tf.data.TextLineDataset(wiki_dir + "wiki.train.raw")
.filter(lambda x: tf.strings.length(x) > 100)
.batch(PRETRAINING_BATCH_SIZE)
)
wiki_val_ds = (
tf.data.TextLineDataset(wiki_dir + "wiki.valid.raw")
.filter(lambda x: tf.strings.length(x) > 100)
.batch(PRETRAINING_BATCH_SIZE)
)
# Take a peak at the sst-2 dataset.
print(sst_train_ds.unbatch().batch(4).take(1).get_single_element())
(<tf.Tensor: shape=(4,), dtype=string, numpy=
array([b'hide new secretions from the parental units ',
b'contains no wit , only labored gags ',
b'that loves its characters and communicates something rather beautiful about human nature ',
b'remains utterly satisfied to remain the same throughout '],
dtype=object)>, <tf.Tensor: shape=(4,), dtype=int32, numpy=array([0, 0, 1, 0], dtype=int32)>)
你可以看到 SST-2
数据集包含相对较短的影评文本片段。我们的目标是预测片段的情感。标签 1 表示积极情感,标签 0 表示消极情感。
第一步,我们将建立一个良好性能的基线。我们实际上不需要 KerasHub 来完成此操作,只需使用核心 Keras 层即可。
我们将训练一个简单的词袋模型,在该模型中,我们为词汇表中的每个词学习一个积极或消极的权重。样本的得分是该样本中所有存在词的权重的总和。
# This layer will turn our input sentence into a list of 1s and 0s the same size
# our vocabulary, indicating whether a word is present in absent.
multi_hot_layer = keras.layers.TextVectorization(
max_tokens=4000, output_mode="multi_hot"
)
multi_hot_layer.adapt(sst_train_ds.map(lambda x, y: x))
multi_hot_ds = sst_train_ds.map(lambda x, y: (multi_hot_layer(x), y))
multi_hot_val_ds = sst_val_ds.map(lambda x, y: (multi_hot_layer(x), y))
# We then learn a linear regression over that layer, and that's our entire
# baseline model!
inputs = keras.Input(shape=(4000,), dtype="int32")
outputs = keras.layers.Dense(1, activation="sigmoid")(inputs)
baseline_model = keras.Model(inputs, outputs)
baseline_model.compile(loss="binary_crossentropy", metrics=["accuracy"])
baseline_model.fit(multi_hot_ds, validation_data=multi_hot_val_ds, epochs=5)
Epoch 1/5
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 2s 698us/step - accuracy: 0.6421 - loss: 0.6469 - val_accuracy: 0.7567 - val_loss: 0.5391
Epoch 2/5
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 1s 493us/step - accuracy: 0.7524 - loss: 0.5392 - val_accuracy: 0.7868 - val_loss: 0.4891
Epoch 3/5
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 1s 513us/step - accuracy: 0.7832 - loss: 0.4871 - val_accuracy: 0.7991 - val_loss: 0.4671
Epoch 4/5
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 1s 475us/step - accuracy: 0.7991 - loss: 0.4543 - val_accuracy: 0.8069 - val_loss: 0.4569
Epoch 5/5
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 1s 476us/step - accuracy: 0.8100 - loss: 0.4313 - val_accuracy: 0.8036 - val_loss: 0.4530
<keras.src.callbacks.history.History at 0x7f13902967a0>
词袋方法速度快且出人意料地强大,特别是在输入示例包含大量词时。对于较短的序列,它可能会达到性能瓶颈。
为了做得更好,我们希望构建一个能够在上下文中评估词的模型。我们不能孤立地评估每个词,而需要使用输入整个有序序列中包含的信息。
这就遇到了一个问题。SST-2
是一个非常小的数据集,根本没有足够的示例文本来尝试构建一个更大、参数更多的、可以在序列上学习的模型。我们将很快开始过拟合并记住训练集,而泛化到未见示例的能力没有任何提高。
这时,预训练登场了,它让我们可以在更大的语料库上学习,并将知识迁移到 SST-2
任务上。而 KerasHub 登场了,它使我们能够轻松地预训练一个特别强大的模型——Transformer。
为了超越我们的基线,我们将利用 WikiText103
数据集,这是一个比 SST-2
大得多的无标注维基百科文章集合。
我们将训练一个 transformer,这是一个高度表达能力的模型,它将学习将输入中的每个词嵌入为低维向量。我们的维基百科数据集没有标注,因此我们将使用一个称为掩码语言建模 (MaskedLM) 的无监督训练目标。
本质上,我们将玩一个大型的“猜缺失词”游戏。对于每个输入样本,我们将遮蔽 25% 的输入数据,并训练模型来预测我们遮住的部分。
MaskedLM 任务的文本预处理将分两个阶段进行。
为了进行分词,我们可以使用 keras_hub.tokenizers.Tokenizer
——这是 KerasHub 中用于将文本转换为整型 token ID 序列的构建块。
特别是,我们将使用 keras_hub.tokenizers.WordPieceTokenizer
,它执行子词分词。子词分词在大型文本语料库上训练模型时很受欢迎。本质上,它允许我们的模型从不常见的词中学习,而不需要一个包含训练集中所有词的庞大词汇表。
我们需要做的第二件事是为 MaskedLM 任务对输入进行掩码。为此,我们可以使用 keras_hub.layers.MaskedLMMaskGenerator
,它将随机选择每个输入中的一组 token 并将其掩码掉。
分词器和掩码层都可以用在对 tf.data.Dataset.map 的调用中。我们可以使用 tf.data
在 CPU 上高效地预计算每个批次,而 GPU 或 TPU 则处理上一个批次的训练。由于我们的掩码层每次都会选择新的词进行掩码,因此每次遍历数据集都会给我们一组全新的标签进行训练。
# Setting sequence_length will trim or pad the token outputs to shape
# (batch_size, SEQ_LENGTH).
tokenizer = keras_hub.tokenizers.WordPieceTokenizer(
vocabulary=vocab_file,
sequence_length=SEQ_LENGTH,
lowercase=True,
strip_accents=True,
)
# Setting mask_selection_length will trim or pad the mask outputs to shape
# (batch_size, PREDICTIONS_PER_SEQ).
masker = keras_hub.layers.MaskedLMMaskGenerator(
vocabulary_size=tokenizer.vocabulary_size(),
mask_selection_rate=MASK_RATE,
mask_selection_length=PREDICTIONS_PER_SEQ,
mask_token_id=tokenizer.token_to_id("[MASK]"),
)
def preprocess(inputs):
inputs = tokenizer(inputs)
outputs = masker(inputs)
# Split the masking layer outputs into a (features, labels, and weights)
# tuple that we can use with keras.Model.fit().
features = {
"token_ids": outputs["token_ids"],
"mask_positions": outputs["mask_positions"],
}
labels = outputs["mask_ids"]
weights = outputs["mask_weights"]
return features, labels, weights
# We use prefetch() to pre-compute preprocessed batches on the fly on the CPU.
pretrain_ds = wiki_train_ds.map(
preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
pretrain_val_ds = wiki_val_ds.map(
preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
# Preview a single input example.
# The masks will change each time you run the cell.
print(pretrain_val_ds.take(1).get_single_element())
({'token_ids': <tf.Tensor: shape=(128, 128), dtype=int32, numpy=
array([[7570, 7849, 2271, ..., 9673, 103, 7570],
[7570, 7849, 103, ..., 1007, 1012, 2023],
[1996, 2034, 3940, ..., 0, 0, 0],
...,
[2076, 1996, 2307, ..., 0, 0, 0],
[3216, 103, 2083, ..., 0, 0, 0],
[ 103, 2007, 1045, ..., 0, 0, 0]], dtype=int32)>, 'mask_positions': <tf.Tensor: shape=(128, 32), dtype=int64, numpy=
array([[ 5, 6, 7, ..., 118, 120, 126],
[ 2, 3, 14, ..., 105, 106, 113],
[ 4, 9, 10, ..., 0, 0, 0],
...,
[ 4, 11, 19, ..., 117, 118, 0],
[ 1, 14, 17, ..., 0, 0, 0],
[ 0, 3, 6, ..., 0, 0, 0]])>}, <tf.Tensor: shape=(128, 32), dtype=int32, numpy=
array([[ 1010, 2124, 2004, ..., 2095, 11300, 1012],
[ 2271, 13091, 2303, ..., 2029, 2027, 1010],
[23976, 2007, 1037, ..., 0, 0, 0],
...,
[ 1010, 1996, 1010, ..., 1999, 7511, 0],
[ 2225, 1998, 10722, ..., 0, 0, 0],
[ 9794, 1030, 2322, ..., 0, 0, 0]], dtype=int32)>, <tf.Tensor: shape=(128, 32), dtype=float32, numpy=
array([[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 1., 1., 1.],
[1., 1., 1., ..., 0., 0., 0.],
...,
[1., 1., 1., ..., 1., 1., 0.],
[1., 1., 1., ..., 0., 0., 0.],
[1., 1., 1., ..., 0., 0., 0.]], dtype=float32)>)
以上代码块将数据集整理成一个 (features, labels, weights)
元组,可以直接传递给 keras.Model.fit()
。
我们有两个特征:
"token_ids"
,其中一些 token 已被我们的掩码 token ID 替换。"mask_positions"
,用于跟踪我们掩码了哪些 token。我们的标签就是我们掩码掉的 ID。
因为并非所有序列都具有相同数量的掩码,所以我们还保留了一个 sample_weight
张量,通过给填充标签零权重来将其从损失函数中移除。
KerasHub 提供了快速构建 Transformer 编码器的所有构建块。
我们使用 keras_hub.layers.TokenAndPositionEmbedding
首先嵌入我们的输入 token ID。该层同时学习两种嵌入——一种用于句子中的词,另一种用于句子中的整数位置。输出嵌入是两者的简单求和。
然后我们可以添加一系列 keras_hub.layers.TransformerEncoder
层。这些是 Transformer 模型的核心,它们使用注意力机制关注输入句子的不同部分,然后是一个多层感知器块。
该模型的输出将是每个输入 token ID 的编码向量。与我们用作基线的词袋模型不同,该模型将根据上下文来嵌入每个 token。
inputs = keras.Input(shape=(SEQ_LENGTH,), dtype="int32")
# Embed our tokens with a positional embedding.
embedding_layer = keras_hub.layers.TokenAndPositionEmbedding(
vocabulary_size=tokenizer.vocabulary_size(),
sequence_length=SEQ_LENGTH,
embedding_dim=MODEL_DIM,
)
outputs = embedding_layer(inputs)
# Apply layer normalization and dropout to the embedding.
outputs = keras.layers.LayerNormalization(epsilon=NORM_EPSILON)(outputs)
outputs = keras.layers.Dropout(rate=DROPOUT)(outputs)
# Add a number of encoder blocks
for i in range(NUM_LAYERS):
outputs = keras_hub.layers.TransformerEncoder(
intermediate_dim=INTERMEDIATE_DIM,
num_heads=NUM_HEADS,
dropout=DROPOUT,
layer_norm_epsilon=NORM_EPSILON,
)(outputs)
encoder_model = keras.Model(inputs, outputs)
encoder_model.summary()
Model: "functional_3"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓ ┃ Layer (type) ┃ Output Shape ┃ Param # ┃ ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩ │ input_layer_1 (InputLayer) │ (None, 128) │ 0 │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ token_and_position_embedding │ (None, 128, 256) │ 7,846,400 │ │ (TokenAndPositionEmbedding) │ │ │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ layer_normalization │ (None, 128, 256) │ 512 │ │ (LayerNormalization) │ │ │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ dropout (Dropout) │ (None, 128, 256) │ 0 │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ transformer_encoder │ (None, 128, 256) │ 527,104 │ │ (TransformerEncoder) │ │ │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ transformer_encoder_1 │ (None, 128, 256) │ 527,104 │ │ (TransformerEncoder) │ │ │ ├─────────────────────────────────┼───────────────────────────┼────────────┤ │ transformer_encoder_2 │ (None, 128, 256) │ 527,104 │ │ (TransformerEncoder) │ │ │ └─────────────────────────────────┴───────────────────────────┴────────────┘
Total params: 9,428,224 (287.73 MB)
Trainable params: 9,428,224 (287.73 MB)
Non-trainable params: 0 (0.00 B)
你可以将 encoder_model
视为一个独立的模块化单元,它是我们在下游任务中真正感兴趣的模型部分。然而,我们仍然需要设置编码器来训练 MaskedLM 任务;为此,我们连接一个 keras_hub.layers.MaskedLMHead
。
该层将把 token 编码作为输入之一,并将我们在原始输入中掩码掉的位置作为另一个输入。它将收集我们掩码的 token 编码,并将其转换回对我们整个词汇表的预测。
至此,我们已准备好编译并运行预训练。如果你在 Colab 中运行此代码,请注意这将花费大约一个小时。Transformer 的训练以计算密集型而闻名,因此即使是这个相对较小的 Transformer 也需要一些时间。
# Create the pretraining model by attaching a masked language model head.
inputs = {
"token_ids": keras.Input(shape=(SEQ_LENGTH,), dtype="int32", name="token_ids"),
"mask_positions": keras.Input(
shape=(PREDICTIONS_PER_SEQ,), dtype="int32", name="mask_positions"
),
}
# Encode the tokens.
encoded_tokens = encoder_model(inputs["token_ids"])
# Predict an output word for each masked input token.
# We use the input token embedding to project from our encoded vectors to
# vocabulary logits, which has been shown to improve training efficiency.
outputs = keras_hub.layers.MaskedLMHead(
token_embedding=embedding_layer.token_embedding,
activation="softmax",
)(encoded_tokens, mask_positions=inputs["mask_positions"])
# Define and compile our pretraining model.
pretraining_model = keras.Model(inputs, outputs)
pretraining_model.compile(
loss="sparse_categorical_crossentropy",
optimizer=keras.optimizers.AdamW(PRETRAINING_LEARNING_RATE),
weighted_metrics=["sparse_categorical_accuracy"],
jit_compile=True,
)
# Pretrain the model on our wiki text dataset.
pretraining_model.fit(
pretrain_ds,
validation_data=pretrain_val_ds,
epochs=PRETRAINING_EPOCHS,
)
# Save this base model for further finetuning.
encoder_model.save("encoder_model.keras")
Epoch 1/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 242s 41ms/step - loss: 5.4679 - sparse_categorical_accuracy: 0.1353 - val_loss: 3.4570 - val_sparse_categorical_accuracy: 0.3522
Epoch 2/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 234s 40ms/step - loss: 3.6031 - sparse_categorical_accuracy: 0.3396 - val_loss: 3.0514 - val_sparse_categorical_accuracy: 0.4032
Epoch 3/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 232s 40ms/step - loss: 3.2609 - sparse_categorical_accuracy: 0.3802 - val_loss: 2.8858 - val_sparse_categorical_accuracy: 0.4240
Epoch 4/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 233s 40ms/step - loss: 3.1099 - sparse_categorical_accuracy: 0.3978 - val_loss: 2.7897 - val_sparse_categorical_accuracy: 0.4375
Epoch 5/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 235s 40ms/step - loss: 3.0145 - sparse_categorical_accuracy: 0.4090 - val_loss: 2.7504 - val_sparse_categorical_accuracy: 0.4419
Epoch 6/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 252s 43ms/step - loss: 2.9530 - sparse_categorical_accuracy: 0.4157 - val_loss: 2.6925 - val_sparse_categorical_accuracy: 0.4474
Epoch 7/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 232s 40ms/step - loss: 2.9088 - sparse_categorical_accuracy: 0.4210 - val_loss: 2.6554 - val_sparse_categorical_accuracy: 0.4513
Epoch 8/8
5857/5857 ━━━━━━━━━━━━━━━━━━━━ 236s 40ms/step - loss: 2.8721 - sparse_categorical_accuracy: 0.4250 - val_loss: 2.6389 - val_sparse_categorical_accuracy: 0.4548
预训练后,我们现在可以在 SST-2
数据集上微调我们的模型。我们可以利用构建的编码器在上下文中预测词的能力来提高我们在下游任务上的性能。
微调的预处理比我们预训练的 MaskedLM 任务简单得多。我们只需对输入句子进行分词,就可以进行训练了!
def preprocess(sentences, labels):
return tokenizer(sentences), labels
# We use prefetch() to pre-compute preprocessed batches on the fly on our CPU.
finetune_ds = sst_train_ds.map(
preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
finetune_val_ds = sst_val_ds.map(
preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
# Preview a single input example.
print(finetune_val_ds.take(1).get_single_element())
(<tf.Tensor: shape=(32, 128), dtype=int32, numpy=
array([[ 2009, 1005, 1055, ..., 0, 0, 0],
[ 4895, 10258, 2378, ..., 0, 0, 0],
[ 4473, 2149, 2000, ..., 0, 0, 0],
...,
[ 1045, 2018, 2000, ..., 0, 0, 0],
[ 4283, 2000, 3660, ..., 0, 0, 0],
[ 1012, 1012, 1012, ..., 0, 0, 0]], dtype=int32)>, <tf.Tensor: shape=(32,), dtype=int32, numpy=
array([1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 0,
0, 1, 1, 0, 0, 1, 0, 0, 1, 0], dtype=int32)>)
为了将编码的 token 输出转换为分类预测,我们需要在我们的 Transformer 模型上附加另一个“头”。在这里我们可以保持简单。我们将编码的 token 汇集在一起,并使用一个简单的全连接层进行预测。
# Reload the encoder model from disk so we can restart fine-tuning from scratch.
encoder_model = keras.models.load_model("encoder_model.keras", compile=False)
# Take as input the tokenized input.
inputs = keras.Input(shape=(SEQ_LENGTH,), dtype="int32")
# Encode and pool the tokens.
encoded_tokens = encoder_model(inputs)
pooled_tokens = keras.layers.GlobalAveragePooling1D()(encoded_tokens[0])
# Predict an output label.
outputs = keras.layers.Dense(1, activation="sigmoid")(pooled_tokens)
# Define and compile our fine-tuning model.
finetuning_model = keras.Model(inputs, outputs)
finetuning_model.compile(
loss="binary_crossentropy",
optimizer=keras.optimizers.AdamW(FINETUNING_LEARNING_RATE),
metrics=["accuracy"],
)
# Finetune the model for the SST-2 task.
finetuning_model.fit(
finetune_ds,
validation_data=finetune_val_ds,
epochs=FINETUNING_EPOCHS,
)
Epoch 1/3
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 21s 9ms/step - accuracy: 0.7500 - loss: 0.4891 - val_accuracy: 0.8036 - val_loss: 0.4099
Epoch 2/3
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 16s 8ms/step - accuracy: 0.8826 - loss: 0.2779 - val_accuracy: 0.8482 - val_loss: 0.3964
Epoch 3/3
2105/2105 ━━━━━━━━━━━━━━━━━━━━ 16s 8ms/step - accuracy: 0.9176 - loss: 0.2066 - val_accuracy: 0.8549 - val_loss: 0.4142
<keras.src.callbacks.history.History at 0x7f12d85c21a0>
预训练足以将我们的性能提升到 84%,但这远非 Transformer 模型的上限。你可能已经注意到,在预训练期间,我们的验证性能仍在稳步提升。我们的模型仍然远未充分训练。训练更多轮次、训练更大的 Transformer 以及在更多无标注文本上训练都将显著提高性能。
KerasHub 的主要目标之一是为 NLP 模型构建提供模块化方法。我们在这里展示了一种构建 Transformer 的方法,但 KerasHub 支持越来越多的用于文本预处理和构建模型的组件。我们希望它能让你更容易地实验解决自然语言问题的方法。