热启动嵌入层矩阵

在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 上查看 下载笔记本

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

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

嵌入矩阵

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

词汇表

一组唯一单词被称为词汇表。要构建文本模型，您需要选择一个固定的词汇表。通常，您会从数据集中的高频词构建词汇表。词汇表允许我们通过 ID 序列来表示每一段文本，您可以在嵌入矩阵中查找这些 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.keras.utils.text_dataset_from_directory 创建一个 tf.data.Dataset。您可以在此 文本分类教程 中阅读有关使用此实用程序的更多信息。

使用 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

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

您也可以从检查点（checkpoint）加载模型并更新模型架构，如下所示。

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

继续进行热启动训练

注意训练是如何热启动的。第一轮训练的准确率约为 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 语言理解模型 教程。