TensorFlow Lite Model Maker による BERT 質疑応答

TensorFlow.org で表示

Google Colab で実行

GitHub でソースを表示

ノートブックをダウンロード

TensorFlow Lite Model Maker ライブラリは、TensorFlow モデルを適合し、オンデバイス ML アプリケーションにこのモデルをデプロイする際に特定の入力データに変換するプロセスを単純化します。

このノートブックでは、Model Maker ライブラリを使用したエンドツーエンドの例を示し、質疑応答タスクで一般的に使用される質疑応答モデルの適合と変換を説明します。

BERT 質疑応答タスクの基礎

このライブラリでサポートされているタスクは、抽出型の質疑応答タスクです。特定の文章と質問に対する回答が文章に含まれていることになります。以下の画像は、質疑応答の例を示します。

回答は文章に含まれている（画像提供: SQuAD ブログ）

質疑応答タスクのモデルでは、入力は、すでに前処理されている文章と質問のペアで、出力は、文章の各トークンの開始ロジットと終了ロジットです。入力のサイズは設定可能で、文章と質問の長さに応じて調整することができます。

エンドツーエンドの概要

次のコードスニペットでは、数行のコード内でモデルを取得する方法を示します。全体的なプロセスには、（1）モデルの選択、（2）データの読み込み、（3）モデルの再トレーニング、（4）評価、（5）TensorFlow Lite 形式へのエクスポート、という 5 つのステップが含まれます。

# 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 ベースモデルをサポートしています。

サポートされているモデル	model_spec の名前	モデルの説明
MobileBERT	'mobilebert_qa'	BERT ベースより 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')

オンデバイス ML アプリ固有の入力データを読み込み、データを前処理する

TriviaQA は、読解問題のデータセットで、質問、回答、エビデンスの 3 つを 1 組とした 65 万個を超えるデータが含まれます。このチュートリアルでは、このデータセットのサブセットを使用して、Model Maker ライブラリの使用方法を学習します。

データを読み込むには、--sample_size=8000 としたコンバータ用 Python スクリプトと一連の 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 メソッドを使用して、SQuAD 形式データを読み込み、特定の model_spec に従って前処理します。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 関数は次のステップで構成されています。

model_spec に基づいて質疑応答のモデルを作成します。
質疑応答モデルをトレーニングします。デフォルトのエポックとデフォルトのバッチサイズは、default_training_epochs と default_batch_size の 2 つの変数に従って model_spec オブジェクトに設定されています。

model = question_answer.create(train_data, model_spec=spec)

モデル構造を詳しく確認します。

model.summary()

カスタマイズ済みのモデルを評価する

検証データでモデルを評価し、f1 スコアや exact match などを含むメトリクスの dict を取得します。SQuAD1.1 と SQuAD2.0 のメトリクスは異なることに注意してください。

model.evaluate(validation_data)

TensorFlow Lite モデルをエクスポートする

トレーニングされたモデルをメタデータで TensorFlow Lite モデル形式に変換し、後でオンデバイス ML アプリケーションで使用できるようにします。語彙ファイルはメタデータに埋め込まれています。デフォルトの TFLite ファイル名は model.tflite です。

多くのオンデバイス ML アプリケーションでは、モデルサイズが重要な要因です。そのため、モデルの量子化を適用して小さくし、実行速度を高められるようにすることをお勧めします。デフォルトのポストトレーニング量子化手法は、BERT および MobileBERT モデルのダイナミックレンジ量子化です。

model.export(export_dir='.')

Colab の左サイドバーからダウンロードすることで、bert_qa 参照アプリで BertQuestionAnswerer API（TensorFlow Lite Task ライブラリ）を使って TensorFlow Lite モデルファイルを使用できます。

次のいずれかのエクスポートフォーマットを使用できます。

デフォルトでは、メタデータとともに TensorFlow Lite モデルをエクスポートするだけです。さまざまなファイルを選択的にエクスポートすることも可能です。たとえば、vocab ファイルのみをエクスポートする場合は、次のように行います。

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

また、evaluate_tflite メソッドを使って tflite モデルを評価することも可能です。このステップには長い時間がかかります。

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

高度な使用

create 関数は、model_spec パラメータがモデルの仕様を定義するため、このライブラリでは重要な部分です。現在、BertQASpec クラスがサポートされています。モデルには MobileBERT モデルと BERT ベースモデルの 2 つがあります。create 関数は次のステップで構成されています。

model_spec に基づいて質疑応答のモデルを作成します。
質疑応答モデルをトレーニングします。

このセクションでは、モデルの調整やトレーニングハイパーパラメータの調整など、いくつかの高度なトピックを説明します。

モデルの調整

BertQASpec クラスの seq_len や query_len パラメータなどのモデルインフラストラクチャを調整できます。

モデルの調整可能なパラメータは次のとおりです。

seq_len: モデルにフィードする文章の長さ。
query_len: モデルにフィードする質問の長さ。
doc_stride: ドキュメントの塊を取るためのスライドウィンドウアプローチを実行する際のストライド。
initializer_range: すべての重み行列を初期化する truncated_normal_initializer の stdev。
trainable: トレーニング済みレイヤーがトレーニング可能かどうかを示すブール型。

トレーニングパイプラインの調整可能なパラメータは次のとおりです。

model_dir: モデルチェックポイントファイルの場所。設定されていない場合、一時ディレクトリが使用されます。
dropout_rate: ドロップアウトのレート。
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: エポック数が多いほどパフォーマンスが改善されますが、過適合が発生する可能性があります。
batch_size: 1 つのトレーニングステップに使用するサンプル数。

たとえば、エポック数とバッチサイズを増やしてトレーニングすることができます。

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

モデルアーキテクチャを変更する

model_spec を変更して、データをトレーニングするベースモデルを変更できます。たとえば、BERT ベースモデルに変更するには、次を実行します。

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 質疑応答の例をご覧ください。詳細については、以下をご覧ください。

TensorFlow Lite Model Maker のガイドと API リファレンス
タスクライブラリ: デプロイ用の BertQuestionAnswerer
エンドツーエンドリファレンスアプリ: Android およおび iOS