Warm-start embedding layer matrix

##### Copyright 2022 The TensorFlow Authors.


# @title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Warm-start embedding layer matrix#

在 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

加载数据集#

本教程将使用 Large Movie Review Dataset。您将在此数据集上训练情感分类器模型，在此过程中将从头开始学习嵌入向量。请参阅加载文本教程以了解详情。

使用 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_model_input = tf.keras.layers.Input(dtype=tf.string, shape=(1,))
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").get_weights()[0]
)
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").get_weights()[0].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]
)

继续温启动训练#

注意训练是如何温启动的。第一个周期的准确率约为 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 模型教程。