使用 TensorFlow Hub 进行迁移学习

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

查看 TF Hub 模型

TensorFlow Hub 是一个预训练 TensorFlow 模型的存储库。

本教程演示了如何

使用 TensorFlow Hub 中的模型与 tf.keras。
使用 TensorFlow Hub 中的图像分类模型。
对模型进行简单的迁移学习，以便为自己的图像类别微调模型。

设置

import numpy as np
import time

import PIL.Image as Image
import matplotlib.pylab as plt

import tensorflow as tf
import tensorflow_hub as hub

import datetime

%load_ext tensorboard

ImageNet 分类器

您将从使用在 ImageNet 基准数据集上预训练的分类器模型开始 - 无需初始训练！

下载分类器

从 TensorFlow Hub 选择一个 MobileNetV2 预训练模型 from TensorFlow Hub，并使用 hub.KerasLayer 将其包装为 Keras 层。来自 TensorFlow Hub 的任何兼容图像分类器模型都可以在此处使用，包括下拉菜单中提供的示例。

mobilenet_v2 ="https://tfhub.dev/google/tf2-preview/mobilenet_v2/classification/4"
inception_v3 = "https://tfhub.dev/google/imagenet/inception_v3/classification/5"

classifier_model = mobilenet_v2

IMAGE_SHAPE = (224, 224)

classifier = tf.keras.Sequential([
    hub.KerasLayer(classifier_model, input_shape=IMAGE_SHAPE+(3,))
])

在单个图像上运行它

下载单个图像以在模型上进行尝试

grace_hopper = tf.keras.utils.get_file('image.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/grace_hopper.jpg')
grace_hopper = Image.open(grace_hopper).resize(IMAGE_SHAPE)
grace_hopper

grace_hopper = np.array(grace_hopper)/255.0
grace_hopper.shape

添加批次维度（使用 np.newaxis）并将图像传递给模型

result = classifier.predict(grace_hopper[np.newaxis, ...])
result.shape

结果是一个包含 1001 个元素的 logits 向量，对图像中每个类别的概率进行评级。

可以使用 tf.math.argmax 找到最高类别 ID

predicted_class = tf.math.argmax(result[0], axis=-1)
predicted_class

解码预测

获取 predicted_class ID（例如 653）并获取 ImageNet 数据集标签以解码预测

labels_path = tf.keras.utils.get_file('ImageNetLabels.txt','https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt')
imagenet_labels = np.array(open(labels_path).read().splitlines())

plt.imshow(grace_hopper)
plt.axis('off')
predicted_class_name = imagenet_labels[predicted_class]
_ = plt.title("Prediction: " + predicted_class_name.title())

简单的迁移学习

但是，如果您想使用自己的数据集创建自定义分类器，该数据集包含原始 ImageNet 数据集中未包含的类别（预训练模型是在该数据集上训练的），该怎么办？

为此，您可以

从 TensorFlow Hub 选择一个预训练模型；以及
重新训练顶部（最后一个）层以识别来自自定义数据集的类别。

数据集

在本示例中，您将使用 TensorFlow 花卉数据集

import pathlib

data_file = tf.keras.utils.get_file(
  'flower_photos.tgz',
  'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
  cache_dir='.',
   extract=True)

data_root = pathlib.Path(data_file).with_suffix('')

首先，使用磁盘上的图像数据使用 tf.keras.utils.image_dataset_from_directory 将这些数据加载到模型中，这将生成一个 tf.data.Dataset

batch_size = 32
img_height = 224
img_width = 224

train_ds = tf.keras.utils.image_dataset_from_directory(
  str(data_root),
  validation_split=0.2,
  subset="training",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size
)

val_ds = tf.keras.utils.image_dataset_from_directory(
  str(data_root),
  validation_split=0.2,
  subset="validation",
  seed=123,
  image_size=(img_height, img_width),
  batch_size=batch_size
)

花卉数据集有五个类别

class_names = np.array(train_ds.class_names)
print(class_names)

其次，由于 TensorFlow Hub 对图像模型的约定是期望 [0, 1] 范围内的浮点输入，因此使用 tf.keras.layers.Rescaling 预处理层来实现这一点。

normalization_layer = tf.keras.layers.Rescaling(1./255)
train_ds = train_ds.map(lambda x, y: (normalization_layer(x), y)) # Where x—images, y—labels.
val_ds = val_ds.map(lambda x, y: (normalization_layer(x), y)) # Where x—images, y—labels.

第三，使用 Dataset.prefetch 使用缓冲预取来完成输入管道，这样您就可以从磁盘生成数据，而不会出现 I/O 阻塞问题。

这些是加载数据时应该使用的一些最重要的 tf.data 方法。有兴趣的读者可以了解更多关于它们的信息，以及如何将数据缓存到磁盘和其他技术，请参阅使用 tf.data API 提高性能指南。

AUTOTUNE = tf.data.AUTOTUNE
train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)

for image_batch, labels_batch in train_ds:
  print(image_batch.shape)
  print(labels_batch.shape)
  break

在图像批次上运行分类器

现在，在图像批次上运行分类器

result_batch = classifier.predict(train_ds)

predicted_class_names = imagenet_labels[tf.math.argmax(result_batch, axis=-1)]
predicted_class_names

检查这些预测与图像的匹配情况

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(predicted_class_names[n])
  plt.axis('off')
_ = plt.suptitle("ImageNet predictions")

结果远非完美，但考虑到这些不是模型训练的类别（除了“雏菊”），这是合理的。

下载无头模型

TensorFlow Hub 还分发没有顶部分类层的模型。这些模型可用于轻松执行迁移学习。

从 TensorFlow Hub 选择一个 MobileNetV2 预训练模型 from TensorFlow Hub。来自 TensorFlow Hub 的任何兼容图像特征向量模型都可以在此处使用，包括下拉菜单中的示例。

mobilenet_v2 = "https://tfhub.dev/google/tf2-preview/mobilenet_v2/feature_vector/4"
inception_v3 = "https://tfhub.dev/google/tf2-preview/inception_v3/feature_vector/4"

feature_extractor_model = mobilenet_v2

通过使用 hub.KerasLayer 将预训练模型包装为 Keras 层来创建特征提取器。使用 trainable=False 参数冻结变量，以便训练仅修改新的分类器层

feature_extractor_layer = hub.KerasLayer(
    feature_extractor_model,
    input_shape=(224, 224, 3),
    trainable=False)

特征提取器为每个图像返回一个 1280 长的向量（在本示例中，图像批次大小保持为 32）

feature_batch = feature_extractor_layer(image_batch)
print(feature_batch.shape)

附加分类头

要完成模型，请将特征提取器层包装在一个 tf.keras.Sequential 模型中，并添加一个用于分类的全连接层

num_classes = len(class_names)

model = tf.keras.Sequential([
  feature_extractor_layer,
  tf.keras.layers.Dense(num_classes)
])

model.summary()

predictions = model(image_batch)

predictions.shape

训练模型

使用 Model.compile 配置训练过程，并添加一个 tf.keras.callbacks.TensorBoard 回调以创建和存储日志

model.compile(
  optimizer=tf.keras.optimizers.Adam(),
  loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  metrics=['acc'])

log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(
    log_dir=log_dir,
    histogram_freq=1) # Enable histogram computation for every epoch.

现在使用 Model.fit 方法训练模型。

为了使本示例简短，您将只训练 10 个 epoch。为了稍后在 TensorBoard 中可视化训练进度，请创建一个 TensorBoard 回调并存储日志。

NUM_EPOCHS = 10

history = model.fit(train_ds,
                    validation_data=val_ds,
                    epochs=NUM_EPOCHS,
                    callbacks=tensorboard_callback)

启动 TensorBoard 以查看指标在每个 epoch 中的变化方式，以及跟踪其他标量值

%tensorboard --logdir logs/fit

检查预测

从模型预测中获取类名排序列表

predicted_batch = model.predict(image_batch)
predicted_id = tf.math.argmax(predicted_batch, axis=-1)
predicted_label_batch = class_names[predicted_id]
print(predicted_label_batch)

绘制模型预测

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)

for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(predicted_label_batch[n].title())
  plt.axis('off')
_ = plt.suptitle("Model predictions")

导出和重新加载您的模型

现在您已经训练了模型，请将其导出为 SavedModel 以便稍后重新使用。

t = time.time()

export_path = "/tmp/saved_models/{}".format(int(t))
model.save(export_path)

export_path

确认您可以重新加载 SavedModel 并且模型能够输出相同的结果

reloaded = tf.keras.models.load_model(export_path)

result_batch = model.predict(image_batch)
reloaded_result_batch = reloaded.predict(image_batch)

abs(reloaded_result_batch - result_batch).max()

reloaded_predicted_id = tf.math.argmax(reloaded_result_batch, axis=-1)
reloaded_predicted_label_batch = class_names[reloaded_predicted_id]
print(reloaded_predicted_label_batch)

plt.figure(figsize=(10,9))
plt.subplots_adjust(hspace=0.5)
for n in range(30):
  plt.subplot(6,5,n+1)
  plt.imshow(image_batch[n])
  plt.title(reloaded_predicted_label_batch[n].title())
  plt.axis('off')
_ = plt.suptitle("Model predictions")

后续步骤

您可以使用 SavedModel 加载进行推理，或者将其转换为 TensorFlow Lite 模型（用于设备上机器学习）或 TensorFlow.js 模型（用于 JavaScript 中的机器学习）。

探索更多教程以了解如何将来自 TensorFlow Hub 的预训练模型用于图像、文本、音频和视频任务。