使用 TensorFlow Hub 進行遷移學習

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視 下載筆記本 查看 TF Hub 模型

TensorFlow Hub 是預先訓練 TensorFlow 模型的存放區。

本教學課程示範如何

  1. 搭配 tf.keras 使用 TensorFlow Hub 的模型。
  2. 使用 TensorFlow Hub 的圖片分類模型。
  3. 進行簡單的遷移學習，以便針對您自己的圖片類別微調模型。

設定

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 預先訓練模型，並使用 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 資料集中 (預先訓練模型在該資料集上接受訓練)，該怎麼辦？

若要執行此操作，您可以：

  1. 從 TensorFlow Hub 選取預先訓練模型；以及
  2. 重新訓練頂層 (最後一層) 以辨識來自您自訂資料集的類別。

資料集

在此範例中，您將使用 TensorFlow Flowers 資料集

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
)

Flowers 資料集有五個類別

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")

結果遠非完美，但考慮到這些並非模型訓練的類別 (「daisy」除外)，因此還算合理。

下載無頭模型

TensorFlow Hub 也會散佈不含頂層分類層的模型。這些模型可用於輕鬆執行遷移學習。

從 TensorFlow Hub 中選取 MobileNetV2 預先訓練模型。來自 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 的預先訓練模型。