Keras深度学习框架第二十五讲：使用KerasNLP预训练Transformer模型

1、KerasNPL预训练Transformer模型概念

使用KerasNLP来预训练一个Transformer模型涉及多个步骤。由于Keras本身并不直接提供NLP的预训练模型或工具集，我们通常需要结合像TensorFlow Hub、Hugging Face的Transformers库或自定义的Keras层来实现。

以下是一个简化的步骤概述，用于说明如何使用Keras和相关的NLP工具来预训练一个Transformer模型：

• 环境准备：

    - 安装TensorFlow和Keras。
    - 安装任何你需要的NLP相关库，如Hugging Face的Transformers。
  

• 数据准备：
  - 收集一个大型的无标签文本数据集，用于预训练。
  - 对文本进行预处理，包括分词、填充/截断、创建词汇表等。

• 定义模型：
  - 使用Keras的Sequential API或函数式API定义Transformer模型的架构。
  - 通常，你会需要实现Transformer的编码器层（包括自注意力机制和前馈神经网络）。
  - 你也可以考虑使用现有的Transformer实现作为起点，如Hugging Face的Transformers库中的BERT或GPT模型。

• 编译和配置模型：
  - 编译模型，设置损失函数（对于预训练，可能是掩码语言建模的损失或句子级别的损失）和优化器。
  - 配置模型训练的超参数，如学习率、批大小、训练轮数等。

• 训练模型：
  - 使用准备好的数据集训练模型。这可能需要很长时间，特别是当数据集很大时。
  - 监控训练过程中的损失和任何其他的评估指标。

• 保存模型：
  - 在训练完成后，保存模型的权重和架构。

• （可选）微调：
  - 如果你打算将预训练的Transformer模型用于特定的NLP任务，你可能需要在有标签的数据集上进行微调。

需要注意的是，预训练一个Transformer模型是一个复杂且资源密集的任务，通常需要大量的计算资源（如GPU或TPU）和时间。此外，由于Keras本身并不直接提供NLP的预训练功能，你可能需要深入了解Transformer架构和相关的NLP技术来实现这一目标。

如果你希望快速开始并使用预训练的Transformer模型，考虑使用Hugging Face的Transformers库，它提供了大量预训练的模型和方便的API来加载和使用这些模型。

KerasNLP旨在简化构建最先进的文本处理模型的过程。在本文中，我们将展示如何使用该库的组件来简化从头开始预训练和微调Transformer模型的流程。

• 设置、任务定义和建立基线。
• 预训练一个Transformer模型。
• 在分类任务上微调Transformer模型。

2、KerasNLP预训练Transformer模型的步骤

2.1 设置

安装及导入

!pip install -q --upgrade keras-nlp
!pip install -q --upgrade keras  # Upgrade to Keras 3.

import os

os.environ["KERAS_BACKEND"] = "jax"  # or "tensorflow" or "torch"


import keras_nlp
import tensorflow as tf
import keras

下载数据集
接下来，我们可以下载两个数据集。

SST-2 是一个文本分类数据集，也是我们的“最终目标”。这个数据集经常被用来作为语言模型的基准测试。
WikiText-103：一个中等规模的英文维基百科精选文章集合，我们将用于预训练。
最后，我们将下载一个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

2.2 加载数据

使用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表示负面情感。

2.3建立基线

作为第一步，我们将建立一个性能良好的基线。其实这一步并不需要KerasNLP，我们只需要使用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)

词袋方法可能既快速又出乎意料地强大，尤其是在输入示例包含大量单词时。但对于较短的序列，它可能会达到性能上限。

为了做得更好，我们希望构建一个能够评估上下文中的单词的模型。我们不能再将每个单词视为孤立的存在，而是需要利用输入中整个有序序列所包含的信息。

这就引出了一个问题。SST-2是一个非常小的数据集，没有足够的示例文本来尝试构建一个更大、参数更多的模型来学习序列。我们很快就会开始过拟合并记住训练集，而没有任何提高我们对未见示例的泛化能力。

这时，预训练就派上了用场，它允许我们在一个更大的语料库上进行学习，并将我们的知识转移到SST-2任务上。同时，KerasNLP的引入使我们能够轻松地预训练一个特别强大的模型——Transformer。

2.4预训练

为了超越我们的基线模型，我们将利用WikiText103数据集，这是一个比SST-2大得多的无标签维基百科文章集合。

我们将训练一个Transformer模型，这是一个高度表达性的模型，它将学习将输入中的每个单词嵌入为低维向量。由于我们的维基百科数据集没有标签，因此我们将使用一个无监督的训练目标，称为掩码语言建模（Masked Language Modeling，MaskedLM）。

基本上，我们将玩一个大型的“猜缺失单词”游戏。对于每个输入样本，我们将遮盖25%的输入数据，并训练我们的模型来预测我们遮盖的部分。

2.4.1为MaskedLM任务预处理数据

我们为MaskedLM任务进行的文本预处理将分为两个阶段。

• 将输入文本分词为token id的整数序列。
• 遮盖输入中的某些位置以进行预测。

为了进行分词，我们可以使用keras_nlp.tokenizers.Tokenizer——KerasNLP中将文本转换为整数token id序列的构建块。

特别是，我们将使用keras_nlp.tokenizers.WordPieceTokenizer进行子词分词。在大文本语料库上训练模型时，子词分词非常流行。基本上，它允许我们的模型从不常见的单词中学习，而不需要包含训练集中每个单词的庞大词汇表。

我们需要做的第二件事是为MaskedLM任务遮盖输入。为此，我们可以使用keras_nlp.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_nlp.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_nlp.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())

上述部分将我们的数据集整理为一个(features, labels, weights)的元组，可以直接传递给keras.Model.fit()函数进行训练。

我们有两个特征：

• “token_ids”，其中一些token已经被替换为我们的掩码token id。
• “mask_positions”，跟踪我们遮盖了哪些token。

我们的标签就是我们遮盖掉的那些token的id。

由于不是所有序列的掩码数量都相同，我们还保留了一个sample_weight张量，通过给填充的标签零权重，来从我们的损失函数中移除它们。

2.4.2创建Transformer编码器

KerasNLP提供了所有构建块，可以快速构建一个Transformer编码器。

我们使用keras_nlp.layers.TokenAndPositionEmbedding首先将输入token id进行嵌入。这个层同时学习两种嵌入——一种是句子中单词的嵌入，另一种是句子中整数位置的嵌入。输出嵌入就是两者的和。

然后我们可以添加一系列的keras_nlp.layers.TransformerEncoder层。这些层是Transformer模型的核心，使用注意力机制来关注输入句子的不同部分，接着是一个多层感知机块。

这个模型的输出是每个输入token id的编码向量。与我们用作基线的词袋模型不同，这个模型在嵌入每个token时会考虑到它出现的上下文。

inputs = keras.Input(shape=(SEQ_LENGTH,), dtype="int32")

# Embed our tokens with a positional embedding.
embedding_layer = keras_nlp.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_nlp.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()

2.4.3预训练Transformer

你可以把encoder_model看作是一个模块化的单元，这是我们对下游任务真正感兴趣的部分。然而，我们仍然需要设置编码器以在MaskedLM任务上进行训练；为此，我们将添加一个keras_nlp.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_nlp.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")

2.5微调

预训练之后，我们现在可以在SST-2数据集上微调我们的模型。我们可以利用我们构建的编码器在上下文中预测单词的能力，以提高我们在下游任务上的性能。

2.5.1分类任务的数据预处理

对于微调任务的数据预处理比预训练的MaskedLM任务简单得多。我们只需要对输入句子进行分词，就可以开始训练了！

具体来说，对于SST-2这样的情感分类任务，我们可能需要执行以下步骤来准备数据：

• 分词：使用与预训练时相同的分词器对文本进行分词，并转换为token IDs。

• 填充和截断：由于模型期望接收固定长度的输入，我们需要将文本填充或截断到模型接受的长度。

• 标签编码：将情感标签转换为数字形式，以便模型可以处理它们作为目标变量。

• 划分数据集：将数据集划分为训练集、验证集（可选）和测试集。

• 构建输入：将分词后的文本（token IDs）和对应的标签组合成训练所需的格式（例如，使用tf.data API构建输入管道）。

• （可选）添加样本权重：如果数据集不平衡，我们可以考虑使用样本权重来平衡不同类别的贡献。

完成这些步骤后，我们就可以将预处理后的数据输入到预训练好的Transformer编码器中进行微调了。微调过程将调整模型的所有权重（或其中的一部分，如只调整顶部几层），以优化在SST-2数据集上的性能。

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())

2.5.2微调Transformer

要从编码后的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,
)

预训练已经将我们的性能提升到了84%，但这对于Transformer模型来说还远未达到上限。在预训练过程中，你可能已经注意到验证性能仍在稳步提高。我们的模型仍然远远没有训练充分。通过增加训练轮次、训练一个更大的Transformer模型，以及在更多的未标记文本上进行训练，都可以继续显著提高性能。

KerasNLP的关键目标之一是提供NLP模型构建的模块化方法。在这里，我们展示了构建Transformer模型的一种方法，但KerasNLP支持越来越多用于文本预处理和模型构建的组件。我们希望它能让您更容易地针对自然语言问题尝试不同的解决方案。

3、总结

本文的讨论主要涉及了Transformer模型的预训练和微调过程，以及如何使用KerasNLP库来构建和实验NLP模型。首先，提到了预训练Transformer模型对于提高下游任务性能的重要性，包括通过MaskedLM任务进行预训练。随后，介绍了在预训练之后，如何对Transformer模型进行微调以适应特定的分类任务（如SST-2情感分类），并强调了数据预处理的重要性。此外，还提到了预训练模型的性能可能远未达到上限，可以通过增加训练轮次、扩大模型规模以及在更多未标记文本上进行训练来进一步提高性能。最后，强调了KerasNLP库在NLP模型构建中的模块化方法，它支持各种用于文本预处理和模型构建的组件，使得实验和尝试不同的解决方案变得更加容易。整个讨论展示了Transformer模型在NLP任务中的强大潜力和灵活性，以及使用KerasNLP库进行NLP模型开发的优势。