BERT 問答與 TensorFlow Lite Model Maker

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

TensorFlow Lite Model Maker 程式庫簡化了當部署裝置端機器學習應用程式時，調整 TensorFlow 模型並將其轉換為特定輸入資料的流程。

這個筆記本示範了一個端對端範例，說明如何運用 Model Maker 程式庫調整並轉換常用的問答模型，以執行問答任務。

BERT 問答任務簡介

這個程式庫中支援的任務是抽取式問答任務，也就是說，在給定段落和問題的情況下，答案是段落中的一段文字。下圖顯示了一個問答範例。

答案是段落中的文字範圍 (圖片來源：SQuAD 部落格)

對於問答任務模型，輸入應為已預先處理的段落和問題配對，輸出應為段落中每個符記的起始和結束 logits。輸入大小可以根據段落和問題的長度進行設定和調整。

端對端總覽

以下程式碼片段示範如何在幾行程式碼內取得模型。整體流程包含 5 個步驟：(1) 選擇模型，(2) 載入資料，(3) 重新訓練模型，(4) 評估，以及 (5) 匯出為 TensorFlow Lite 格式。

# Chooses a model specification that represents the model.
spec = model_spec.get('mobilebert_qa')

# Gets the training data and validation data.
train_data = DataLoader.from_squad(train_data_path, spec, is_training=True)
validation_data = DataLoader.from_squad(validation_data_path, spec, is_training=False)

# Fine-tunes the model.
model = question_answer.create(train_data, model_spec=spec)

# Gets the evaluation result.
metric = model.evaluate(validation_data)

# Exports the model to the TensorFlow Lite format with metadata in the export directory.
model.export(export_dir)

以下章節將更詳細地說明程式碼。

先決條件

若要執行這個範例，請安裝必要的套件，包括來自 GitHub 存放區的 Model Maker 套件。

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

匯入必要的套件。

import numpy as np
import os

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

from tflite_model_maker import model_spec
from tflite_model_maker import question_answer
from tflite_model_maker.config import ExportFormat
from tflite_model_maker.question_answer import DataLoader

「端對端總覽」示範了一個簡單的端對端範例。以下章節將逐步說明範例，以呈現更多細節。

選擇代表問答模型的 model_spec

每個 model_spec 物件都代表一個特定的問答模型。Model Maker 目前支援 MobileBERT 和 BERT-Base 模型。

支援的模型 model_spec 名稱 模型說明
MobileBERT 'mobilebert_qa' 比 BERT-Base 小 4.3 倍、快 5.5 倍，同時達到具競爭力的成果，適合裝置端情境。
MobileBERT-SQuAD 'mobilebert_qa_squad' 與 MobileBERT 模型相同的模型架構，且初始模型已在 SQuAD1.1 上重新訓練。
BERT-Base 'bert_qa' 廣泛用於 NLP 任務的標準 BERT 模型。

在本教學課程中，MobileBERT-SQuAD 作為範例使用。由於模型已在 SQuAD1.1 上重新訓練，因此可以更快地涵蓋問答任務。

spec = model_spec.get('mobilebert_qa_squad')

載入裝置端機器學習應用程式專屬的輸入資料並預先處理資料

TriviaQA 是一個閱讀理解資料集，包含超過 65 萬個問答證據三元組。在本教學課程中，您將使用這個資料集的子集來學習如何使用 Model Maker 程式庫。

若要載入資料，請執行 converter Python 腳本，並使用 --sample_size=8000 和一組 web 資料，將 TriviaQA 資料集轉換為 SQuAD1.1 格式。稍微修改轉換程式碼：

  • 略過在上下文文件中找不到任何答案的範例；
  • 取得上下文中不含大小寫的原始答案。

下載已轉換資料集的封存版本。

train_data_path = tf.keras.utils.get_file(
    fname='triviaqa-web-train-8000.json',
    origin='https://storage.googleapis.com/download.tensorflow.org/models/tflite/dataset/triviaqa-web-train-8000.json')
validation_data_path = tf.keras.utils.get_file(
    fname='triviaqa-verified-web-dev.json',
    origin='https://storage.googleapis.com/download.tensorflow.org/models/tflite/dataset/triviaqa-verified-web-dev.json')

您也可以使用自己的資料集訓練 MobileBERT 模型。如果您在 Colab 上執行這個筆記本，請使用左側邊欄上傳您的資料。

Upload File

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

使用 DataLoader.from_squad 方法，根據特定的 model_spec 載入和預先處理 SQuAD 格式資料。您可以使用 SQuAD2.0 或 SQuAD1.1 格式。將參數 version_2_with_negative 設定為 True 表示格式為 SQuAD2.0。否則，格式為 SQuAD1.1。預設情況下，version_2_with_negative 為 False。

train_data = DataLoader.from_squad(train_data_path, spec, is_training=True)
validation_data = DataLoader.from_squad(validation_data_path, spec, is_training=False)

自訂 TensorFlow 模型

根據載入的資料建立自訂問答模型。create 函式包含以下步驟

  1. 根據 model_spec 建立問答模型。
  2. 訓練問答模型。預設 epoch 和預設批次大小根據 model_spec 物件中的兩個變數 default_training_epochs 和 default_batch_size 設定。
model = question_answer.create(train_data, model_spec=spec)

查看詳細的模型結構。

model.summary()

評估自訂模型

在驗證資料上評估模型，並取得包含 f1 分數和 exact match 等指標的字典。請注意，SQuAD1.1 和 SQuAD2.0 的指標有所不同。

model.evaluate(validation_data)

匯出為 TensorFlow Lite 模型

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

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

model.export(export_dir='.')

您可以透過從 Colab 左側邊欄下載 TensorFlow Lite 模型檔案，在 bert_qa 參考應用程式中使用 TensorFlow Lite Task Library 中的 BertQuestionAnswerer API。

允許的匯出格式可以是一個或多個下列格式

預設情況下，它只匯出包含中繼資料的 TensorFlow Lite 模型。您也可以選擇性地匯出不同的檔案。例如，僅匯出詞彙檔案，如下所示

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

您也可以使用 evaluate_tflite 方法評估 tflite 模型。這個步驟預計會花費很長時間。

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

進階用法

create 函式是這個程式庫的關鍵部分，其中 model_spec 參數定義了模型規格。目前支援 BertQASpec 類別。有 2 個模型：MobileBERT 模型、BERT-Base 模型。create 函式包含以下步驟

  1. 根據 model_spec 建立問答模型。
  2. 訓練問答模型。

本節說明幾個進階主題，包括調整模型、調整訓練超參數等等。

調整模型

您可以調整模型基礎架構，例如 BertQASpec 類別中的參數 seq_len 和 query_len。

模型的可調整參數

  • seq_len：饋送至模型的段落長度。
  • query_len：饋送至模型的問題長度。
  • doc_stride：執行滑動視窗方法以取得文件區塊時的步幅。
  • initializer_range：用於初始化所有權重矩陣的 truncated_normal_initializer 的標準差。
  • trainable：布林值，預先訓練的層是否可訓練。

訓練管線的可調整參數

  • model_dir：模型檢查點檔案的位置。如果未設定，將使用暫時目錄。
  • dropout_rate：dropout 的比率。
  • learning_rate：Adam 的初始學習率。
  • predict_batch_size：預測的批次大小。
  • tpu：要連線的 TPU 位址。僅在使用 tpu 時使用。

例如，您可以使用更長的序列長度訓練模型。如果您變更模型，必須先建構新的 model_spec。

new_spec = model_spec.get('mobilebert_qa')
new_spec.seq_len = 512

剩餘步驟相同。請注意，您必須重新執行 dataloader 和 create 部分，因為不同的模型規格可能具有不同的預先處理步驟。

調整訓練超參數

您也可以調整訓練超參數，例如 epochs 和 batch_size，以影響模型效能。例如：

  • epochs：更多 epoch 可以實現更好的效能，但可能導致過度擬合。
  • batch_size：在一個訓練步驟中使用的樣本數。

例如，您可以使用更多 epoch 和更大的批次大小進行訓練，例如

model = question_answer.create(train_data, model_spec=spec, epochs=5, batch_size=64)

變更模型架構

您可以透過變更 model_spec 來變更資料訓練所依據的基礎模型。例如，若要變更為 BERT-Base 模型，請執行

spec = model_spec.get('bert_qa')

剩餘步驟相同。

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

深入瞭解

您可以閱讀我們的 BERT 問答範例，以瞭解技術細節。如需更多資訊，請參閱