使用 TensorFlow Hub 进行文本分类:电影评论

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

此笔记本使用电影评论的文本将电影评论分类为正面负面。这是一个二元(或两类)分类的示例,这是一种重要且应用广泛的机器学习问题。

本教程演示了使用 TensorFlow Hub 和 Keras 进行迁移学习的基本应用。

它使用 IMDB 数据集,其中包含来自 互联网电影数据库 的 50,000 条电影评论的文本。这些评论被分成 25,000 条用于训练,25,000 条用于测试。训练集和测试集是平衡的,这意味着它们包含相同数量的正面和负面评论。

此笔记本使用 tf.keras(一个用于在 TensorFlow 中构建和训练模型的高级 API)和 tensorflow_hub(一个用于使用一行代码从 TFHub 加载训练模型的库)。有关使用 tf.keras 进行更高级的文本分类教程,请参阅 MLCC 文本分类指南

pip install tensorflow-hub
pip install tensorflow-datasets
import os
import numpy as np

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds

print("Version: ", tf.__version__)
print("Eager mode: ", tf.executing_eagerly())
print("Hub version: ", hub.__version__)
print("GPU is", "available" if tf.config.list_physical_devices("GPU") else "NOT AVAILABLE")

下载 IMDB 数据集

IMDB 数据集可在 imdb reviewsTensorFlow 数据集 上获取。以下代码将 IMDB 数据集下载到您的机器(或 colab 运行时)。

# Split the training set into 60% and 40% to end up with 15,000 examples
# for training, 10,000 examples for validation and 25,000 examples for testing.
train_data, validation_data, test_data = tfds.load(
    name="imdb_reviews", 
    split=('train[:60%]', 'train[60%:]', 'test'),
    as_supervised=True)

探索数据

让我们花点时间了解数据的格式。每个示例都是代表电影评论的句子和相应的标签。句子没有经过任何预处理。标签是一个整数,值为 0 或 1,其中 0 表示负面评论,1 表示正面评论。

让我们打印前 10 个示例。

train_examples_batch, train_labels_batch = next(iter(train_data.batch(10)))
train_examples_batch

我们也打印前 10 个标签。

train_labels_batch

构建模型

神经网络是通过堆叠层创建的,这需要三个主要的架构决策

  • 如何表示文本?
  • 模型中要使用多少层?
  • 每层要使用多少个隐藏单元

在本例中,输入数据由句子组成。要预测的标签是 0 或 1。

表示文本的一种方法是将句子转换为嵌入向量。使用预训练的文本嵌入作为第一层,这将具有三个优点

  • 您不必担心文本预处理,
  • 从迁移学习中获益,
  • 嵌入具有固定大小,因此更易于处理。

在本例中,您使用来自 TensorFlow Hub预训练文本嵌入模型,名为 google/nnlm-en-dim50/2

本教程中可以使用来自 TFHub 的许多其他预训练文本嵌入

还有更多!在 TFHub 上查找更多 文本嵌入模型

让我们首先创建一个 Keras 层,该层使用 TensorFlow Hub 模型来嵌入句子,并在几个输入示例上试用它。请注意,无论输入文本的长度如何,嵌入的输出形状都是:(num_examples, embedding_dimension)

embedding = "https://tfhub.dev/google/nnlm-en-dim50/2"
hub_layer = hub.KerasLayer(embedding, input_shape=[], 
                           dtype=tf.string, trainable=True)
hub_layer(train_examples_batch[:3])

现在让我们构建完整的模型

model = tf.keras.Sequential()
model.add(hub_layer)
model.add(tf.keras.layers.Dense(16, activation='relu'))
model.add(tf.keras.layers.Dense(1))

model.summary()

层按顺序堆叠以构建分类器

  1. 第一层是 TensorFlow Hub 层。此层使用预训练的 Saved Model 将句子映射到其嵌入向量。您正在使用的预训练文本嵌入模型 (google/nnlm-en-dim50/2) 将句子拆分为标记,嵌入每个标记,然后组合嵌入。结果维度为:(num_examples, embedding_dimension)。对于此 NNLM 模型,embedding_dimension 为 50。
  2. 此固定长度的输出向量通过具有 16 个隐藏单元的全连接 (Dense) 层传递。
  3. 最后一层是与单个输出节点密集连接的。

让我们编译模型。

损失函数和优化器

模型需要损失函数和优化器进行训练。由于这是一个二元分类问题,并且模型输出 logits(具有线性激活的单单元层),因此您将使用 binary_crossentropy 损失函数。

这不是损失函数的唯一选择,例如,您可以选择 mean_squared_error。但是,通常情况下,binary_crossentropy 更适合处理概率,它测量概率分布之间的“距离”,或者在我们的例子中,测量地面真实分布和预测之间的距离。

稍后,当您探索回归问题(例如,预测房价)时,您将看到如何使用另一个称为均方误差的损失函数。

现在,配置模型以使用优化器和损失函数

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

训练模型

以 512 个样本的迷你批次训练模型 10 个 epochs。这表示对 x_trainy_train 张量中的所有样本进行 10 次迭代。在训练期间,监控模型在验证集中的 10,000 个样本上的损失和准确率

history = model.fit(train_data.shuffle(10000).batch(512),
                    epochs=10,
                    validation_data=validation_data.batch(512),
                    verbose=1)

评估模型

让我们看看模型的性能。将返回两个值。损失(一个代表我们误差的数字,值越低越好)和准确率。

results = model.evaluate(test_data.batch(512), verbose=2)

for name, value in zip(model.metrics_names, results):
  print("%s: %.3f" % (name, value))

这种相当简单的做法实现了大约 87% 的准确率。使用更高级的方法,模型应该更接近 95% 的准确率。

进一步阅读

# MIT License
#
# Copyright (c) 2017 François Chollet
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.