TensorFlow Lite Model Maker によるテキスト分類

TensorFlow.org で表示 Google Colab で実行 GitHub でソースを表示 ノートブックをダウンロード

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

このノートブックでは、Model Maker ライブラリを使用して、モバイルデバイスで映画レビューを分類するために一般的に使用されるテキスト分類モデルの適応と変換を示す、エンドツーエンドの例を説明します。テキスト分類モデルは、テキストを既定のカテゴリに分類します。入力は前処理済みのテキストで、出力はカテゴリの確率です。このチュートリアルで使用されるデータセットは肯定的な映画レビューと否定的な映画レビューです。

前提条件

必要なパッケージをインストールする

この例を実行するには、GitHub リポジトリ から、Model Maker パッケージを含む必要なパッケージをインストールする必要があります。

pip install -q tflite-model-maker

必要なパッケージをインポートします。

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

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

サンプルデータをダウンロードする

このチュートリアルでは、SST-2(Stanford Sentiment Treebank)を使用します。これは GLUE ベンチマークのタスクの 1 つで、トレーニング用の 67,349 件の映画レビューとテスト用の 872 件の映画レビューが含まれています。データセットには肯定的なレビューと否定的なレビューの 2 つのクラスが含まれます。

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 はタブ \t 区切りであるのに対し、CSV はカンマ , 区切りであることです。

次に、トレーニングデータセットの最初の 5 行を示します。label=0 は否定的、label=1 は肯定的です。

文章 ラベル
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 データフレームに読み込み、現在のラベル名(01)を人間が理解しやすいもの(negativepositive)に変更し、モデルのトレーニングにそれらを使用します。

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')

クイックスタート

テキスト分類モデルのトレーニングには 5 つのステップがあります。

ステップ 1. テキスト分類モデルアーキテクチャを選択する。

ここでは、平均的な単語埋め込みモデルアーキテクチャを使用します。これで、適度な正確度のある小さな高速モデルが生成されます。

spec = model_spec.get('average_word_vec')

Model Maker では、BERT などのモデルもサポートしています。ほかのアーキテクチャについても興味がある場合は、以下の「テキスト分類器のモデルアーキテクチャを選択する」セクションをご覧ください。

ステップ 2. training と test データを読み込んで、特定の model_spec に従ってそれらを処理する。

Model Maker は CSV 形式の入力データを取ることができます。先ほど作成した、人が理解できるラベル名を使った training データセットと test データセットを読み込みます。

それぞれのモデルアーキテクチャでは、入力データを特定の方法で処理する必要があります。DataLoadermodel_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. training データを使用して TensorFlow モデルをトレーニングする。

平均的な単語埋め込みモデルではデフォルトで、batch_size = 32 が使用されます。そのため、training データセットの 67,349 文を処理するのに 2104 ステップが掛かります。モデルを 10 エポックでトレーニングするため、トレーニングデータセットは 10 回処理されます。

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

ステップ 4. test データでモデルを評価する。

training データセット内の文章を使用してテキスト分類モデルをトレーニングしたら、test データセットにある残りの 872 文を使用して、モデルがそれまでに見たことのない新しいデータに対してどのように実行するのかを評価します。

デフォルトのバッチサイズは 32 であるため、test データセット内の 872 文の処理に 28 ステップが掛かります。

loss, acc = model.evaluate(test_data)

ステップ 5. TensorFlow Lite モデルをエクスポートする。

撮れ人ぐ下テキスト分類を TensorFlow Lite 形式でエクスポートしましょう。モデルの手クスポート先のフォルダを指定します。デフォルトでは、TFLite モデルは平均的な単語埋め込みモデルアーキテクチャ向けにエクスポートされます。

model.export(export_dir='average_word_vec')

TensorFlow Lite モデルファイルは、Colab の左サイドバーを使ってダウンロードできます。上記の export_dir パラメーターに指定したとおり、average_word_vec フォルダに移動して、model.tflite ファイルを右クリックして Download を選択すると、ローカルコンピューターにダウンロードされます。

このモデルは、NLClassifier APITensorFlow Lite Task ライブラリ)を使って Android または iOS アプリに統合することができます。

実動アプリでモデルがどのように使用されているかについては、「TFLite テキスト分類のサンプルアプリ」をご覧ください。

注意 1: Android Studio Model Binding はテキスト分類をまだサポートしていないため、TensorFlow Lite Task ライブラリを使用してください。

注意 2: TFLite モデルと同じフォルダに model.json ファイルがあります。このファイルには、TensorFlow Lite モデル内にバンドルされたメタデータの JSON 表現が含まれています。モデルのメタデータは、TFLite Task ライブラリがモデルが何を行い、どのようにモデルのデータを事前処理/事後処理するのかを伝える上で役立ちます。model.json ファイルは情報提供の目的でのみ存在し、そのコンテンツは TFLite 内に存在するため、ダウンロードする必要はありません。

注意 3: MobileBERT または BERT-Base アーキテクチャを使ってテキスト分類モデルをトレーニングする場合は、BertNLClassifier API を代わりに使って、トレーニング済みのモデルをモバイルアプリに組み込む必要があります。

次のウォークスルーでは、例の詳細を手順を追って説明します。

テキスト分類器のモデルアーキテクチャを選択する

model_spec オブジェクトは、テキスト分類器の特定のモデルを表現します。TensorFlow Lite Model Maker は現在、MobileBERT、平均的な単語埋め込み、および BERT-Base モデルをサポートしています。

サポートされているモデル model_spec の名前 モデルの説明 モデルサイズ
MobileBERT 'mobilebert_classifier' BERT ベースより 4.3 倍小さく、5.5 倍高速ですが、オンデバイスアプリケーションに適した、優位性のある結果を達成します。 1MB 未満
BERT ベース 'bert_classifier' NLP タスクで広く使用される標準的な BERT モデルです。 25MB、量子化あり
100MB、量子化なし
BERT-Base 'average_word_vec' RELU アクティベーションを使った平均テキスト単語埋め込みです。 300MB

クイックスタートでは、兵器的な単語埋め込みモデルを使用しました。MobileBERT に切り替えて、より高い正確度でモデルをトレーニングしてみましょう。

mb_spec = model_spec.get('mobilebert_classifier')

トレーニングデータを読み込む

このチュートリアルでは独自のデータセットをアップロードして進めることができます。データセットをアップロードするには、Colab の左サイドバーを使ってください。

ファイルのアップロード

データセットをクラウドにアップロードしない場合は、ガイドに従ってローカルでライブラリを実行することもできます。

説明をわかりやすくするために、先ほどダウンロードした SST-2 データセットを再利用することにします。DataLoader.from_csv メソッドを使って、データを読み込みましょう。

モデルのアーキテクチャを変更したため、training データセットと test データセットを読み込み直して新しいプリプロセッシングロジックを適用する必要があります。

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() メソッドをサポートしています。同じクラスのテキストデータは同じサブディレクトリに存在し、サブフォルダ名はクラス名として前提づけられています。各テキストファイルには、1 つの映画レビューサンプルがあります。class_labels パラメータはどのサブフォルダかを指定するために使用されます。

TensorFlow モデルをトレーニングする

training データを使用してテキスト分類モデルをトレーニングします。

注意: MobileBERT は複雑なモデルであるため、各エポックには Colab GPU で約 10 分掛かります。必ず GPU ランタイムを使用してください。

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

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

model.summary()

モデルを評価する

test データ使用して、トレーニングしたばかりのモデルを評価し、損失と正確度の値を測定します。

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 APITensorFlow Lite Task ライブラリ)を使ってモバイルアプリに組み込むことができます。これは、平均的な単語埋め込みモデルアーキテクチャでトレーニングしたテキスト分類を組み込むために使用した NLClassifier API とは異なることに注意してください。

エクスポート形式は次のいずれかを使用できます。

デフォルトでは、モデルのメタデータを含む TensorFlow Lite モデルのみをエクスポートしますが、より詳しく調査するために、モデルに関連するほかのファイルをエクスポートすることを選択することもできます。たとえば、ラベルファイルと語彙ファイルのみをエクスポートする場合は、次のようにします。

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

TFLite モデルを evaluate_tflite メソッドを使って評価し、その正確度を測定することができます。トレーニング済みの TensorFlow モデルを TFLite 形式に変換して量子化を適用すると、その正確度に影響が与えられるため、デプロイ前に TFLite モデルの正確度を評価することをお勧めします。

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

高度な使用

create 関数は、Model Maker ライブラリがモデルの作成に使用するドライバー関数です。model_spec パラメータによってモデルの仕様が定義されています。現在サポートされているクラスは、AverageWordVecSpecBertClassifierSpec です。create 関数の構成は、次の手順で行います。

  1. model_spec に基づいてテキスト分類器のモデルを作成します。
  2. 分類器モデルをトレーニングします。デフォルトのエポックとデフォルトのバッチサイズは、default_training_epochsdefault_batch_size の変数に従って model_spec オブジェクトに設定されています。

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

MobileBERT モデルのハイパーパラメータをカスタマイズする

調整できるモデルパラメータは次のとおりです。

  • seq_len: モデルにフィードするシーケンスの長さ。
  • initializer_range: すべての重み行列を初期化する truncated_normal_initializer の標準偏差。
  • trainable: トレーニング済みレイヤーがトレーニング可能かどうかを示すブール型。

調整できるモデルパラメータは次のとおりです。

  • model_dir: モデルチェックポイントファイルの場所。設定されていない場合、一時ディレクトリが使用されます。
  • dropout_rate: ドロップアウト率。
  • learning_rate: Adam オプティマイザの初期学習率。
  • tpu: 接続先の TPU アドレス。

調整できるトレーニングパイプラインは次のとおりです。

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

平均的な単語埋め込みモデルのハイパーパラメータをカスタマイズする

AverageWordVecModelSpec クラスのwordvec_dimseq_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)

トレーニングハイパーパラメータの調整

モデルの精度に影響する epochsbatch_size などのトレーニングハイパーパラメータを調整できます。次にその例を示します。

  • epochs: エポックを増やすと、より優れた精度を達成できますが、過適合となる可能性があります。
  • batch_size: 1 つのトレーニングステップに使用するサンプル数。

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

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 ベースモデルに変更する例を示します。

テキスト分類器の model_spec を BERT ベースモデルに変更します。

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)

その他の資料

技術的な詳細については、テキスト分類の例をご覧ください。詳細については、以下をご覧ください。