在 TensorFlow.org 上檢視

在 Google Colab 中執行

在 GitHub 上檢視原始碼

下載筆記本

TensorFlow Lite Model Maker 程式庫簡化了在部署此模型以用於裝置端 ML 應用程式時，調整和轉換 TensorFlow 模型以符合特定輸入資料的流程。

此筆記本顯示端對端範例，說明如何運用 Model Maker 程式庫調整和轉換常用的文字分類模型，以便在行動裝置上分類電影評論。文字分類模型會將文字分類到預先定義的類別。輸入應為預先處理的文字，輸出為類別的機率。本教學課程中使用的資料集是正面和負面電影評論。

先決條件

安裝必要套件

若要執行此範例，請安裝必要套件，包括 GitHub 存放區中的 Model Maker 套件。

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

匯入必要套件。

import numpy as np
import os

from tflite_model_maker import model_spec
from tflite_model_maker import text_classifier
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.text_classifier import AverageWordVecSpec
from tflite_model_maker.text_classifier import DataLoader

from tflite_support.task import core
from tflite_support.task import processor
from tflite_support.task import text

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

下載範例訓練資料。

在本教學課程中，我們將使用 SST-2 (史丹佛情感樹庫)，這是 GLUE 基準測試中的其中一項任務。其中包含 67,349 則電影評論用於訓練，以及 872 則電影評論用於測試。資料集有兩個類別：正面和負面電影評論。

data_dir = tf.keras.utils.get_file(
      fname='SST-2.zip',
      origin='https://dl.fbaipublicfiles.com/glue/data/SST-2.zip',
      extract=True)
data_dir = os.path.join(os.path.dirname(data_dir), 'SST-2')

SST-2 資料集以 TSV 格式儲存。TSV 和 CSV 之間的唯一差異在於，TSV 使用 Tab 字元 \t 作為分隔符號，而非 CSV 格式中的逗號 ,。

以下是訓練資料集的前 5 行。label=0 代表負面，label=1 代表正面。

sentence	label
hide new secretions from the parental units	0
contains no wit , only labored gags	0
that loves its characters and communicates something rather beautiful about human nature	1
remains utterly satisfied to remain the same throughout	0
on the worst revenge-of-the-nerds clichés the filmmakers could dredge up	0

接下來，我們會將資料集載入 Pandas DataFrame，並將目前的標籤名稱 (0 和 1) 變更為更容易理解的名稱 (negative 和 positive)，並用於模型訓練。

import pandas as pd

def replace_label(original_file, new_file):
  # Load the original file to pandas. We need to specify the separator as
  # '\t' as the training data is stored in TSV format
  df = pd.read_csv(original_file, sep='\t')

  # Define how we want to change the label name
  label_map = {0: 'negative', 1: 'positive'}

  # Excute the label change
  df.replace({'label': label_map}, inplace=True)

  # Write the updated dataset to a new file
  df.to_csv(new_file)

# Replace the label name for both the training and test dataset. Then write the
# updated CSV dataset to the current folder.
replace_label(os.path.join(os.path.join(data_dir, 'train.tsv')), 'train.csv')
replace_label(os.path.join(os.path.join(data_dir, 'dev.tsv')), 'dev.csv')

快速入門

訓練文字分類模型共有五個步驟

步驟 1. 選擇文字分類模型架構。

在這裡，我們使用平均字詞嵌入模型架構，這會產生一個小型快速且準確度尚可的模型。

spec = model_spec.get('average_word_vec')

Model Maker 也支援其他模型架構，例如 BERT。如果您有興趣瞭解其他架構，請參閱下方的「選擇文字分類器的模型架構」一節。

步驟 2. 載入訓練和測試資料，然後根據特定的 model_spec 預先處理這些資料。

Model Maker 可以 CSV 格式取得輸入資料。我們會載入先前建立的訓練和測試資料集，其中包含容易理解的標籤名稱。

每個模型架構都要求以特定方式處理輸入資料。DataLoader 會從 model_spec 讀取需求，並自動執行必要的預先處理。

train_data = DataLoader.from_csv(
      filename='train.csv',
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      is_training=True)
test_data = DataLoader.from_csv(
      filename='dev.csv',
      text_column='sentence',
      label_column='label',
      model_spec=spec,
      is_training=False)

步驟 3. 使用訓練資料訓練 TensorFlow 模型。

平均字詞嵌入模型預設使用 batch_size = 32。因此，您會看到需要 2104 個步驟才能完成訓練資料集中 67,349 個句子的處理。我們會將模型訓練 10 個週期，也就是處理訓練資料集 10 次。

model = text_classifier.create(train_data, model_spec=spec, epochs=10)

步驟 4. 使用測試資料評估模型。

在使用訓練資料集中的句子訓練文字分類模型後，我們會使用測試資料集中剩餘的 872 個句子，評估模型在新資料 (模型從未見過的資料) 上的效能。

由於預設批次大小為 32，因此需要 28 個步驟才能完成測試資料集中 872 個句子的處理。

loss, acc = model.evaluate(test_data)

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

讓我們匯出以 TensorFlow Lite 格式訓練的文字分類模型。我們會指定匯出模型的資料夾。根據預設，系統會針對平均字詞嵌入模型架構匯出 float TFLite 模型。

model.export(export_dir='average_word_vec')

您可以使用 Colab 左側邊欄下載 TensorFlow Lite 模型檔案。前往 export_dir 參數中指定的 average_word_vec 資料夾，在 model.tflite 檔案上按一下滑鼠右鍵，然後選擇「Download」即可將檔案下載到本機電腦。

這個模型可以使用 NLClassifier API 的 TensorFlow Lite Task Library 整合到 Android 或 iOS 應用程式中。

如需如何在實際運作的應用程式中使用模型的詳細資訊，請參閱 TFLite 文字分類範例應用程式。

注意 1：Android Studio Model Binding 尚不支援文字分類，因此請使用 TensorFlow Lite Task Library。

注意 2：在與 TFLite 模型相同的資料夾中，有一個 model.json 檔案。其中包含 TensorFlow Lite 模型中捆綁的 中繼資料 的 JSON 表示法。模型中繼資料可協助 TFLite Task Library 瞭解模型的功能，以及如何預先處理/後續處理模型的資料。您不需要下載 model.json 檔案，因為它僅供參考，且其內容已包含在 TFLite 檔案中。

注意 3：如果您使用 MobileBERT 或 BERT-Base 架構訓練文字分類模型，則需要改用 BertNLClassifier API，才能將訓練後的模型整合到行動應用程式中。

以下章節將逐步引導您完成範例，以顯示更多詳細資訊。

步驟 6：使用 TFLite Task Library 示範如何使用訓練後的模型

將 dev.csv 檔案讀取到句子資料中，以便使用訓練後的模型進行預測

sentence_data = pd.read_csv('/content/dev.csv', index_col=0)
sentence_data

模型設定參數

# Name of the TFLite text classification model.
_MODEL = '/content/average_word_vec/model.tflite'
# Whether to run the model on EdgeTPU.
_ENABLE_EDGETPU = False
# Number of CPU threads to run the model.
_NUM_THREADS = 4

初始化模型

我們也可以變更 file_name、use_coral 和 num_threads 等參數，這些參數可能會影響模型結果。您可以調整的參數如下：

• file_name：TFLite 圖片分類模型的名稱。
• use_coral：若為 true，推論將委派給已連線的 Coral Edge TPU 裝置。
• num_threads：執行模型的 CPU 執行緒數。

# Initialize the text classification model.
base_options = core.BaseOptions(file_name=_MODEL, use_coral=_ENABLE_EDGETPU, num_threads=_NUM_THREADS)
options = text.NLClassifierOptions(base_options)

# Create NLClassifier from options.
classifier = text.NLClassifier.create_from_options(options)

使用 TFLite Task Library 進行預測

for idx in range(20):
  sentence = sentence_data['sentence'].iloc[idx]
  label = sentence_data['label'].iloc[idx]
  text_classification_result = classifier.classify(sentence)
  classification_list = text_classification_result.classifications[0].categories

  # Sort output by probability descending.
  predict_label = sorted(
      classification_list, key=lambda item: item.score, reverse=True)[0]

  print('truth_label: {} -----> predict_label: {}'.format(label, predict_label.category_name))

選擇文字分類器的模型架構

每個 model_spec 物件都代表文字分類器的特定模型。TensorFlow Lite Model Maker 目前支援 MobileBERT、平均字詞嵌入和 BERT-Base 模型。

支援的模型	model_spec 名稱	模型說明	模型大小
平均字詞嵌入	'average_word_vec'	平均文字字詞嵌入，搭配 RELU 啟動。	<1MB
MobileBERT	'mobilebert_classifier'	比 BERT-Base 小 4.3 倍且快 5.5 倍，同時達到具有競爭力的成果，適合裝置端應用程式。	25MB (含量化) 100MB (不含量化)
BERT-Base	'bert_classifier'	廣泛用於 NLP 任務的標準 BERT 模型。	300MB

在快速入門中，我們使用了平均字詞嵌入模型。讓我們切換到 MobileBERT 來訓練具有更高準確度的模型。

mb_spec = model_spec.get('mobilebert_classifier')

載入訓練資料

您可以上傳自己的資料集，以完成本教學課程。使用 Colab 左側邊欄上傳您的資料集。

如果您不想將資料集上傳到雲端，也可以按照指南在本機執行程式庫。

為了簡化操作，我們將重複使用先前下載的 SST-2 資料集。讓我們使用 DataLoader.from_csv 方法載入資料。

請注意，由於我們變更了模型架構，因此需要重新載入訓練和測試資料集，以套用新的預先處理邏輯。

train_data = DataLoader.from_csv(
      filename='train.csv',
      text_column='sentence',
      label_column='label',
      model_spec=mb_spec,
      is_training=True)
test_data = DataLoader.from_csv(
      filename='dev.csv',
      text_column='sentence',
      label_column='label',
      model_spec=mb_spec,
      is_training=False)

Model Maker 程式庫也支援 from_folder() 方法來載入資料。它會假設相同類別的文字資料位於相同的子目錄中，且子資料夾名稱為類別名稱。每個文字檔包含一個電影評論範例。class_labels 參數用於指定子資料夾。

訓練 TensorFlow 模型

使用訓練資料訓練文字分類模型。

model = text_classifier.create(train_data, model_spec=mb_spec, epochs=3)

檢查詳細的模型結構。

model.summary()

評估模型

評估我們剛使用測試資料訓練的模型，並測量損失和準確度值。

loss, acc = model.evaluate(test_data)

匯出為 TensorFlow Lite 模型

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

在許多裝置端 ML 應用程式中，模型大小是重要的考量因素。因此，建議您量化模型以使其更小，並可能更快執行。BERT 和 MobileBERT 模型的預設訓練後量化技術是動態範圍量化。

model.export(export_dir='mobilebert/')

TensorFlow Lite 模型檔案可以使用 BertNLClassifier API (在 TensorFlow Lite Task Library 中) 整合到行動應用程式中。請注意，這與用於整合以平均字詞向量模型架構訓練的文字分類模型的 NLClassifier API 不同。

匯出格式可以是下列其中一種或清單

• ExportFormat.TFLITE
• ExportFormat.LABEL
• ExportFormat.VOCAB
• ExportFormat.SAVED_MODEL

根據預設，它只會匯出包含模型中繼資料的 TensorFlow Lite 模型檔案。您也可以選擇匯出與模型相關的其他檔案，以便更妥善地檢查。例如，僅匯出標籤檔案和詞彙檔案，如下所示：

model.export(export_dir='mobilebert/', export_format=[ExportFormat.LABEL, ExportFormat.VOCAB])

您可以使用 evaluate_tflite 方法評估 TFLite 模型，以測量其準確度。將訓練後的 TensorFlow 模型轉換為 TFLite 格式並套用量化可能會影響其準確度，因此建議您在部署前評估 TFLite 模型的準確度。

accuracy = model.evaluate_tflite('mobilebert/model.tflite', test_data)
print('TFLite model accuracy: ', accuracy)

進階用法

create 函式是 Model Maker 程式庫用來建立模型的驅動程式函式。model_spec 參數定義模型規格。目前支援 AverageWordVecSpec 和 BertClassifierSpec 類別。create 函式包含下列步驟：

1. 根據 model_spec 建立文字分類器的模型。
2. 訓練分類器模型。預設週期和預設批次大小是由 model_spec 物件中的 default_training_epochs 和 default_batch_size 變數設定。

本節涵蓋進階用法主題，例如調整模型和訓練超參數。

自訂 MobileBERT 模型超參數

您可以調整的模型參數如下：

• seq_len：饋送至模型的序列長度。
• initializer_range：用於初始化所有權重矩陣的 truncated_normal_initializer 的標準差。
• trainable：布林值，指定預先訓練的層是否可訓練。

您可以調整的訓練管道參數如下：

• model_dir：模型檢查點檔案的位置。如果未設定，將使用暫時目錄。
• dropout_rate：Dropout 率。
• learning_rate：Adam 最佳化工具的初始學習率。
• tpu：要連線的 TPU 位址。

例如，您可以設定 seq_len=256 (預設值為 128)。這可讓模型分類較長的文字。

new_model_spec = model_spec.get('mobilebert_classifier')
new_model_spec.seq_len = 256

自訂平均字詞嵌入模型超參數

您可以調整模型基礎架構，例如 AverageWordVecSpec 類別中的 wordvec_dim 和 seq_len 變數。

例如，您可以使用較大的 wordvec_dim 值訓練模型。請注意，如果您修改模型，則必須建構新的 model_spec。

new_model_spec = AverageWordVecSpec(wordvec_dim=32)

取得預先處理的資料。

new_train_data = DataLoader.from_csv(
      filename='train.csv',
      text_column='sentence',
      label_column='label',
      model_spec=new_model_spec,
      is_training=True)

訓練新模型。

model = text_classifier.create(new_train_data, model_spec=new_model_spec)

調整訓練超參數

您也可以調整訓練超參數，例如影響模型準確度的 epochs 和 batch_size。例如：

• epochs：更多週期可以達到更好的準確度，但也可能導致過度配適。
• batch_size：在一個訓練步驟中使用的樣本數。

例如，您可以使用更多週期進行訓練。

model = text_classifier.create(new_train_data, model_spec=new_model_spec, epochs=20)

使用 20 個訓練週期評估新重新訓練的模型。

new_test_data = DataLoader.from_csv(
      filename='dev.csv',
      text_column='sentence',
      label_column='label',
      model_spec=new_model_spec,
      is_training=False)

loss, accuracy = model.evaluate(new_test_data)

變更模型架構

您可以透過變更 model_spec 來變更模型。以下說明如何變更為 BERT-Base 模型。

將文字分類器的 model_spec 變更為 BERT-Base 模型。

spec = model_spec.get('bert_classifier')

其餘步驟相同。

自訂 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)

閱讀更多資訊

您可以閱讀我們的 文字分類範例，以瞭解技術詳細資訊。如需更多資訊，請參閱：

• TensorFlow Lite Model Maker 指南和 API 參考資料。
• Task Library：用於部署的 NLClassifier 和 BertNLClassifier。
• 端對端參考應用程式：Android 和 iOS。