重新訓練圖片分類器

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

簡介

圖片分類模型有數百萬個參數。從頭開始訓練這些模型需要大量的標籤訓練資料和大量的運算能力。遷移學習是一種技術,可藉由擷取已針對相關任務訓練的模型部分,並在新模型中重複使用,來大幅縮減所需資源。

這個 Colab 示範如何建構 Keras 模型,以使用來自 TensorFlow Hub 的預先訓練 TF2 SavedModel 進行圖片特徵擷取,藉此對五種花卉進行分類,而該模型已針對更大且更通用的 ImageNet 資料集進行訓練。或者,也可以與新加入的分類器一同訓練 (「微調」) 特徵擷取器。

改為尋找工具?

這是 TensorFlow 編碼教學課程。如果您想要一種工具,能直接為您建構 TensorFlow 或 TFLite 模型,請查看 make_image_classifier 命令列工具,此工具可透過 PIP 套件 tensorflow-hub[make_image_classifier] 安裝,或查看這個 TFLite colab。

設定

import itertools
import os

import matplotlib.pylab as plt
import numpy as np

import tensorflow as tf
import tensorflow_hub as hub

print("TF version:", tf.__version__)
print("Hub version:", hub.__version__)
print("GPU is", "available" if tf.config.list_physical_devices('GPU') else "NOT AVAILABLE")

選取要使用的 TF2 SavedModel 模組

首先,請使用 https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/feature_vector/4。程式碼中可以使用相同的網址來識別 SavedModel,瀏覽器中則可顯示其文件。(請注意,TF1 Hub 格式的模型在此處無法運作。)

您可以在這裡找到更多產生圖片特徵向量的 TF2 模型。

您可以嘗試多種可能的模型。您只需在下方的儲存格中選取不同的模型,然後繼續使用筆記本即可。

設定 Flowers 資料集

輸入會針對選取的模組適當調整大小。資料集擴增 (亦即,每次讀取圖片時隨機扭曲圖片) 可改善訓練效果,尤其是在微調時。

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

定義模型

只需將線性分類器放在具有 Hub 模組的 feature_extractor_layer 之上即可。

為了加快速度,我們從不可訓練的 feature_extractor_layer 開始,但您也可以啟用微調以提高準確性。

do_fine_tuning = False
print("Building model with", model_handle)
model = tf.keras.Sequential([
    # Explicitly define the input shape so the model can be properly
    # loaded by the TFLiteConverter
    tf.keras.layers.InputLayer(input_shape=IMAGE_SIZE + (3,)),
    hub.KerasLayer(model_handle, trainable=do_fine_tuning),
    tf.keras.layers.Dropout(rate=0.2),
    tf.keras.layers.Dense(len(class_names),
                          kernel_regularizer=tf.keras.regularizers.l2(0.0001))
])
model.build((None,)+IMAGE_SIZE+(3,))
model.summary()

訓練模型

model.compile(
  optimizer=tf.keras.optimizers.SGD(learning_rate=0.005, momentum=0.9), 
  loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1),
  metrics=['accuracy'])
steps_per_epoch = train_size // BATCH_SIZE
validation_steps = valid_size // BATCH_SIZE
hist = model.fit(
    train_ds,
    epochs=5, steps_per_epoch=steps_per_epoch,
    validation_data=val_ds,
    validation_steps=validation_steps).history
plt.figure()
plt.ylabel("Loss (training and validation)")
plt.xlabel("Training Steps")
plt.ylim([0,2])
plt.plot(hist["loss"])
plt.plot(hist["val_loss"])

plt.figure()
plt.ylabel("Accuracy (training and validation)")
plt.xlabel("Training Steps")
plt.ylim([0,1])
plt.plot(hist["accuracy"])
plt.plot(hist["val_accuracy"])

在驗證資料中的圖片上試用模型

x, y = next(iter(val_ds))
image = x[0, :, :, :]
true_index = np.argmax(y[0])
plt.imshow(image)
plt.axis('off')
plt.show()

# Expand the validation image to (1, 224, 224, 3) before predicting the label
prediction_scores = model.predict(np.expand_dims(image, axis=0))
predicted_index = np.argmax(prediction_scores)
print("True label: " + class_names[true_index])
print("Predicted label: " + class_names[predicted_index])

最後,可以儲存訓練後的模型,以便部署到 TF Serving 或 TFLite (在行動裝置上),如下所示。

saved_model_path = f"/tmp/saved_flowers_model_{model_name}"
tf.saved_model.save(model, saved_model_path)

選用:部署至 TensorFlow Lite

TensorFlow Lite 可讓您將 TensorFlow 模型部署到行動裝置和 IoT 裝置。以下程式碼說明如何將訓練後的模型轉換為 TFLite,並套用 TensorFlow Model Optimization Toolkit 的訓練後工具。最後,會在 TFLite Interpreter 中執行,以檢查產生的品質

  • 在不進行最佳化的情況下轉換,會產生與先前相同的結果 (最多到捨入誤差)。
  • 在不使用任何資料進行最佳化的情況下轉換,會將模型權重量化為 8 位元,但推論仍會對神經網路啟動使用浮點運算。這幾乎可將模型大小縮減 4 倍,並改善行動裝置上的 CPU 延遲時間。
  • 此外,如果提供小型參考資料集來校準量化範圍,則神經網路啟動的運算也可以量化為 8 位元整數。在行動裝置上,這會進一步加速推論,並可讓您在 Edge TPU 等加速器上執行。

最佳化設定

interpreter = tf.lite.Interpreter(model_content=lite_model_content)
# This little helper wraps the TFLite Interpreter as a numpy-to-numpy function.
def lite_model(images):
  interpreter.allocate_tensors()
  interpreter.set_tensor(interpreter.get_input_details()[0]['index'], images)
  interpreter.invoke()
  return interpreter.get_tensor(interpreter.get_output_details()[0]['index'])
num_eval_examples = 50 
eval_dataset = ((image, label)  # TFLite expects batch size 1.
                for batch in train_ds
                for (image, label) in zip(*batch))
count = 0
count_lite_tf_agree = 0
count_lite_correct = 0
for image, label in eval_dataset:
  probs_lite = lite_model(image[None, ...])[0]
  probs_tf = model(image[None, ...]).numpy()[0]
  y_lite = np.argmax(probs_lite)
  y_tf = np.argmax(probs_tf)
  y_true = np.argmax(label)
  count +=1
  if y_lite == y_tf: count_lite_tf_agree += 1
  if y_lite == y_true: count_lite_correct += 1
  if count >= num_eval_examples: break
print("TFLite model agrees with original model on %d of %d examples (%g%%)." %
      (count_lite_tf_agree, count, 100.0 * count_lite_tf_agree / count))
print("TFLite model is accurate on %d of %d examples (%g%%)." %
      (count_lite_correct, count, 100.0 * count_lite_correct / count))