预热嵌入层矩阵

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

本教程展示了在更改词汇表时如何使用 tf.keras.utils.warmstart_embedding_matrix API 对文本情感分类进行“预热”训练。

首先，您将使用基本词汇表训练一个简单的 Keras 模型，然后在更新词汇表后继续训练该模型。这称为“预热”训练，您需要针对新词汇表重新映射文本嵌入矩阵。

嵌入矩阵

嵌入提供了一种使用高效、密集的表示方式，其中相似的词汇标记具有相似的编码。它们是可训练参数（模型在训练期间学习的权重，就像模型学习密集层的权重一样）。对于小型数据集，嵌入通常是 8 维的，而对于大型数据集，嵌入可以高达 1024 维。维度较高的嵌入可以捕捉单词之间的细粒度关系，但学习起来可能需要更多数据。

词汇表

一组唯一单词称为词汇表。要构建文本模型，您需要选择一个固定的词汇表。通常，您会从数据集中最常见的单词中构建词汇表。词汇表允许我们通过嵌入矩阵中可以查找的 ID 序列来表示每一段文本。词汇表允许我们通过其中出现的特定单词来表示每一段文本。

为什么要预热嵌入矩阵？

使用一组嵌入训练模型，这些嵌入表示给定的词汇表。如果需要更新或改进模型，则可以通过重复使用上次运行中的权重来显著加快收敛速度。使用上次运行中的嵌入矩阵更加困难。问题在于词汇表的任何更改都会使单词到 ID 的映射无效。

tf.keras.utils.warmstart_embedding_matrix 通过使用来自基础词表的嵌入矩阵为新词表创建一个嵌入矩阵，来解决此问题。如果一个单词同时存在于两个词表中，则基础嵌入向量将被复制到新嵌入矩阵中的正确位置。这允许你在词汇表的大小或顺序发生任何更改后，对训练进行热启动。

设置

pip install --pre -U "tensorflow>2.10"  # Requires 2.11

import io
import numpy as np
import os
import re
import shutil
import string
import tensorflow as tf

from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, Embedding, GlobalAveragePooling1D
from tensorflow.keras.layers import TextVectorization

加载数据集

本教程使用 大型电影评论数据集。你将在该数据集上训练一个情感分类器模型，并在此过程中从头开始学习嵌入。请参阅 加载文本教程 以了解更多信息。

使用 Keras 文件实用程序下载数据集并查看目录。

url = "https://ai.stanford.edu/~amaas/data/sentiment/aclImdb_v1.tar.gz"

dataset = tf.keras.utils.get_file(
    "aclImdb_v1.tar.gz", url, untar=True, cache_dir=".", cache_subdir=""
)

dataset_dir = os.path.join(os.path.dirname(dataset), "aclImdb")
os.listdir(dataset_dir)

train/ 目录具有 pos 和 neg 文件夹，其中分别将电影评论标记为正面和负面。你将使用来自 pos 和 neg 文件夹的评论来训练二元分类模型。

train_dir = os.path.join(dataset_dir, "train")
os.listdir(train_dir)

train 目录还包含其他文件夹，在创建训练集之前应将其删除。

remove_dir = os.path.join(train_dir, "unsup")
shutil.rmtree(remove_dir)

接下来，使用 tf.data.Dataset 创建 tf.keras.utils.text_dataset_from_directory。你可以在 文本分类教程 中阅读有关使用此实用程序的更多信息。

使用 train 目录创建训练集和验证集，其中 20% 用于验证。

batch_size = 1024
seed = 123
train_ds = tf.keras.utils.text_dataset_from_directory(
    "aclImdb/train",
    batch_size=batch_size,
    validation_split=0.2,
    subset="training",
    seed=seed,
)
val_ds = tf.keras.utils.text_dataset_from_directory(
    "aclImdb/train",
    batch_size=batch_size,
    validation_split=0.2,
    subset="validation",
    seed=seed,
)

配置数据集以提高性能

你可以在 Dataset.cache 和 Dataset.prefetch 中了解有关 数据性能指南 中如何将数据缓存到磁盘的更多信息。

AUTOTUNE = tf.data.AUTOTUNE

train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

文本预处理

接下来，定义你的情感分类模型所需的数据集预处理步骤。初始化一个 layers.TextVectorization 层，使用所需的电影评论向量化参数。你可以在 文本分类 教程中了解有关如何使用此层的更多信息。

# Create a custom standardization function to strip HTML break tags '<br />'.
def custom_standardization(input_data):
    lowercase = tf.strings.lower(input_data)
    stripped_html = tf.strings.regex_replace(lowercase, "<br />", " ")
    return tf.strings.regex_replace(
        stripped_html, "[%s]" % re.escape(string.punctuation), ""
    )


# Vocabulary size and number of words in a sequence.
vocab_size = 10000
sequence_length = 100

# Use the text vectorization layer to normalize, split, and map strings to
# integers. Note that the layer uses the custom standardization defined above.
# Set maximum_sequence length as all samples are not of the same length.
vectorize_layer = TextVectorization(
    standardize=custom_standardization,
    max_tokens=vocab_size,
    output_mode="int",
    output_sequence_length=sequence_length,
)

# Make a text-only dataset (no labels) and call `Dataset.adapt` to build the
# vocabulary.
text_ds = train_ds.map(lambda x, y: x)
vectorize_layer.adapt(text_ds)

创建分类模型

使用 Keras Sequential API 定义情感分类模型。

embedding_dim = 16
text_embedding = Embedding(vocab_size, embedding_dim, name="embedding")

text_input = tf.keras.Sequential(
    [vectorize_layer, text_embedding], name="text_input"
)
classifier_head = tf.keras.Sequential(
    [GlobalAveragePooling1D(), Dense(16, activation="relu"), Dense(1)],
    name="classifier_head",
)

model = tf.keras.Sequential([text_input, classifier_head])

编译和训练模型

你将使用 TensorBoard 来可视化包括损失和准确性在内的指标。创建一个 tf.keras.callbacks.TensorBoard。

tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir="logs")

使用 Adam 优化器和 BinaryCrossentropy 损失编译并训练模型。

model.compile(
    optimizer="adam",
    loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
    metrics=["accuracy"],
)

model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=15,
    callbacks=[tensorboard_callback],
)

通过这种方法，模型达到约 85% 的验证准确率

你可以查看模型摘要以了解有关模型的每一层的更多信息。

model.summary()

在 TensorBoard 中可视化模型指标。

# docs_infra: no_execute
%load_ext tensorboard
%tensorboard --logdir logs

词汇重新映射

现在，你将更新词汇表并继续进行热启动训练。

首先，获取基础词汇表和嵌入矩阵。

embedding_weights_base = (
    model.get_layer("text_input").get_layer("embedding").embeddings
)
vocab_base = vectorize_layer.get_vocabulary()

定义一个新的向量化层以生成一个新的更大的词汇表

# Vocabulary size and number of words in a sequence.
vocab_size_new = 10200
sequence_length = 100

vectorize_layer_new = TextVectorization(
    standardize=custom_standardization,
    max_tokens=vocab_size_new,
    output_mode="int",
    output_sequence_length=sequence_length,
)

# Make a text-only dataset (no labels) and call adapt to build the vocabulary.
text_ds = train_ds.map(lambda x, y: x)
vectorize_layer_new.adapt(text_ds)

# Get the new vocabulary
vocab_new = vectorize_layer_new.get_vocabulary()

# View the new vocabulary tokens that weren't in `vocab_base`
set(vocab_base) ^ set(vocab_new)

使用 keras.utils.warmstart_embedding_matrix 实用程序生成更新的嵌入。

# Generate the updated embedding matrix
updated_embedding = tf.keras.utils.warmstart_embedding_matrix(
    base_vocabulary=vocab_base,
    new_vocabulary=vocab_new,
    base_embeddings=embedding_weights_base,
    new_embeddings_initializer="uniform",
)
# Update the model variable
updated_embedding_variable = tf.Variable(updated_embedding)

或者

如果你有一个嵌入矩阵，你想用它来初始化新的嵌入矩阵，请使用 keras.initializers.Constant 作为 new_embeddings 初始化器。将以下代码块复制到代码单元格中以尝试此操作。当你在词汇表中为新单词提供更好的嵌入矩阵初始化时，这将很有帮助。

# generate updated embedding matrix
new_embedding = np.random.rand(len(vocab_new), 16)
updated_embedding = tf.keras.utils.warmstart_embedding_matrix(
            base_vocabulary=vocab_base,
            new_vocabulary=vocab_new,
            base_embeddings=embedding_weights_base,
            new_embeddings_initializer=tf.keras.initializers.Constant(
                new_embedding
            )
        )
# update model variable
updated_embedding_variable = tf.Variable(updated_embedding)

验证嵌入矩阵的形状是否已更改以反映新词汇表。

updated_embedding_variable.shape

现在你有了更新的嵌入矩阵，下一步是更新层权重。

text_embedding_layer_new = Embedding(
    vectorize_layer_new.vocabulary_size(), embedding_dim, name="embedding"
)
text_embedding_layer_new.build(input_shape=[None])
text_embedding_layer_new.embeddings.assign(updated_embedding)
text_input_new = tf.keras.Sequential(
    [vectorize_layer_new, text_embedding_layer_new], name="text_input_new"
)
text_input_new.summary()

# Verify the shape of updated weights
# The new weights shape should reflect the new vocabulary size
text_input_new.get_layer("embedding").embeddings.shape

修改模型架构以使用新的文本向量化层。

你还可以从检查点加载模型，并按如下所示更新模型架构。

warm_started_model = tf.keras.Sequential([text_input_new, classifier_head])
warm_started_model.summary()

您已成功更新模型以接受新词汇。嵌入层已更新，以将旧词汇词映射到旧嵌入，并初始化新词汇词的嵌入以供学习。模型其余部分的学习权重将保持不变。模型已预热启动，以继续从上次中断的地方进行训练。

您现在可以验证重新映射是否有效。获取同时存在于基础词汇和新词汇中的词汇词“the”的索引，并比较嵌入值。它们应该相等。

# New vocab words
base_vocab_index = vectorize_layer("the")[0]
new_vocab_index = vectorize_layer_new("the")[0]
print(
    warm_started_model.get_layer("text_input_new").get_layer("embedding")(
        new_vocab_index
    )
    == embedding_weights_base[base_vocab_index]
)

继续进行预热启动训练

请注意训练是如何预热启动的。第一个 epoch 的准确率约为 85%。这接近于先前训练结束时的准确率。

model.compile(
    optimizer="adam",
    loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
    metrics=["accuracy"],
)
model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=15,
    callbacks=[tensorboard_callback],
)

可视化预热启动训练

# docs_infra: no_execute
%reload_ext tensorboard
%tensorboard --logdir logs

后续步骤

在本教程中，您将学习如何

• 从头开始在小型词汇数据集上训练情感分类模型。
• 当词汇量大小时，更新模型架构并预热启动嵌入矩阵。
• 随着数据集的扩展，持续提高模型准确率

要了解有关嵌入的更多信息，请查看Word2Vec和用于语言理解的 Transformer 模型教程。