作者：开发者布道师 Sara Robinson

我们经常看到迁移学习应用于计算机视觉模型，但对于文本分类呢？进入 TensorFlow Hub，一个用于通过迁移学习增强您的 TF 模型的库。迁移学习是指使用已经过大量数据训练的预先存在的模型的权重和变量，并将其用于您自己的数据和预测任务的过程。

迁移学习的众多好处之一是，您不需要像从头开始一样提供那么多的训练数据。但是，这些预先存在的模型来自哪里呢？这就是 TensorFlow Hub 派上用场的地方：它提供了各种模型（图像、文本等）的现有模型检查点的完整存储库。在本文中，我将引导您使用 TensorFlow Hub 文本模块构建一个模型，根据电影的描述预测其类型。

您可以使用 Colab 在浏览器中运行此模型，无需任何设置。

导入和预处理数据

对于此模型，我们将使用来自 Kaggle 的这个很棒的公共领域 电影数据集。它包含超过 45,000 部电影的数据。每部电影都有很多数据，但为了简单起见，我们只使用此数据集中电影的描述（称为“概述”）及其类型。以下是 Kaggle 中数据集的预览

首先，我们将导入用于构建此模型的库

import numpy as np
import pandas as pd

import tensorflow as tf
import tensorflow_hub as hub

from sklearn.preprocessing import MultiLabelBinarizer

我已经将来自此数据集的 CSV 文件放在一个公共云存储桶中。我们可以运行以下命令将数据下载到我们的 Colab 实例并将其读取为 Pandas 数据框

!wget 'https://storage.googleapis.com/movies_data/movies_metadata.csv'
data = pd.read_csv('movies_metadata.csv')

descriptions = data['overview']
genres = data['genres']

为了简单起见，我们将可能的类型限制在以下几种

top_genres = ['Comedy', 'Thriller', 'Romance', 'Action', 'Horror', 'Crime', 'Documentary', 'Adventure', 'Science Fiction']

我们将数据集限制在这些类型中描述不为空的电影，然后我们可以使用 80% / 20% 的训练/测试分割将数据分成训练集和测试集

train_size = int(len(descriptions) * .8)

train_descriptions = descriptions[:train_size]
train_genres = genres[:train_size]

test_descriptions = descriptions[train_size:]
test_genres = genres[train_size:]

使用 TF Hub 构建嵌入层

使用 TF Hub 创建嵌入需要令人惊讶的少量的代码。我们的模型将只有一个特征（描述），并且它将被表示为一个嵌入列。文本嵌入提供了一种在向量空间中表示文本片段的方法，以便相似的词或句子在嵌入空间中更靠近（您可以在这里阅读更多关于它们的 信息）。您可以使用完全来自您自己的数据从头开始构建文本嵌入向量。TF Hub 通过 提供文本嵌入 来简化此过程，这些文本嵌入已经过各种文本数据的训练。

对于英文文本，TF Hub 提供了在不同类型的文本数据上训练的各种嵌入

• 通用句子编码器：适用于较长的文本输入
  
• ELMo：在 1B 词基准上训练的深度嵌入
  
• 神经网络语言模型嵌入：在 Google 新闻上训练
  
• Word2vec：在维基百科上训练
  

您选择的预训练文本嵌入是模型中的超参数，因此最好尝试不同的嵌入，看看哪个嵌入能产生最高的准确率。从在最接近您的文本上训练的模型开始。由于我们的电影描述是较长的输入，我发现使用 通用句子编码器 嵌入获得了最高的准确率。这将把我们的描述编码成高维文本向量。请注意，此特定模型非常大，将占用 1 GB。

我们可以使用 hub.text_embedding_column 在一行代码中为这一层创建一个特征列，并将我们的层的名称（“movie_descriptions”）和我们将使用的 TF Hub 模型的 URL 传递给它

description_embeddings = hub.text_embedding_column(
  "movie_descriptions", 
  module_spec="https://tfhub.dev/google/universal-sentence-encoder/2"
)

请注意，运行此单元格可能需要一些时间，因为它正在下载预训练的嵌入。

最棒的是，我们不需要进行任何预处理即可将文本描述馈送到预训练的词嵌入中。如果我们从头开始构建此模型，则需要自己将描述转换为向量，但使用 TF Hub 列，我们可以直接将描述字符串传递给模型。

将标签转换为多热编码

由于一部电影通常有多种类型，因此我们的模型将为每部电影返回 多个可能的标签。我们的类型目前是每部电影的字符串列表（例如 ['Action', 'Adventure']）。由于每个标签都需要相同长度，因此我们将这些列表转换为 1 和 0 的多热向量，对应于特定描述中存在的类型。动作和冒险电影的多热向量将如下所示

# Genre lookup, each genre corresponds to an index
top_genres = ['Comedy', 'Thriller', 'Romance', 'Action', 'Horror', 'Crime', 'Documentary', 'Adventure', 'Science Fiction']

# Multi-hot label for an action and adventure movie
[0 0 0 1 0 0 0 1 0]

为了将我们的字符串标签转换为多热向量，我们将在几行代码中使用名为 MultiLabelBinarizer 的 scikit learn 实用程序

encoder = MultiLabelBinarizer()
encoder.fit_transform(train_genres)
train_encoded = encoder.transform(train_genres)
test_encoded = encoder.transform(test_genres)
num_classes = len(encoder.classes_)

您可以打印 encoder.classes_ 以查看模型预测的所有字符串类的列表。

构建和训练 DNNEstimator 模型

对于我们的模型，我们将使用 DNNEstimator 来构建一个深层神经网络，该网络返回一个多热向量，因为每部电影都可以有 0 个或更多个可能的标签（这与每个输入恰好有一个标签的模型不同）。我们传递给 DNNEstimator 的第一个参数称为 head，它定义了模型应期望的标签类型。由于我们希望我们的模型输出多个标签，因此我们将在这里使用 multi_label_head

multi_label_head = tf.contrib.estimator.multi_label_head(
    num_classes,
    loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE
)

现在，我们可以在实例化 DNNEstimator 时将其传递进来。hidden_units 参数指示网络中将有多少层。此模型有 2 层，第一层有 64 个神经元，第二层有 10 个神经元。层数和层大小是一个超参数，因此您应该尝试不同的值以查看哪个值最适合您的数据集。最后，我们将特征列传递给 Estimator。在本例中，我们只有一个（描述），并且我们已经在上面将其定义为 TF Hub 嵌入列，因此我们可以将其作为列表传递到这里

estimator = tf.contrib.estimator.DNNEstimator(
    head=multi_label_head,
    hidden_units=[64,10],
    feature_columns=[description_embeddings]
)

我们几乎准备好训练模型了。在我们可以训练我们的 Estimator 实例之前，我们需要定义我们的训练输入函数。输入函数将我们的数据连接到模型。在这里，我们将使用 numpy_input_fn 并将我们的数据作为 numpy 数组馈送到我们的模型

# Format our data for the numpy_input_fn
features = {
  "descriptions": np.array(train_descriptions)
}
labels = np.array(train_encoded)

train_input_fn = tf.estimator.inputs.numpy_input_fn(
    features, 
    labels, 
    shuffle=True, 
    batch_size=32, 
    num_epochs=20
)

输入函数中的 batch_size 和 num_epochs 参数都是超参数。batch_size 告诉我们的模型在一次迭代期间将多少个示例传递到我们的模型，而 num_epochs 是我们的模型遍历整个训练集的次数。

是时候训练我们的模型了。我们可以用一行代码做到这一点

estimator.train(input_fn=train_input_fn)

为了评估模型的准确率，我们使用测试数据创建一个 eval input_function，并调用 estimator.evaluate()

eval_input_fn = tf.estimator.inputs.numpy_input_fn({"descriptions": np.array(test_descriptions).astype(np.str)}, test_encoded.astype(np.int32), shuffle=False)

estimator.evaluate(input_fn=eval_input_fn)

此模型的 AUC 为 91.5%，精度/召回率为 74%。您的结果可能略有不同。

在训练后的模型上生成预测

是时候进行最棒的部分了：在模型从未见过的數據上生成预测。首先，让我们设置一个包含几个描述的数组（我从 IMDB 获取了这些描述）

raw_test = [
    "An examination of our dietary choices and the food we put in our bodies. Based on Jonathan Safran Foer's memoir.", # Documentary
    "A teenager tries to survive the last week of her disastrous eighth-grade year before leaving to start high school.", # Comedy
    "Ethan Hunt and his IMF team, along with some familiar allies, race against time after a mission gone wrong." # Action, Adventure
]

然后，我们将定义我们的预测输入函数并调用 predict()

predict_input_fn = tf.estimator.inputs.numpy_input_fn({"descriptions": np.array(raw_test).astype(np.str)}, shuffle=False)

results = estimator.predict(predict_input_fn)

最后，我们可以遍历结果并显示每部电影找到的前 2 种类型及其置信度值

for movie_genres in results:
  top_2 = movie_genres['probabilities'].argsort()[-2:][::-1]
  for genre in top_2:
    text_genre = encoder.classes_[genre]
    print(text_genre + ': ' + str(round(movie_genres['probabilities'][genre] * 100, 2)) + '%')

我们的模型能够正确地标记上面所有电影描述。

入门

想要开始使用 TF Hub 构建自己的模型吗？查看 文档 和 教程。您可以在 GitHub 或 Colab 上找到此处概述的模型的完整代码。在以后的文章中，我将介绍如何导出此模型以便在 TensorFlow Serving 或 Cloud ML Engine 上提供服务，以及构建一个应用程序，该应用程序可以根据新的描述生成预测。

如果您有任何问题或反馈，请在 Twitter 上告诉我 @SRobTweets。