使用 TensorFlow Lite Model Maker 進行圖片分類

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

TensorFlow Lite Model Maker 程式庫簡化了在為裝置端機器學習應用程式部署模型時,調整和轉換 TensorFlow 神經網路模型以符合特定輸入資料的流程。

這個筆記本示範了端對端範例,說明如何運用 Model Maker 程式庫來調整和轉換常用的圖片分類模型,以便在行動裝置上進行花卉分類。

先決條件

若要執行這個範例,我們首先需要安裝幾個必要的套件,包括 GitHub 存放區中的 Model Maker 套件。

sudo apt -y install libportaudio2
pip install -q tflite-model-maker

匯入必要的套件。

import os

import numpy as np

import tensorflow as tf
assert tf.__version__.startswith('2')

from tflite_model_maker import model_spec
from tflite_model_maker import image_classifier
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.config import QuantizationConfig
from tflite_model_maker.image_classifier import DataLoader

import matplotlib.pyplot as plt

簡易端對端範例

取得資料路徑

讓我們取得一些圖片來試用這個簡易端對端範例。數百張圖片是 Model Maker 的良好起點,而更多資料則可實現更佳的準確度。

您可以將 image_path 替換成您自己的圖片資料夾。如要將資料上傳至 Colab,您可以在下圖紅色矩形標示的左側邊欄中找到上傳按鈕。不妨試著上傳 zip 檔並解壓縮。根檔案路徑為目前路徑。

Upload File

如果您不想將圖片上傳到雲端,可以嘗試按照 GitHub 中的指南在本機執行程式庫。

執行範例

此範例僅由 4 行程式碼組成,如下所示,每一行程式碼代表整體流程的一個步驟。

步驟 1. 載入裝置端機器學習應用程式專用的輸入資料。將其分割為訓練資料和測試資料。

data = DataLoader.from_folder(image_path)
train_data, test_data = data.split(0.9)

步驟 2. 自訂 TensorFlow 模型。

model = image_classifier.create(train_data)

步驟 3. 評估模型。

loss, accuracy = model.evaluate(test_data)

步驟 4. 匯出為 TensorFlow Lite 模型。

在此,我們匯出具有中繼資料的 TensorFlow Lite 模型,中繼資料為模型描述提供標準。標籤檔案已嵌入在中繼資料中。預設的訓練後量化技術是圖片分類任務的全整數量化。

您可以在左側邊欄中下載,方法與您自己的用途的上傳部分相同。

model.export(export_dir='.')

完成這 4 個簡單步驟後,我們可以進一步在裝置端應用程式中使用 TensorFlow Lite 模型檔案,例如在圖片分類參考應用程式中。

詳細流程

目前,我們支援多種模型,例如 EfficientNet-Lite* 模型、MobileNetV2、ResNet50 作為圖片分類的預先訓練模型。但是,只需幾行程式碼,即可彈性地將新的預先訓練模型新增至此程式庫。

以下逐步說明這個端對端範例,以顯示更多詳細資訊。

步驟 1:載入裝置端機器學習應用程式專用的輸入資料

花卉資料集包含 3670 張圖片,屬於 5 個類別。下載資料集的封存版本並解壓縮。

資料集具有下列目錄結構

flower_photos
|__ daisy
    |______ 100080576_f52e8ee070_n.jpg
    |______ 14167534527_781ceb1b7a_n.jpg
    |______ ...
|__ dandelion
    |______ 10043234166_e6dd915111_n.jpg
    |______ 1426682852_e62169221f_m.jpg
    |______ ...
|__ roses
    |______ 102501987_3cdb8e5394_n.jpg
    |______ 14982802401_a3dfb22afb.jpg
    |______ ...
|__ sunflowers
    |______ 12471791574_bb1be83df4.jpg
    |______ 15122112402_cafa41934f.jpg
    |______ ...
|__ tulips
    |______ 13976522214_ccec508fe7.jpg
    |______ 14487943607_651e8062a1_m.jpg
    |______ ...
image_path = tf.keras.utils.get_file(
      'flower_photos.tgz',
      'https://storage.googleapis.com/download.tensorflow.org/example_images/flower_photos.tgz',
      extract=True)
image_path = os.path.join(os.path.dirname(image_path), 'flower_photos')

使用 DataLoader 類別載入資料。

至於 from_folder() 方法,它可以從資料夾載入資料。它假設相同類別的圖片資料位於相同的子目錄中,而子資料夾名稱是類別名稱。目前支援 JPEG 編碼圖片和 PNG 編碼圖片。

data = DataLoader.from_folder(image_path)

將其分割為訓練資料 (80%)、驗證資料 (10%,選用) 和測試資料 (10%)。

train_data, rest_data = data.split(0.8)
validation_data, test_data = rest_data.split(0.5)

顯示 25 個帶有標籤的圖片範例。

plt.figure(figsize=(10,10))
for i, (image, label) in enumerate(data.gen_dataset().unbatch().take(25)):
  plt.subplot(5,5,i+1)
  plt.xticks([])
  plt.yticks([])
  plt.grid(False)
  plt.imshow(image.numpy(), cmap=plt.cm.gray)
  plt.xlabel(data.index_to_label[label.numpy()])
plt.show()

步驟 2:自訂 TensorFlow 模型

根據載入的資料建立自訂圖片分類器模型。預設模型為 EfficientNet-Lite0。

model = image_classifier.create(train_data, validation_data=validation_data)

查看詳細模型結構。

model.summary()

步驟 3:評估自訂模型

評估模型結果,取得模型的損失和準確度。

loss, accuracy = model.evaluate(test_data)

我們可以在 100 張測試圖片中繪製預測結果。紅色預測標籤為錯誤預測結果,其他為正確。

# A helper function that returns 'red'/'black' depending on if its two input
# parameter matches or not.
def get_label_color(val1, val2):
  if val1 == val2:
    return 'black'
  else:
    return 'red'

# Then plot 100 test images and their predicted labels.
# If a prediction result is different from the label provided label in "test"
# dataset, we will highlight it in red color.
plt.figure(figsize=(20, 20))
predicts = model.predict_top_k(test_data)
for i, (image, label) in enumerate(test_data.gen_dataset().unbatch().take(100)):
  ax = plt.subplot(10, 10, i+1)
  plt.xticks([])
  plt.yticks([])
  plt.grid(False)
  plt.imshow(image.numpy(), cmap=plt.cm.gray)

  predict_label = predicts[i][0][0]
  color = get_label_color(predict_label,
                          test_data.index_to_label[label.numpy()])
  ax.xaxis.label.set_color(color)
  plt.xlabel('Predicted: %s' % predict_label)
plt.show()

如果準確度不符合應用程式需求,可以參考進階用法,探索替代方案,例如變更為較大型模型、調整重新訓練參數等。

步驟 4:匯出為 TensorFlow Lite 模型

將訓練後的模型轉換為具有中繼資料的 TensorFlow Lite 模型格式,以便您稍後可以在裝置端機器學習應用程式中使用。標籤檔案和詞彙檔案已嵌入在中繼資料中。預設的 TFLite 檔案名稱為 model.tflite

在許多裝置端機器學習應用程式中,模型大小是一個重要因素。因此,建議您量化模型,使其更小且可能執行更快。圖像分類任務的預設訓練後量化技術為全整數量化。

model.export(export_dir='.')

請參閱圖片分類範例指南,以取得關於如何將 TensorFlow Lite 模型整合到行動應用程式的更多詳細資訊。

此模型可以使用 ImageClassifier API,整合到 Android 或 iOS 應用程式中,TensorFlow Lite Task Library

允許的匯出格式可以是以下其中一種或多種

預設情況下,它只匯出帶有中繼資料的 TensorFlow Lite 模型。您也可以選擇性地匯出不同的檔案。例如,僅匯出標籤檔案,如下所示

model.export(export_dir='.', export_format=ExportFormat.LABEL)

您也可以使用 evaluate_tflite 方法評估 tflite 模型。

model.evaluate_tflite('model.tflite', test_data)

進階用法

create 函數是此程式庫的關鍵部分。它使用遷移學習,並使用類似於教學課程的預先訓練模型。

create 函數包含以下步驟

  1. 根據參數 validation_ratiotest_ratio 將資料分割為訓練、驗證、測試資料。validation_ratiotest_ratio 的預設值分別為 0.10.1
  2. 從 TensorFlow Hub 下載 Image Feature Vector 作為基礎模型。預設的預先訓練模型為 EfficientNet-Lite0。
  3. 在標頭層和預先訓練模型之間新增一個帶有 dropout_rate 的 Dropout Layer 的分類器標頭。預設 dropout_rate 是 TensorFlow Hub 的 make_image_classifier_lib 中的預設 dropout_rate 值。
  4. 預處理原始輸入資料。目前的預處理步驟包括將每個圖片像素的值正規化為模型輸入比例,並將其調整為模型輸入大小。EfficientNet-Lite0 的輸入比例為 [0, 1],輸入圖片大小為 [224, 224, 3]
  5. 將資料饋送到分類器模型中。預設情況下,訓練參數 (例如訓練週期、批次大小、學習率、動量) 是 TensorFlow Hub 的 make_image_classifier_lib 中的預設值。僅訓練分類器標頭。

在本節中,我們描述了幾個進階主題,包括切換到不同的圖片分類模型、變更訓練超參數等。

自訂 TensorFLow Lite 模型上的訓練後量化

訓練後量化是一種轉換技術,可以減少模型大小和推論延遲,同時還可以提高 CPU 和硬體加速器推論速度,而模型準確度只會稍微降低。因此,它被廣泛用於最佳化模型。

Model Maker 程式庫在匯出模型時套用預設的訓練後量化技術。如果您想要自訂訓練後量化,Model Maker 也支援使用 QuantizationConfig 多個訓練後量化選項。讓我們以 float16 量化為例。首先,定義量化設定。

config = QuantizationConfig.for_float16()

然後我們使用此設定匯出 TensorFlow Lite 模型。

model.export(export_dir='.', tflite_filename='model_fp16.tflite', quantization_config=config)

在 Colab 中,您可以從左側邊欄下載名為 model_fp16.tflite 的模型,方法與上述的上傳部分相同。

變更模型

變更為此程式庫中支援的模型。

此程式庫目前支援 EfficientNet-Lite 模型、MobileNetV2、ResNet50。EfficientNet-Lite 是一系列圖片分類模型,可以達到最先進的準確度,並且適用於邊緣裝置。預設模型為 EfficientNet-Lite0。

我們可以透過在 create 方法中將參數 model_spec 設定為 MobileNetV2 模型規格,將模型切換為 MobileNetV2。

model = image_classifier.create(train_data, model_spec=model_spec.get('mobilenet_v2'), validation_data=validation_data)

評估新重新訓練的 MobileNetV2 模型,以查看測試資料中的準確度和損失。

loss, accuracy = model.evaluate(test_data)

變更為 TensorFlow Hub 中的模型

此外,我們也可以切換到其他新模型,這些模型輸入圖片並輸出具有 TensorFlow Hub 格式的特徵向量。

Inception V3 模型為例,我們可以定義 inception_v3_spec,它是 image_classifier.ModelSpec 的一個物件,並且包含 Inception V3 模型的規格。

我們需要指定模型名稱 name、TensorFlow Hub 模型的網址 uri。同時,input_image_shape 的預設值為 [224, 224]。我們需要將其變更為 [299, 299] 以用於 Inception V3 模型。

inception_v3_spec = image_classifier.ModelSpec(
    uri='https://tfhub.dev/google/imagenet/inception_v3/feature_vector/1')
inception_v3_spec.input_image_shape = [299, 299]

然後,透過在 create 方法中將參數 model_spec 設定為 inception_v3_spec,我們可以重新訓練 Inception V3 模型。

剩餘步驟完全相同,我們最終可以獲得自訂的 InceptionV3 TensorFlow Lite 模型。

變更您自己的自訂模型

如果我們想要使用不在 TensorFlow Hub 中的自訂模型,我們應該在 TensorFlow Hub 中建立並匯出 ModelSpec

然後開始像上述流程一樣定義 ModelSpec 物件。

變更訓練超參數

我們也可以變更訓練超參數,例如 epochsdropout_ratebatch_size,這可能會影響模型準確度。您可以調整的模型參數包括:

  • epochs:更多週期可以達到更好的準確度,直到收斂,但訓練過多週期可能會導致過度擬合。
  • dropout_rate:Dropout 率,避免過度擬合。預設為 None。
  • batch_size:一個訓練步驟中要使用的樣本數。預設為 None。
  • validation_data:驗證資料。如果為 None,則跳過驗證過程。預設為 None。
  • train_whole_model:如果為 true,則 Hub 模組會與頂部的分類層一起訓練。否則,僅訓練頂部分類層。預設為 None。
  • learning_rate:基礎學習率。預設為 None。
  • momentum:轉發給最佳化器的 Python float。僅當 use_hub_library 為 True 時使用。預設為 None。
  • shuffle:布林值,資料是否應該洗牌。預設為 False。
  • use_augmentation:布林值,使用資料擴增進行預處理。預設為 False。
  • use_hub_library:布林值,使用 tensorflow hub 中的 make_image_classifier_lib 重新訓練模型。此訓練管線可以針對具有許多類別的複雜資料集實現更好的效能。預設為 True。
  • warmup_steps:學習率預熱排程的預熱步驟數。如果為 None,則使用預設 warmup_steps,即兩個週期的總訓練步驟數。僅當 use_hub_library 為 False 時使用。預設為 None。
  • model_dir:選用,模型檢查點檔案的位置。僅當 use_hub_library 為 False 時使用。預設為 None。

預設值為 None 的參數 (例如 epochs) 將在 TensorFlow Hub 程式庫的 make_image_classifier_libtrain_image_classifier_lib 中取得具體的預設參數。

例如,我們可以訓練更多週期。

model = image_classifier.create(train_data, validation_data=validation_data, epochs=10)

評估新重新訓練的具有 10 個訓練週期的模型。

loss, accuracy = model.evaluate(test_data)

閱讀更多資訊

您可以閱讀我們的圖片分類範例以學習技術細節。如需更多資訊,請參閱