开发者指南 / KerasHub / KerasHub 入门

KerasHub 入门

作者: Jonathan Bischof
创建日期 2022/12/15
上次修改日期 2023/07/01
描述:KerasHub API 简介。

在 Colab 中查看 GitHub 源代码


简介

KerasHub 是一个自然语言处理库,它在用户的整个开发周期中提供支持。我们的工作流程构建于模块化组件之上,这些组件在开箱即用时具有最先进的预设权重和架构,并且在需要更多控制时易于自定义。

该库是核心 Keras API 的扩展;所有高级模块都是 LayersModels。如果您熟悉 Keras,恭喜!您已经了解了 KerasHub 的大部分内容。

KerasHub 使用 Keras 3 与 TensorFlow、Pytorch 和 Jax 中的任何一个一起工作。在下面的指南中,我们将使用 jax 后端来训练我们的模型,并使用 tf.data 来高效地运行我们的输入预处理。但请随意混合搭配!本指南在 TensorFlow 或 PyTorch 后端中运行,无需任何更改,只需更新下面的 KERAS_BACKEND 即可。

本指南使用六个复杂程度级别的观点分析示例演示了我们的模块化方法。

  • 使用预训练分类器进行推理
  • 微调预训练主干网络
  • 使用用户控制的预处理进行微调
  • 微调自定义模型
  • 预训练主干模型
  • 从头开始构建和训练自己的 Transformer

在整个指南中,我们使用 Keras 官方吉祥物 Professor Keras 作为材料复杂度的视觉参考。

drawing

!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 keras

# Use mixed precision to speed up all training in this guide.
keras.mixed_precision.set_global_policy("mixed_float16")


API 快速入门

我们最高级别的 API 是 keras_hub.models。这些符号涵盖了将字符串转换为标记、标记转换为密集特征以及密集特征转换为特定任务输出的完整用户旅程。对于每个 XX 架构(例如,Bert),我们提供以下模块。

  • 分词器keras_hub.models.XXTokenizer
    • 作用:将字符串转换为标记 ID 序列。
    • 重要性:字符串的原始字节维度过高,无法用作有用的特征,因此我们首先将其映射到少量标记,例如 "The quick brown fox" 映射到 ["the", "qu", "##ick", "br", "##own", "fox"]
    • 继承自keras.layers.Layer
  • 预处理器keras_hub.models.XXPreprocessor
    • 作用:将字符串转换为由主干网络使用的预处理张量字典,从分词开始。
    • 重要性:每个模型都使用特殊的标记和额外的张量来理解输入,例如分隔输入片段和识别填充标记。将每个序列填充到相同的长度可以提高计算效率。
    • 包含XXTokenizer
    • 继承自keras.layers.Layer
  • 主干网络keras_hub.models.XXBackbone
    • 作用:将预处理的张量转换为密集特征。不处理字符串;请先调用预处理器。
    • 重要性:主干网络将输入标记提炼成密集特征,这些特征可用于下游任务。它通常在使用大量未标记数据的语言建模任务上进行预训练。将这些信息转移到新任务是现代 NLP 的一项重大突破。
    • 继承自keras.Model
  • 任务:例如,keras_hub.models.XXClassifier
    • 作用:将字符串转换为特定任务的输出(例如,分类概率)。
    • 重要性:任务模型将字符串预处理和主干模型与特定任务的 Layers 相结合,以解决句子分类、标记分类或文本生成等问题。额外的 Layers 必须在标记数据上进行微调。
    • 包含XXBackboneXXPreprocessor
    • 继承自keras.Model

以下是 BertClassifier 的模块化层次结构(所有关系都是组合关系)。

drawing

所有模块都可以独立使用,并且除了标准构造函数之外,还具有 from_preset() 方法,该方法使用**预设**架构和权重实例化类(请参见下面的示例)。


数据

我们将使用 IMDB 电影评论的情感分析作为运行示例。在此任务中,我们使用文本预测评论是正面(label = 1)还是负面(label = 0)。

我们使用 keras.utils.text_dataset_from_directory 加载数据,该函数利用功能强大的 tf.data.Dataset 格式来表示示例。

!curl -O https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz
!tar -xf aclImdb_v1.tar.gz
!# Remove unsupervised examples
!rm -r aclImdb/train/unsup
BATCH_SIZE = 16
imdb_train = keras.utils.text_dataset_from_directory(
    "aclImdb/train",
    batch_size=BATCH_SIZE,
)
imdb_test = keras.utils.text_dataset_from_directory(
    "aclImdb/test",
    batch_size=BATCH_SIZE,
)

# Inspect first review
# Format is (review text tensor, label tensor)
print(imdb_train.unbatch().take(1).get_single_element())
  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
                                 Dload  Upload   Total   Spent    Left  Speed
100 80.2M  100 80.2M    0     0  88.0M      0 --:--:-- --:--:-- --:--:-- 87.9M

Found 25000 files belonging to 2 classes.
Found 25000 files belonging to 2 classes.
(<tf.Tensor: shape=(), dtype=string, numpy=b'This is a very, very early Bugs Bunny cartoon. As a result, the character is still in a transition period--he is not drawn as elongated as he later was and his voice isn\'t quite right. In addition, the chemistry between Elmer and Bugs is a little unusual. Elmer is some poor sap who buys Bugs from a pet shop--there is no gun or desire on his part to blast the bunny to smithereens! However, despite this, this is still a very enjoyable film. The early Bugs was definitely more sassy and cruel than his later incarnations. In later films, he messed with Elmer, Yosimite Sam and others because they started it--they messed with the rabbit. But, in this film, he is much more like Daffy Duck of the late 30s and early 40s--a jerk who just loves irritating others!! A true "anarchist" instead of the hero of the later cartoons. While this isn\'t among the best Bug Bunny cartoons, it sure is fun to watch and it\'s interesting to see just how much he\'s changed over the years.'>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)

使用预训练分类器进行推理

drawing

KerasHub 中的最高级模块是**任务**。**任务**是一个 keras.Model,它由一个(通常是预训练的)**主干网络**模型和特定任务的层组成。以下是一个使用 keras_hub.models.BertClassifier 的示例。

注意:输出是每个类别的 logits(例如,[0, 0] 表示正面概率为 50%)。对于二元分类,输出为 [负面,正面]。

classifier = keras_hub.models.BertClassifier.from_preset("bert_tiny_en_uncased_sst2")
# Note: batched inputs expected so must wrap string in iterable
classifier.predict(["I love modular workflows in keras-hub!"])
 1/1 ━━━━━━━━━━━━━━━━━━━━ 1s 689ms/step

array([[-1.539,  1.543]], dtype=float16)

所有**任务**都具有 from_preset 方法,该方法使用预设的预处理、架构和权重构建 keras.Model 实例。这意味着我们可以以 keras.Model 接受的任何格式传递原始字符串,并获得特定于我们任务的输出。

此特定**预设**是一个在 sst2 上微调的 "bert_tiny_uncased_en" **主干网络**(另一个电影评论情感分析,这次来自 Rotten Tomatoes)。出于演示目的,我们使用 tiny 架构,但对于 SoTA 性能,建议使用更大的模型。有关 BertClassifier 可用的所有特定任务预设,请参阅我们的 keras.io 模型页面

让我们在 IMDB 数据集上评估我们的分类器。您会注意到我们不需要在这里调用 keras.Model.compile。所有像 BertClassifier 这样的**任务**模型都带有编译默认值,这意味着我们只需直接调用 keras.Model.evaluate 即可。您始终可以像往常一样调用 compile 来覆盖这些默认值(例如,添加新的指标)。

下面的输出为 [损失,准确率],

classifier.evaluate(imdb_test)
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 4s 2ms/step - loss: 0.4610 - sparse_categorical_accuracy: 0.7882

[0.4630218744277954, 0.783519983291626]

我们的结果是 78% 的准确率,而无需训练任何内容。还不错!


微调预训练的 BERT 主干网络

drawing

当特定于我们任务的标记文本可用时,微调自定义分类器可以提高性能。如果我们想要预测 IMDB 评论的情感,使用 IMDB 数据应该比使用 Rotten Tomatoes 数据表现更好!并且对于许多任务,将没有相关的预训练模型可用(例如,对客户评论进行分类)。

微调的工作流程与上面几乎相同,只是我们请求**主干网络**模型的**预设**,而不是整个分类器。当传递**主干网络****预设**时,**任务**Model 将随机初始化所有特定任务的层,以准备进行训练。有关 BertClassifier 可用的所有**主干网络**预设,请参阅我们的 keras.io 模型页面

要训练您的分类器,请使用 keras.Model.fit,就像使用任何其他 keras.Model 一样。与我们的推理示例一样,我们可以依赖**任务**的编译默认值并跳过 keras.Model.compile。由于预处理已包含在内,因此我们再次传递原始数据。

classifier = keras_hub.models.BertClassifier.from_preset(
    "bert_tiny_en_uncased",
    num_classes=2,
)
classifier.fit(
    imdb_train,
    validation_data=imdb_test,
    epochs=1,
)
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 16s 9ms/step - loss: 0.5202 - sparse_categorical_accuracy: 0.7281 - val_loss: 0.3254 - val_sparse_categorical_accuracy: 0.8621

<keras.src.callbacks.history.History at 0x7f281ffc9f90>

在这里,我们看到验证准确率有了显著提高(0.78 -> 0.87),即使 IMDB 数据集远小于 sst2,也只需进行一次训练周期。


使用用户控制的预处理进行微调

drawing

对于某些高级训练场景,用户可能更喜欢直接控制预处理。对于大型数据集,可以预先预处理示例并保存到磁盘,或者使用 tf.data.experimental.service 由单独的工作池进行预处理。在其他情况下,需要自定义预处理来处理输入。

preprocessor=None 传递给**任务**Model 的构造函数以跳过自动预处理,或者传递自定义的 BertPreprocessor

从同一预设中分离预处理

每个模型架构都有一个平行的**预处理器**Layer,它有自己的 from_preset 构造函数。使用此 Layer 的相同**预设**将返回与**任务**匹配的**预处理器**。

在这个工作流程中,我们使用tf.data.Dataset.cache()对模型训练三个周期,该函数会在拟合开始前计算一次预处理并缓存结果。

注意:我们可以在使用 Jax 或 PyTorch 后端时使用tf.data进行预处理。输入数据集会在拟合过程中自动转换为后端原生张量类型。事实上,鉴于tf.data在运行预处理方面的效率,这在所有后端上都是良好的实践。

import tensorflow as tf

preprocessor = keras_hub.models.BertPreprocessor.from_preset(
    "bert_tiny_en_uncased",
    sequence_length=512,
)

# Apply the preprocessor to every sample of train and test data using `map()`.
# [`tf.data.AUTOTUNE`](https://tensorflowcn.cn/api_docs/python/tf/data/AUTOTUNE) and `prefetch()` are options to tune performance, see
# https://tensorflowcn.cn/guide/data_performance for details.

# Note: only call `cache()` if you training data fits in CPU memory!
imdb_train_cached = (
    imdb_train.map(preprocessor, tf.data.AUTOTUNE).cache().prefetch(tf.data.AUTOTUNE)
)
imdb_test_cached = (
    imdb_test.map(preprocessor, tf.data.AUTOTUNE).cache().prefetch(tf.data.AUTOTUNE)
)

classifier = keras_hub.models.BertClassifier.from_preset(
    "bert_tiny_en_uncased", preprocessor=None, num_classes=2
)
classifier.fit(
    imdb_train_cached,
    validation_data=imdb_test_cached,
    epochs=3,
)
Epoch 1/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 15s 8ms/step - loss: 0.5194 - sparse_categorical_accuracy: 0.7272 - val_loss: 0.3032 - val_sparse_categorical_accuracy: 0.8728
Epoch 2/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 10s 7ms/step - loss: 0.2871 - sparse_categorical_accuracy: 0.8805 - val_loss: 0.2809 - val_sparse_categorical_accuracy: 0.8818
Epoch 3/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 10s 7ms/step - loss: 0.2134 - sparse_categorical_accuracy: 0.9178 - val_loss: 0.3043 - val_sparse_categorical_accuracy: 0.8790

<keras.src.callbacks.history.History at 0x7f281ffc87f0>

三个周期后,我们的验证准确率仅提升至 0.88。这既与我们数据集的小规模有关,也与我们的模型有关。要超过 90% 的准确率,请尝试使用更大的预设,例如"bert_base_en_uncased"。有关BertClassifier可用的所有主干预设,请参阅我们的 keras.io 模型页面

自定义预处理

在需要自定义预处理的情况下,我们提供了对Tokenizer类的直接访问权限,该类将原始字符串映射到标记。它还有一个from_preset()构造函数,用于获取与预训练匹配的词汇表。

注意:BertTokenizer默认情况下不会填充序列,因此输出是不规则的(每个序列的长度都不同)。下面的MultiSegmentPacker处理将这些不规则序列填充到密集张量类型(例如tf.Tensortorch.Tensor)。

tokenizer = keras_hub.models.BertTokenizer.from_preset("bert_tiny_en_uncased")
tokenizer(["I love modular workflows!", "Libraries over frameworks!"])

# Write your own packer or use one of our `Layers`
packer = keras_hub.layers.MultiSegmentPacker(
    start_value=tokenizer.cls_token_id,
    end_value=tokenizer.sep_token_id,
    # Note: This cannot be longer than the preset's `sequence_length`, and there
    # is no check for a custom preprocessor!
    sequence_length=64,
)


# This function that takes a text sample `x` and its
# corresponding label `y` as input and converts the
# text into a format suitable for input into a BERT model.
def preprocessor(x, y):
    token_ids, segment_ids = packer(tokenizer(x))
    x = {
        "token_ids": token_ids,
        "segment_ids": segment_ids,
        "padding_mask": token_ids != 0,
    }
    return x, y


imdb_train_preprocessed = imdb_train.map(preprocessor, tf.data.AUTOTUNE).prefetch(
    tf.data.AUTOTUNE
)
imdb_test_preprocessed = imdb_test.map(preprocessor, tf.data.AUTOTUNE).prefetch(
    tf.data.AUTOTUNE
)

# Preprocessed example
print(imdb_train_preprocessed.unbatch().take(1).get_single_element())
({'token_ids': <tf.Tensor: shape=(64,), dtype=int32, numpy=
array([  101,  2023,  2003,  2941,  2028,  1997,  2026,  5440,  3152,
        1010,  1045,  2052, 16755,  2008,  3071, 12197,  2009,  1012,
        2045,  2003,  2070,  2307,  3772,  1999,  2009,  1998,  2009,
        3065,  2008,  2025,  2035,  1000,  2204,  1000,  3152,  2024,
        2137,  1012,  1012,  1012,  1012,   102,     0,     0,     0,
           0,     0,     0,     0,     0,     0,     0,     0,     0,
           0,     0,     0,     0,     0,     0,     0,     0,     0,
           0], dtype=int32)>, 'segment_ids': <tf.Tensor: shape=(64,), dtype=int32, numpy=
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
      dtype=int32)>, 'padding_mask': <tf.Tensor: shape=(64,), dtype=bool, numpy=
array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True, False, False, False,
       False, False, False, False, False, False, False, False, False,
       False, False, False, False, False, False, False, False, False,
       False])>}, <tf.Tensor: shape=(), dtype=int32, numpy=1>)

使用自定义模型进行微调

drawing

对于更高级的应用,可能无法获得合适的任务Model。在这种情况下,我们提供了对主干Model的直接访问权限,该模型具有自己的from_preset构造函数,并且可以与自定义Layer组合使用。详细示例可在我们的迁移学习指南中找到。

主干Model不包含自动预处理,但可以与使用相同预设的匹配预处理器配对,如前一个工作流程所示。

在这个工作流程中,我们尝试冻结主干模型并添加两个可训练的 Transformer 层以适应新的输入。

注意:我们可以忽略关于pooled_dense层梯度的警告,因为我们使用的是 BERT 的序列输出。

preprocessor = keras_hub.models.BertPreprocessor.from_preset("bert_tiny_en_uncased")
backbone = keras_hub.models.BertBackbone.from_preset("bert_tiny_en_uncased")

imdb_train_preprocessed = (
    imdb_train.map(preprocessor, tf.data.AUTOTUNE).cache().prefetch(tf.data.AUTOTUNE)
)
imdb_test_preprocessed = (
    imdb_test.map(preprocessor, tf.data.AUTOTUNE).cache().prefetch(tf.data.AUTOTUNE)
)

backbone.trainable = False
inputs = backbone.input
sequence = backbone(inputs)["sequence_output"]
for _ in range(2):
    sequence = keras_hub.layers.TransformerEncoder(
        num_heads=2,
        intermediate_dim=512,
        dropout=0.1,
    )(sequence)
# Use [CLS] token output to classify
outputs = keras.layers.Dense(2)(sequence[:, backbone.cls_token_index, :])

model = keras.Model(inputs, outputs)
model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.AdamW(5e-5),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
    jit_compile=True,
)
model.summary()
model.fit(
    imdb_train_preprocessed,
    validation_data=imdb_test_preprocessed,
    epochs=3,
)
Model: "functional_1"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)         Output Shape       Param #  Connected to         ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ padding_mask        │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ segment_ids         │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ token_ids           │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ bert_backbone_3     │ [(None, 128),     │ 4,385,… │ padding_mask[0][0],  │
│ (BertBackbone)      │ (None, None,      │         │ segment_ids[0][0],   │
│                     │ 128)]             │         │ token_ids[0][0]      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ transformer_encoder │ (None, None, 128) │ 198,272 │ bert_backbone_3[0][ │
│ (TransformerEncode… │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ transformer_encode… │ (None, None, 128) │ 198,272 │ transformer_encoder… │
│ (TransformerEncode… │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ get_item_4          │ (None, 128)       │       0 │ transformer_encoder… │
│ (GetItem)           │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ dense (Dense)       │ (None, 2)         │     258 │ get_item_4[0][0]     │
└─────────────────────┴───────────────────┴─────────┴──────────────────────┘
 Total params: 4,782,722 (18.24 MB)
 Trainable params: 396,802 (1.51 MB)
 Non-trainable params: 4,385,920 (16.73 MB)
Epoch 1/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 17s 10ms/step - loss: 0.6208 - sparse_categorical_accuracy: 0.6612 - val_loss: 0.6119 - val_sparse_categorical_accuracy: 0.6758
Epoch 2/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 12s 8ms/step - loss: 0.5324 - sparse_categorical_accuracy: 0.7347 - val_loss: 0.5484 - val_sparse_categorical_accuracy: 0.7320
Epoch 3/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 12s 8ms/step - loss: 0.4735 - sparse_categorical_accuracy: 0.7723 - val_loss: 0.4874 - val_sparse_categorical_accuracy: 0.7742

<keras.src.callbacks.history.History at 0x7f2790170220>

尽管该模型的可训练参数仅为我们的BertClassifier模型的 10%,但它仍实现了合理的准确率。即使考虑到缓存的预处理,每个训练步骤也大约只需要 1/3 的时间。


预训练主干模型

drawing

您是否有权访问您所在领域的大型未标记数据集?它们的大小是否与用于训练流行主干(如 BERT、RoBERTa 或 GPT2)的数据集(XX+ GiB)相似?如果是这样,您可能会受益于您自己的主干模型的特定领域预训练。

NLP 模型通常在语言建模任务上进行预训练,即根据输入句子中可见的词预测被掩盖的词。例如,给定输入"The fox [MASK] over the [MASK] dog",模型可能会被要求预测["jumped", "lazy"]。然后,该模型的较低层被打包为主干,以与与新任务相关的层组合使用。

KerasHub 库提供了 SoTA 主干标记器,以便从头开始训练,无需预设。

在这个工作流程中,我们使用 IMDB 影评文本对 BERT 主干进行预训练。我们跳过了“下一个句子预测”(NSP)损失,因为它增加了数据处理的复杂性,并且在 RoBERTa 等后来的模型中被删除了。有关如何复制原始论文的分步详细信息,请参阅我们的端到端Transformer 预训练

预处理

# All BERT `en` models have the same vocabulary, so reuse preprocessor from
# "bert_tiny_en_uncased"
preprocessor = keras_hub.models.BertPreprocessor.from_preset(
    "bert_tiny_en_uncased",
    sequence_length=256,
)
packer = preprocessor.packer
tokenizer = preprocessor.tokenizer

# keras.Layer to replace some input tokens with the "[MASK]" token
masker = keras_hub.layers.MaskedLMMaskGenerator(
    vocabulary_size=tokenizer.vocabulary_size(),
    mask_selection_rate=0.25,
    mask_selection_length=64,
    mask_token_id=tokenizer.token_to_id("[MASK]"),
    unselectable_token_ids=[
        tokenizer.token_to_id(x) for x in ["[CLS]", "[PAD]", "[SEP]"]
    ],
)


def preprocess(inputs, label):
    inputs = preprocessor(inputs)
    masked_inputs = masker(inputs["token_ids"])
    # Split the masking layer outputs into a (features, labels, and weights)
    # tuple that we can use with keras.Model.fit().
    features = {
        "token_ids": masked_inputs["token_ids"],
        "segment_ids": inputs["segment_ids"],
        "padding_mask": inputs["padding_mask"],
        "mask_positions": masked_inputs["mask_positions"],
    }
    labels = masked_inputs["mask_ids"]
    weights = masked_inputs["mask_weights"]
    return features, labels, weights


pretrain_ds = imdb_train.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE).prefetch(
    tf.data.AUTOTUNE
)
pretrain_val_ds = imdb_test.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)

# Tokens with ID 103 are "masked"
print(pretrain_ds.unbatch().take(1).get_single_element())
({'token_ids': <tf.Tensor: shape=(256,), dtype=int32, numpy=
array([  101,   103,  2332,   103,  1006,   103,   103,  2332,  2370,
        1007,   103,  2029,   103,  2402,  2155,  1010, 24159,  2000,
        3541,  7081,  1010,  2424,  2041,  2055,  1996,  9004,  4528,
         103,   103,  2037,  2188,   103,  1996,  2269,  1006,  8512,
        3054,   103,  4246,  1007,  2059,  4858,  1555,  2055,  1996,
       23025, 22911,  8940,  2598,  3458,  1996, 25483,  4528,  2008,
        2038,   103,  1997, 15218,  1011,   103,  1997,   103,  2505,
        3950,  2045,  3310,  2067,  2025,  3243,  2157,  1012,   103,
        7987,  1013,  1028,   103,  7987,  1013,  1028,  2917,   103,
        1000,  5469,  1000,   103,   103,  2041, 22902,  1010, 23979,
        1010,  1998,  1999, 23606,   103,  1998,  4247,  2008,  2126,
        2005,  1037,  2096,  1010,  2007,  1996,   103,  5409,   103,
        2108,  3054,  3211,  4246,  1005,  1055, 22692,  2836,  1012,
        2009,   103,  1037,  2210,  2488,   103,   103,  2203,  1010,
        2007,   103,   103,  9599,  1012,   103,  2391,  1997,  2755,
        1010,  1996,  2878,  3185,  2003,  2428,   103,  1010,   103,
         103,   103,  1045,  2064,  1005,  1056,  3294, 19776,  2009,
        1011,  2012,  2560,  2009,  2038,  2242,  2000,   103,  2009,
       13432,  1012, 11519,  4637,  4616,  2011,  5965,  1043, 11761,
         103,   103,  2004,   103,  7968,  3243,  4793, 11429,  1010,
        1998,  8226,  2665, 18331,  1010,  1219,  1996,  4487, 22747,
        8004, 12165,  4382,  5125,   103,  3597,   103,  2024,  2025,
        2438,  2000,   103,  2417, 21564,  2143,   103,   103,  7987,
        1013,  1028,  1026,   103,  1013,  1028,  2332,  2038,   103,
        5156, 12081,  2004,  1996,   103,  1012,  1026, 14216,   103,
         103,  1026,  7987,  1013,  1028,   184,  2011,  1037,  8297,
        2036,   103,  2011,  2984,   103,  1006,  2003,  2009,  2151,
        4687,  2008,  2016,  1005,  1055,  2018,  2053,  7731,   103,
         103,  2144,  1029,   102], dtype=int32)>, 'segment_ids': <tf.Tensor: shape=(256,), dtype=int32, numpy=
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], dtype=int32)>, 'padding_mask': <tf.Tensor: shape=(256,), dtype=bool, numpy=
array([ True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True,  True,  True,  True,  True,  True,
        True,  True,  True,  True])>, 'mask_positions': <tf.Tensor: shape=(64,), dtype=int64, numpy=
array([  1,   3,   5,   6,  10,  12,  13,  27,  28,  31,  37,  42,  51,
        55,  59,  61,  65,  71,  75,  80,  83,  84,  85,  94, 105, 107,
       108, 118, 122, 123, 127, 128, 131, 141, 143, 144, 145, 149, 160,
       167, 170, 171, 172, 174, 176, 185, 193, 195, 200, 204, 205, 208,
       210, 215, 220, 223, 224, 225, 230, 231, 235, 238, 251, 252])>}, <tf.Tensor: shape=(64,), dtype=int32, numpy=
array([ 4459,  6789, 22892,  2011,  1999,  1037,  2402,  2485,  2000,
        1012,  3211,  2041,  9004,  4204,  2069,  2607,  3310,  1026,
        1026,  2779,  1000,  3861,  4627,  1010,  7619,  5783,  2108,
        4152,  2646,  1996, 15958, 14888,  1999, 14888,  2029,  2003,
        2339,  1056,  2191,  2011, 11761,  2638,  1010,  1996,  2214,
        2004, 14674,  2860,  2428,  1012,  1026,  1028,  7987,  2010,
        2704,  7987,  1013,  1028,  2628,  2011,  2856, 12838,  2143,
        2147], dtype=int32)>, <tf.Tensor: shape=(64,), dtype=float16, numpy=
array([1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
       1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
       1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.,
       1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.], dtype=float16)>)

预训练模型

# BERT backbone
backbone = keras_hub.models.BertBackbone(
    vocabulary_size=tokenizer.vocabulary_size(),
    num_layers=2,
    num_heads=2,
    hidden_dim=128,
    intermediate_dim=512,
)

# Language modeling head
mlm_head = keras_hub.layers.MaskedLMHead(
    token_embedding=backbone.token_embedding,
)

inputs = {
    "token_ids": keras.Input(shape=(None,), dtype=tf.int32, name="token_ids"),
    "segment_ids": keras.Input(shape=(None,), dtype=tf.int32, name="segment_ids"),
    "padding_mask": keras.Input(shape=(None,), dtype=tf.int32, name="padding_mask"),
    "mask_positions": keras.Input(shape=(None,), dtype=tf.int32, name="mask_positions"),
}

# Encoded token sequence
sequence = backbone(inputs)["sequence_output"]

# 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 = mlm_head(sequence, mask_positions=inputs["mask_positions"])

# Define and compile our pretraining model.
pretraining_model = keras.Model(inputs, outputs)
pretraining_model.summary()
pretraining_model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.AdamW(learning_rate=5e-4),
    weighted_metrics=[keras.metrics.SparseCategoricalAccuracy()],
    jit_compile=True,
)

# Pretrain on IMDB dataset
pretraining_model.fit(
    pretrain_ds,
    validation_data=pretrain_val_ds,
    epochs=3,  # Increase to 6 for higher accuracy
)
Model: "functional_3"
┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Layer (type)         Output Shape       Param #  Connected to         ┃
┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ mask_positions      │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ padding_mask        │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ segment_ids         │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ token_ids           │ (None, None)      │       0 │ -                    │
│ (InputLayer)        │                   │         │                      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ bert_backbone_4     │ [(None, 128),     │ 4,385,… │ mask_positions[0][0… │
│ (BertBackbone)      │ (None, None,      │         │ padding_mask[0][0],  │
│                     │ 128)]             │         │ segment_ids[0][0],   │
│                     │                   │         │ token_ids[0][0]      │
├─────────────────────┼───────────────────┼─────────┼──────────────────────┤
│ masked_lm_head      │ (None, None,      │ 3,954,… │ bert_backbone_4[0][ │
│ (MaskedLMHead)      │ 30522)            │         │ mask_positions[0][0] │
└─────────────────────┴───────────────────┴─────────┴──────────────────────┘
 Total params: 4,433,210 (16.91 MB)
 Trainable params: 4,433,210 (16.91 MB)
 Non-trainable params: 0 (0.00 B)
Epoch 1/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 22s 12ms/step - loss: 5.7032 - sparse_categorical_accuracy: 0.0566 - val_loss: 5.0685 - val_sparse_categorical_accuracy: 0.1044
Epoch 2/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 13s 8ms/step - loss: 5.0701 - sparse_categorical_accuracy: 0.1096 - val_loss: 4.9363 - val_sparse_categorical_accuracy: 0.1239
Epoch 3/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 13s 8ms/step - loss: 4.9607 - sparse_categorical_accuracy: 0.1240 - val_loss: 4.7913 - val_sparse_categorical_accuracy: 0.1417

<keras.src.callbacks.history.History at 0x7f2738299330>

预训练后,保存您的backbone子模型以在新任务中使用!


从头开始构建和训练自己的 Transformer

drawing

想要实现新颖的 Transformer 架构?KerasHub 库提供了用于在我们models API 中构建 SoTA 架构的所有底层模块。这包括keras_hub.tokenizers API,它允许您使用WordPieceTokenizerBytePairTokenizerSentencePieceTokenizer训练自己的子词标记器。

在这个工作流程中,我们使用 IMDB 数据训练了一个自定义标记器,并设计了一个具有自定义 Transformer 架构的主干。为简单起见,我们随后直接在分类任务上进行训练。有兴趣了解更多详细信息?我们在keras.io上编写了有关预训练和微调自定义 Transformer 的完整指南。

从 IMDB 数据训练自定义词汇表

vocab = keras_hub.tokenizers.compute_word_piece_vocabulary(
    imdb_train.map(lambda x, y: x),
    vocabulary_size=20_000,
    lowercase=True,
    strip_accents=True,
    reserved_tokens=["[PAD]", "[START]", "[END]", "[MASK]", "[UNK]"],
)
tokenizer = keras_hub.tokenizers.WordPieceTokenizer(
    vocabulary=vocab,
    lowercase=True,
    strip_accents=True,
    oov_token="[UNK]",
)

使用自定义标记器预处理数据

packer = keras_hub.layers.StartEndPacker(
    start_value=tokenizer.token_to_id("[START]"),
    end_value=tokenizer.token_to_id("[END]"),
    pad_value=tokenizer.token_to_id("[PAD]"),
    sequence_length=512,
)


def preprocess(x, y):
    token_ids = packer(tokenizer(x))
    return token_ids, y


imdb_preproc_train_ds = imdb_train.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)
imdb_preproc_val_ds = imdb_test.map(
    preprocess, num_parallel_calls=tf.data.AUTOTUNE
).prefetch(tf.data.AUTOTUNE)

print(imdb_preproc_train_ds.unbatch().take(1).get_single_element())
(<tf.Tensor: shape=(512,), dtype=int32, numpy=
array([    1,   102,    11,    61,    43,   771,    16,   340,   916,
        1259,   155,    16,   135,   207,    18,   501, 10568,   344,
          16,    51,   206,   612,   211,   232,    43,  1094,    17,
         215,   155,   103,   238,   202,    18,   111,    16,    51,
         143,  1583,   131,   100,    18,    32,   101,    19,    34,
          32,   101,    19,    34,   102,    11,    61,    43,   155,
         105,  5337,    99,   120,     6,  1289,     6,   129,    96,
         526,    18,   111,    16,   193,    51,   197,   102,    16,
          51,   252,    11,    62,   167,   104,   642,    98,     6,
        8572,     6,   154,    51,   153,  1464,   119,  3005,   990,
        2393,    18,   102,    11,    61,   233,   404,   103,   104,
         110,    18,    18,    18,   233,  1259,    18,    18,    18,
         154,    51,   659, 16273,   867,   192,  1632,   133,   990,
        2393,    18,    32,   101,    19,    34,    32,   101,    19,
          34,    96,   110,  2886,   761,   114,  4905,   293, 12337,
          97,  2375,    18,   113,   143,   158,   179,   104,  4905,
         610,    16, 12585,    97,   516,   725,    18,   113,   323,
          96,   651,   146,   104,   207, 17649,    16,    96,   176,
       16022,   136,    16,  1414,   136,    18,   113,   323,    96,
        2184,    18,    97,   150,   651,    51,   242,   104,   100,
       11722,    18,   113,   151,   543,   102,   171,   115,  1081,
         103,    96,   222,    18,    18,    18,    18,   102,   659,
        1081,    18,    18,    18,   102,    11,    61,   115,   299,
          18,   113,   323,    96,  1579,    98,   203,  4438,  2033,
         103,    96,   222,    18,    18,    18,    32,   101,    19,
          34,    32,   101,    19,    34,   111,    16,    51,   455,
         174,    99,   859,    43,  1687,  3330,    99,   104,  1021,
          18,    18,    18,    51,   181,    11,    62,   214,   138,
          96,   155,   100,   115,   916,    14,  1286,    14,    99,
         296,    96,   642,   105,   224,  4598,   117,  1289,   156,
         103,   904,    16,   111,   115,   103,  1628,    18,   113,
         181,    11,    62,   119,    96,  1054,   155,    16,   111,
         156, 14665,    18,   146,   110,   139,   742,    16,    96,
        4905,   293, 12337,    97,  7042,  1104,   106,   557,   103,
         366,    18,   128,    16,   150,  2446,   135,    96,   960,
          98,    96,  4905,    18,   113,   323,   156,    43,  1174,
         293,   188,    18,    18,    18,    43,   639,   293,    96,
         455,   108,   207,    97,  1893,    99,  1081,   104,  4905,
          18,    51,   194,   104,   440,    98, 12337,    99,  7042,
        1104,   654,   122,    30,     6,    51,   276,    99,   663,
          18,    18,    18,    97,   138,   113,   207,   163,    16,
         113,   171,   172,   107,    51,  1027,   113,     6,    18,
          32,   101,    19,    34,    32,   101,    19,    34,   104,
         110,   171,   333, 10311,   141,  1311,   135,   140,   100,
         207,    97,   140,   100,    99,   120,  1632,    18,    18,
          18,    97,   210,    11,    61,    96,  6236,   293,   188,
          18,    51,   181,    11,    62,   214,   138,    96,   421,
          98,   104,   110,   100,     6,   207, 14129,   122,    18,
          18,    18,   151,  1128,    97,  1632,  1675,     6,   133,
           6,   207,   100,   404,    18,    18,    18,   150,   646,
         179,   133,   210,     6,    18,   111,   103,   152,   744,
          16,   104,   110,   100,   557,    43,  1120,   108,    96,
         701,   382,   105,   102,   260,   113,   194,    18,    18,
          18,     2,     0,     0,     0,     0,     0,     0,     0,
           0,     0,     0,     0,     0,     0,     0,     0,     0,
           0,     0,     0,     0,     0,     0,     0,     0,     0,
           0,     0,     0,     0,     0,     0,     0,     0],
      dtype=int32)>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)

设计一个小型 Transformer

token_id_input = keras.Input(
    shape=(None,),
    dtype="int32",
    name="token_ids",
)
outputs = keras_hub.layers.TokenAndPositionEmbedding(
    vocabulary_size=len(vocab),
    sequence_length=packer.sequence_length,
    embedding_dim=64,
)(token_id_input)
outputs = keras_hub.layers.TransformerEncoder(
    num_heads=2,
    intermediate_dim=128,
    dropout=0.1,
)(outputs)
# Use "[START]" token to classify
outputs = keras.layers.Dense(2)(outputs[:, 0, :])
model = keras.Model(
    inputs=token_id_input,
    outputs=outputs,
)

model.summary()
Model: "functional_5"
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Layer (type)                     Output Shape                  Param # ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ token_ids (InputLayer)          │ (None, None)              │          0 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ token_and_position_embedding    │ (None, None, 64)          │  1,259,648 │
│ (TokenAndPositionEmbedding)     │                           │            │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ transformer_encoder_2           │ (None, None, 64)          │     33,472 │
│ (TransformerEncoder)            │                           │            │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ get_item_6 (GetItem)            │ (None, 64)                │          0 │
├─────────────────────────────────┼───────────────────────────┼────────────┤
│ dense_1 (Dense)                 │ (None, 2)                 │        130 │
└─────────────────────────────────┴───────────────────────────┴────────────┘
 Total params: 1,293,250 (4.93 MB)
 Trainable params: 1,293,250 (4.93 MB)
 Non-trainable params: 0 (0.00 B)

直接在分类目标上训练 Transformer

model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.AdamW(5e-5),
    metrics=[keras.metrics.SparseCategoricalAccuracy()],
    jit_compile=True,
)
model.fit(
    imdb_preproc_train_ds,
    validation_data=imdb_preproc_val_ds,
    epochs=3,
)
Epoch 1/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 8s 4ms/step - loss: 0.7790 - sparse_categorical_accuracy: 0.5367 - val_loss: 0.4420 - val_sparse_categorical_accuracy: 0.8120
Epoch 2/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 5s 3ms/step - loss: 0.3654 - sparse_categorical_accuracy: 0.8443 - val_loss: 0.3046 - val_sparse_categorical_accuracy: 0.8752
Epoch 3/3
 1563/1563 ━━━━━━━━━━━━━━━━━━━━ 5s 3ms/step - loss: 0.2471 - sparse_categorical_accuracy: 0.9019 - val_loss: 0.3060 - val_sparse_categorical_accuracy: 0.8748

<keras.src.callbacks.history.History at 0x7f26d032a4d0>

令人兴奋的是,我们的自定义分类器与微调"bert_tiny_en_uncased"的性能相似!要看到预训练的优势并超过 90% 的准确率,我们需要使用更大的预设,例如"bert_base_en_uncased"