使用 TensorFlow Hub 和估算器构建文本分类模型

关注

原创最新推荐文章于 2026-09-07 18:40:30 发布 984 阅读 AI 写同款· GEO 优化 ›

TensorFlow Hub安装及使用 1.介绍 TensorFlow Hub是一个库，用于实现机器学习模型可重用部分的发行，发行和使用。一个模块是一个独立的tensorflow图，还有它的权重等，可以在迁移学习过程中的不同的任务里重复使用。模块包含使用大型数据集对任务进行预训练的变量。通过在相关任务上重用模块，您可以：用较小的数据集训练模型提高泛化能力显著加快训练速度这有一个例子，使用英语嵌入模块，将一个字符串数组... 阅读详情

文 / 开发技术推广工程师 Sara Robinson

来源 | TensorFlow 公众号

将迁移学习应用于计算机视觉模型的例子很常见，但如果将其用于文本分类，情况又会如何呢？不妨进入 TensorFlow Hub，这是通过迁移学习来增强 TF 模型的内容库。迁移学习是这样一种过程：选用已借助大量数据训练好的已有模型的权重和变量，并将其运用于自己的数据和预测任务。

迁移学习具有诸多好处，其中之一就是您无需像从头开始训练时一样提供自己的大量训练数据。但这些已有模型又来自何处呢？这正是 TensorFlow Hub 大显身手之处：它可以为各类模型（图像、文本等）提供现有模型检查点的完整存储区。在本篇博文中，我将介绍如何使用 TensorFlow Hub 文本模块构建一个模型，以根据相关描述预测电影类型。

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

导入数据并进行预处理

在这个模型中，我们将使用 Kaggle 电影数据集这个出色的公共领域资源。该数据集包含逾 45000 部电影的数据，每部电影均有大量相关数据。为简便起见，我们仅使用该数据集中的电影描述（称作 “简介”）和电影类型。下面是该数据集在 Kaggle 中的预览：
注： Kaggle 电影数据集链接
https://www.kaggle.com/rounakbanik/the-movies-dataset/home

这里写图片描述

首先，我们将导入要使用的内容库，以构建此模型：

1    import numpy as np 
2    import pandas as pd
3
4    import tensorflow as tf
5    import tensorflow_hub as hub
6
7    from sklearn.preprocessing import MultiLabelBinarizer

我已经将这个数据集的 CSV 文件放在一个公共 Cloud Storage 存储分区中。我们可以运行如下命令，将数据下载到我们的 Colab 实例中，并读取为 Pandas dataframe 格式：

1    !wget 'https://storage.googleapis.com/movies_data/movies_metadata.csv'
2    data = pd.read_csv('movies_metadata.csv')
3
4    descriptions = data['overview']
5    genres = data['genres']

为简便起见，我们将可能的类型限定为下列几种：

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

我们将数据集限定为这些类型中描述非空白的电影，然后按照 80% 训练和 20% 测试的比例将数据划分为训练数据集和测试数据集：

1    train_size = int(len(descriptions) * .8)
2
3    train_descriptions = descriptions[:train_size]
4    train_genres = genres[:train_size]
5
6    test_descriptions = descriptions[train_size:]
7    test_genres = genres[train_size:]

使用 TF Hub 构建嵌入层

以 TF Hub 创建嵌入层时仅需使用极少量代码。我们的模型仅有一个特征（描述），并将表示为一个嵌入列。文本嵌入提供了在向量空间表示文本内容的方法，如此一来，嵌入空间中相似的字词或句子会更靠近（您可在此处阅读更多相关内容）。您可以完全使用自己的数据从头开始构建文本嵌入向量。TF Hub 可提供已使用各种文本数据训练过的文本嵌入，因而能够简化这一过程。
注：文本嵌入链接
https://www.tensorflow.org/hub/modules/text

对于英文文本，TF Hub 提供了各种已使用不同种类文本数据训练过的嵌入：

通用语句编码器 (Universal sentence encoder)：用于较长的文本输入
ELMo：使用十亿单词基准 (1B Word Benchmark) 训练过的深度嵌入
神经网络语言模型 (Neural Network Language Model) 嵌入：通过 Google 新闻训练
Word2vec：通过 Wikipedia 训练

您所选择的预训练文本嵌入是您模型中的一个超参数，所以最好用不同的文本嵌入进行试验，看看哪个的准确性最高。先从用与您的文本最接近的文本训练过的模型开始。由于我们的电影描述都是较长的输入，因此，我发现使用通用语句编码器嵌入的准确性最高。这可以将我们的描述编码为高维文本向量。请注意，这一特定模型很大，会占用 1GB 容量。
注：通用语句编码器链接
https://www.tensorflow.org/hub/modules/google/universal-sentence-encoder/2

我们可以使用 hub.text_embedding_column，以一行代码为该层创建一个特征列，并向其传递我们层的名称 (“movie_descriptions”) 和要使用的 TF Hub 模型网址：

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

请注意，该单元正在下载预训练过的嵌入，因此需要一些时间来运行。

此操作最大的好处在于，我们无需进行任何预处理，即可将文本描述馈送至预训练过的字词嵌入。如果从头开始构建此模型，我们就需要自己将描述转换为向量，但使用 TF Hub 列，我们可以将描述字符串直接传递至模型。

将标签变为 multi-hot 编码

由于一部电影往往具有多种类型，所以我们的模型会为每部电影返回多个可能的标签。我们的类型目前是每部电影有一个字符串列表（例如 [‘Action’, ‘Adventure’]）。由于每个标签的长度必须相同，所以我们要将这些列表转换为由 1 和 0（与特定描述中的类型相对应）组成的 multi-hot 向量。动作冒险片的 multi-hot 向量如下所示：
注：多个可能的标签链接
https://en.wikipedia.org/wiki/Multi-label_classification

1    # Genre lookup, each genre corresponds to an index    
2    top_genres = ['Comedy', 'Thriller', 'Romance', 'Action', 'Horror', 'Crime', 'Documentary', 'Adventure', 'Science Fiction']
3
4    # Multi-hot label for an action and adventure movie 
5    [0 0 0 1 0 0 0 1 0]

要用短短几行代码将字符串标签转换为 multi-hot 向量，我们需要使用名为 MultiLabelBinarizer 的 Scikit Learn 实用程序：

1    encoder = MultiLabelBinarizer()    
2    encoder.fit_transform(train_genres)    
3    train_encoded = encoder.transform(train_genres)    
4    test_encoded = encoder.transform(test_genres)    
5    num_classes = len(encoder.classes_)

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

构建并训练 DNNEstimator 模型

针对我们的模型，我们将使用 DNNEstimator 构建能够返回 multi-hot 向量的深度神经网络，这是因为每部电影会具有 0 个或多个可能的标签（这与每个输出正好有一个标签的模型不同）。我们传递至 DNNEstimator 的第一个参数称作 head，且此参数会定义我们的模型预期具有的标签类型。我们希望我们的模型可以输出多个标签，所以我们在这里使用 multi_label_head：

1    multi_label_head = tf.contrib.estimator.multi_label_head( 
2            num_classes,
3            loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE
4    )

现在，当我们实例化 DNNEstimator 时，便可对其进行传递。hidden_units 参数表示我们网络中的层数。此模型有 2 个层，第一层有 64 个神经元，第二层有 10 个。层数和层大小是超参数，所以您应当尝试不同的值，看看哪个最适合您的数据集。最后，我们将特征列传递至估算器。在本例中，我们只有一个特征列（即描述），而且在上文中已将其定义为 TF Hub 嵌入列，所以在此我们可以将其作为列表传递：

1    estimator = tf.contrib.estimator.DNNEstimator(
2            head=multi_label_head, 
3            hidden_units=[64,10],
4            feature_columns=[description_embeddings]
5    )

我们基本准备就绪，很快就可以开始训练模型了。在训练估算器实例之前，我们需要定义训练输入函数。输入函数可以将我们的数据与模型联系起来。在这里，我们将使用 numpy_input_fn，并将我们的数据作为 Numpy 数组馈送至模型：

1    # Format our data for the numpy_input_fn    
2    features = {    
3        "descriptions": np.array(train_descriptions)    
4    }    
5    labels = np.array(train_encoded)    
6    
7    train_input_fn = tf.estimator.inputs.numpy_input_fn(
8            features,    
9            labels,    
10            shuffle=True,    
11            batch_size=32,    
12            num_epochs=20    
13    )

我们输入函数中的 batch_size 和 num_epochs 参数都是超函数。batch_size 可告知我们的模型在一次迭代中会有多少示例传递至模型，而 num_epochs 是指我们的模型完成整个训练集的次数。

现在可以开始训练我们的模型了。只用一行代码即可：

1    estimator.train(input_fn=train_input_fn)

为了评估模型的准确性，我们用自己的测试数据创建一个 eval 函数 input_function，然后调用 estimator.evaluate()：

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

此模型的 AUC 达到 91.5%，而查准率/查全率为 74%。您的结果可能稍有不同。

使用我们已训练的模型生成预测结果

现在到了最精彩的部分：根据我们的模型从未见过的数据生成预测结果。首先，我们设置一个包含一些描述的数组（我从 IMDB 中获取这些描述）：

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

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

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

最后，我们可以迭代访问结果，并显示为每部电影找到的前 2 个类型及其置信度值：

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

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