使用 TF HUB 进行模型复用

我们在之前的学习之中学习过如何进行模型的保存与重复的利用,之后我们又学习了如何进行迁移学习。那么今天我们就来认识一个专门用于 TensorFlow 模型复用的库 —— TensorFlow Hub。

1. 什么是 TensorFlow Hub 以及它的安装

TensorFlow Hub 是一个针对可重复使用的机器学习模块的库。 —— 官方定义

也就是说,TensorFlow Hub 内部含有一些已经预训练好的各种常用的网络模型与相应的权重,我们只需要调用其接口就可以得到相应的模型。

TensorFlow Hub 的出现可以说极大地降低了我们使用 TensorFlow 进行迁移学习的难度。因此,当想要使用 TensorFlow 进行迁移学习的时候,我们的第一选择应该是 TensorFlow Hub

TensorFlow Hub 的安装也非常简单,我们只需要使用 pip 安装即可:

pip install tensorflow-hub

倘若下载速度缓慢,可以查阅之前的 TensorFlow 安装教程来使用国内镜像源进行下载。

2. TensorFlow Hub 之中包含哪些模块

TensorFlow HUB 之中包含的模块大致可以分为四个种类,他们分别是:文本模型、图像模型、音频模型、视频模型

1. TensorFlow Hub 之中的文本模型

文本大类中又包含一些小的分类,比如分类、情感分析、语义分析等,这里介绍一些常用的模型:

  • gnews-swivel-20dim/1 模型:文本任务的一个基本的迁移模型,包含 20 维的词向量;
  • nnlm-en-dim50/1 模型:巨大的模型,包含 1M 个词汇量与 50 维的词向量;
  • nnlm-en-dim128/1 模型:非常巨大的模型,包含 1M 个词汇量与 128 维词向量;
  • BERT:迄今为止最大的文本处理模型。

使用以上几种模型,可以满足我们绝大多数的文本任务的迁移学习的需求。

2. TensorFlow Hub 之中的图像模型

图像模型最主要是进行特征提取而使用的模型,我们一般会摒弃掉这些模型的最后几个网络层,从而达到特征提取的目的。

  • mobilenet_v2:一个小型的迁移学习模型,可以胜任一些基本的任务;
  • arbitrary-image-stylization:一个用来进行图像风格处理的模型;
  • ResNet:规模从小到大都有,是最常用的特征提取模型。

3. TensorFlow Hub 之中的音频模型

在音频的处理之中,我们最常用的网络模型就是 “spice” 模型,该模型能够很好地对音频信号进行特征的提取。

另外,我们可以使用 Librosa 等工具提取出 MFCC 等特征之后,再使用 ResNet 等图像特征提取模型进行下一步的处理。

4. TensorFlow Hub 之中的视频模型

在视频处理之中,我们最常用的方式是将视频分为独立的帧,然后对于每一帧使用图像模型进行特征的提取,从而将其转化为一个图片特征提取问题。

3. 如何使用 TensorFlow Hub

我们要使用 TensorFlow Hub 之中的网络模型,我们只需要知道该模型对应的地址即可,而这些地址可以通过 TensorFlow Hub 官方网站进行查询。

在使用的过程之中,我们只需要经过如下几个步骤:

  • 找到模型对应的 URL
  • 使用 hub 对应的 API 来下载该网络模型并进行载入,比如我们常用的 Keras 模型可以通过 hub.KerasLayer API 来进行调用
  • 将加载的预训练模型加入到自己的网络之中
import tensorflow as tf
import tensorflow-hub as hub

model = "https://tfhub.dev/google/tf2-preview/gnews-swivel-20dim/1"
hub_layer = hub.KerasLayer(model, output_shape=[20], input_shape=[],  dtype=tf.string, trainable=True)
model = tf.keras.Sequential()
model.add(hub_layer)
model.add(tf.keras.layers.Dense(16, activation='relu'))
model.add(tf.keras.layers.Dense(1))

如上述代码所示,我们就通过 hub.KerasLayer API 载入了一个预训练模型,并且加入到了自己的模型之中。

4. 使用 TensorFlow Hub 进行迁移学习的文本分类实例

在这一小结,我们会使用 TensorFlow Hub 的一个内置的文本网络模型来进行文本的分类任务的学习。

在这里,我们依然采用 IMDB 分类的任务来进行演示。

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

# 定义基本参数
EPOCHS = 30
BATCH_SIZE = 64

# 获取数据
train_data, test_data = tfds.load(name="imdb_reviews", split=["train", "test"], batch_size=BATCH_SIZE, as_supervised=True)

train_examples, train_labels = tfds.as_numpy(train_data)
test_examples, test_labels = tfds.as_numpy(test_data)


# 定义模型
model = "https://tfhub.dev/google/tf2-preview/gnews-swivel-20dim/1"
hub_layer = hub.KerasLayer(model, output_shape=[20], input_shape=[],  dtype=tf.string, trainable=True)
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()

# 编译模型
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 训练
history = model.fit(train_examples, train_labels, epochs=EPOCHS, batch_size=BATCH_SIZE)

# 测试
results = model.evaluate(test_examples, test_labels)

print(results)

在这里,我们首先使用 tesnorflow_datasets 获取了 IMDB 的数据集;然后我们便定义了我们的模型,这里的模型采用的是上一小节定义的迁移模型,最后我们进行了训练。

最终我们的得到的结果为:

Model: "sequential_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
keras_layer_2 (KerasLayer)   (None, 20)                400020    
_________________________________________________________________
dense_4 (Dense)              (None, 16)                336       
_________________________________________________________________
dense_5 (Dense)              (None, 1)                 17        
=================================================================
Total params: 400,373
Trainable params: 400,373
Non-trainable params: 0
_________________________________________________________________
Epoch 1/30
391/391 [==============================] - 5s 12ms/step - loss: 1.5181 - accuracy: 0.6126
Epoch 2/30
391/391 [==============================] - 5s 13ms/step - loss: 0.7655 - accuracy: 0.6738
.........
Epoch 30/30
391/391 [==============================] - 5s 12ms/step - loss: 0.0454 - accuracy: 0.9963
782/782 [==============================] - 3s 4ms/step - loss: 1.5898 - accuracy: 0.8366
[1.589829683303833, 0.8365600109100342]

可以看到,我们的模型一共包括三层,刚好符合我们的预期,同时我们的模型最终也在测试集合上达到了 83% 的准确率,而这是一个良好的结果。

4. 小结

在这节课的学习之中,我们了解了什么是 TensorFlow Hub,同时我们了解了 TensorFlow Hub 内部所包含的模型的大致类别;然后我们学习了如何使用这些网络模型;最终我们通过一个简单的示例了解了其具体的用法。