TensorFlow Liteの自然言語処理の新機能
以下の記事を参考に書いてます。
・What’s new in TensorFlow Lite for NLP
1. はじめに
「TensorFlow Lite」は、携帯電話、マイクロコントローラ、Edge TPUなどのEdgeデバイスに機械学習機能を提供するために、多くのアプリで広く採用されています。人々の生活をより簡単に、より生産的にする人気の高いアプリの中でも、「自然言語理解」は、研究コミュニティと業界の両方から注目を集めている重要な分野の一つです。2019年の「TensorFlow World」でオンデバイスの「質問応答」のデモを行った後、このようなNLPモデルをより多くのオンデバイス推論で利用できるようにすることについて、コミュニティから多くの関心とフィードバックをもらいました。
そのフィードバックに触発されて、本日、「TensorFlow Lite」をベースにしたNLPタスクのエンドツーエンドサポートを発表することができました。このインフラ作業により、より多くのNLPモデルが携帯電話上で実行できるようになり、ユーザーは個人データをオンデバイスに保ちながら、NLPモデルの利点を享受できるようになります。今回のブログでは、それを可能にする新機能を紹介します。
(1) 新しい事前学習NLPモデルの使用
(2) 独自のNLPモデルの作成
(3) TensorFlow NLPモデルのTensorFlow Lite形式への変換
(4) NLPモデルをオンデバイスで実行
2. 新しい事前学習NLPモデルの使用
◎ リファレンスアプリ
リファレンスアプリは、あらかじめ学習された「機械学習モデル」「推論コード」「実行可能なデモ」をカプセル化したオープンソースのモバイルアプリのセットです。「Android Studio」「Xcode」と統合されたリファレンスアプリを提供しているので、開発者はワンクリックでビルドし、Android端末やiOS端末で利用できます。
以下のリファレンスアプリを使用して、モバイル開発者は、既存のNLPモデル（BERT、MobileBERT、ALBERT）を統合し、生のテキストデータを変換し、モデルの入出力を接続して予測結果を生成するまでのエンドツーエンドの流れを学ぶことができます。
・テキスト分類 : 与えられたテキストに基づいてラベルを予測。
・質問回答 : 記事と質問が与えられると、モデルは記事内の質問に答える。
・スマートリプライ : 与えられたコンテキストに基づいて自動返信を生成。
上記のリファレンスアプリで使用した事前学習モデルは、TensorFlow Hubで利用できます。下図は、モデル間の「レイテンシ」「サイズ」「F1スコア」を比較したものです。
3. オンデバイスに最適化されたNLPモデル
オンデバイスモデルは、サーバサイドモデルとは異なる制約があります。ここでは、NLPタスクのためにモデルを最適化する方法の例をいくつか紹介します。
◎ Quantized MobileBERT
「MobileBERT」は、GitHubで公開されているコンパクトなBERTモデルです。BERT（float32）よりも4.3倍小さく、5.5倍高速でありながら、GLUEおよびSQuADデータセットで同等の結果を達成しています。
初期リリース後、量子化を利用してモデルサイズと性能を最適化することで、GPU/DSPなどのアクセラレータが利用できるように改良しました。量子化された「MobileBERT」は、BERTよりも16倍小さく、8倍速く、精度の損失はほとんどありません。「MLPerf for Mobile」コミュニティは、量子化された「MobileBERT」をモバイル推論のベンチマークに活用しており、このモデルはTensorFlow.jsを使用してChrome上で実行することもできます。
BERT（416MB）と比較すると、以下の表は、同じ設定の下で量子化された「MobileBERT」の性能を示しています。
◎ 投影法を用いたEmbedding-free NLPモデル
「言語識別」は、与えられたテキストの言語を分類するタスクです。最近、私たちは投影法を用いた2つのモデル、「SGNN」と「PRADO」をオープンソース化しました。
「SGNN」を用いて、「Tensorflow Lite」をNLPタスクに利用することがいかに簡単で効率的であるかを示しました。「SGNN」は、テキストを固定長の特徴量に投影し、それに続いて全結合層に投影します。「TensorFlow Lite」のコンバータに「TF.Text API」を融合させるようにアノテーションを与えることで、「TensorFlow Lite」でより効率的な推論モデルを得ることができます。以前のベンチマークでは、モデルの実行に1332.87μsかかっていましたが、融合後は同じマシン上で64.06μsとなりました。これは、20倍のスピードアップという魔法をもたらします。
また、「PRADO」と呼ばれるモデルアーキテクチャのデモも行っています。「PRADO」は、まず単語トークンのシーケンスから学習可能な投影特徴量を計算し、次に畳み込みとアノテーションを適用して特徴量を固定長のエンコーディングにマッピングします。「PRADO」は、投影層、畳み込み層、アテンションエンコーダーを組み合わせることで、LSTMと同様の精度を達成しますが、モデルサイズは100倍小さくなります。
これらのモデルの背景にあるアイデアは、テキストから特徴を計算するために投影を使用することで、モデルがテキストの特徴を埋め込むために大きな埋め込みテーブルを維持する必要がないようにすることです。このようにして、同様の性能と推論の待ち時間を維持しながら、モデルが埋め込みベースのモデルよりもはるかに小さくなることが証明されています。
4. 独自のNLPモデルの作成
「TensorFlow Lite」では、事前学習モデルを使用するだけでなく、「Model Maker」などのツールを使用して、既存のモデルを独自にカスタマイズすることができます。
◎ TensorFlow Lite Model Maker
「TensorFlow Lite Model Maker」は、最先端の機械学習モデルをデータセットに適応させるための転送学習ツールです。モバイル開発者が機械学習の専門知識がなくてもモデルを作成することができ、転移学習によって学習時間を短縮することができます。
Visionタスクに焦点を当てた初期リリースの後、最近、「Model Maker」に2つの新しいNLPタスクを追加しました。「テキスト分類」「質問応答」については、Colabとガイドに従ってください。
「Model Maker」をインストールするコマンドは、次のとおりです。
pip install tflite-model-makerモデルをカスタマイズするために、開発者は以下のように数行の python コードを書く必要があります。
# データの読み込み
train_data = TextClassifierDataLoader.from_csv(train_csv_file, mode_spec=spec)
test_data = TextClassifierDataLoader.from_csv(test_csv_file, mode_spec=spec)
# TensorFlowモデルのカスタマイズ
model = text_classifier.create(train_data, model_spec=spec)
# TensorFlow Liteモデルの評価
loss, acc = model.evaluate(test_data)
# TensorFlow Liteモデルのエクスポート
model.export(export_dir, quantization_config=config)5. TensorFlow NLPモデルのTensorFlow Lite形式への変換
「TensorFlow Lite」のオペレータは、「TensorFlow」のオペレータのサブセットしかサポートしていないため、NLPモデルを「TensorFlow Lite」形式に変換する際に、オペレータが見つからなかったり、サポートされていないデータ型（RaggedTensor、ハッシュテーブル、アセットファイルの処理など）が原因で問題に遭遇したことがあるかもしれません。このような場合の変換問題を解決する方法について、いくつかのヒントを紹介します。
◎ TensorFlow LiteでTensorFlow opsとTF.text opsを実行
「Select TensorFlow ops」を強化し、様々なケースに対応できるようにしました。「Select TensorFlow ops」を使用すると、開発者は「TensorFlow ops」を利用して「TensorFlow Lite」上でモデルを実行できます。「TensorFlow Lite」に相当するopsがない場合には、「TensorFlow ops」を利用します。例えば、TensorFlowモデルを学習する際に、「TF.Text ops」や「RaggedTensor」を使うのが一般的ですが、それらのモデルを「TensorFlow Lite」に簡単に変換して必要なopsで実行できるようになりました。
さらに、モバイルデプロイ用にトリミングされたバイナリが得られるように、op選択的ビルドを使用する解決策も提供しています。これは、最終的なビルドターゲットで使用されるopsの小さなセットを選択することができ、その結果、デプロイ時のバイナリサイズを小さくすることができます。
◎ より効率的でフレンドリーなカスタムオペレーション
「TensorFlow Lite」では、「Ngram」「SentencePieceTokenizer」「WordPieceTokenizer」「WhitespaceTokenizer」など、NLPのためのいくつかの新しいモバイルフレンドリーなオペレーションを提供します。
以前は、「SentencePiece」を使用したモデルが「TensorFlow Lite」に変換されるのをブロックするいくつかの制限がありました。モバイル用の新しい「SentencePieceTokenizer API」は、これらの課題を解決し、同時に実装を最適化して高速に動作するようにしました。
同様に、「Ngram」「WhitespaceTokenizer」もサポートされるようになっただけでなく、オンデバイスでより効率的に実行されるようになります。
「TensorFlow Lite」は先日、MLIRとのオペレーションの融合を発表しました。同じメカニズムを使って「TF.Text API」を「TensorFlow Lite」のカスタムオペレーションに融合させ、推論効率を大幅に向上させました。例えば、「WhitespaceTokenizer API」は複数の演算で構成されており、「TensorFlow Lite」のオリジナルグラフでの実行に0.9msかかっていました。これらの演算を1つの演算に融合した後、0.04msで終了し、23倍のスピードアップを実現しています。このアプローチにより、上述のSGNNモデルにおいて推論レイテンシが大幅に向上することが証明されています。
◎ ハッシュテーブルのサポート
「ハッシュテーブル」は多くのNLPモデルにとって重要です。なぜなら、通常、言語モデルでは単語を「トークンID」に変換したり、その逆を行ったりして数値計算を利用する必要があるからです。「ハッシュテーブル」はまもなく「TensorFlow Lite」で有効になります。これは、「TensorFlow Lite」形式でアセットファイルをネイティブに扱い、「TensorFlow Lite」の組み込みオペレータとしてopカーネルを提供することでサポートされます。
6. NLPモデルをオンデバイスで実行
「TensorFlow Lite」を使った推論の実行は、以前よりもはるかに簡単になりました。事前構築された推論APIを使用して5行のコードでモデルを統合したり、ユーティリティを使用して独自のAndroid・iOSの推論APIを構築したりすることができます。
◎ TensorFlow Liteタスクライブラリを使用した簡単なモデル展開
「TensorFlow Lite Task Library」は、ML推論の前処理と後処理を簡単に利用できる強力で使いやすいライブラリです。アプリ開発者は「TensorFlow Lite」を使って機械学習機能を簡単に実装することができます。「Task Library」でサポートされているテキストAPIは3つあり、上記のユースケースやモデルに対応しています。
・NLClassifier : 入力テキストを，既知のカテゴリの集合に分類。
・BertNLClassifier : BERT系モデルに最適化されたテキストを分類。
・BertQuestionAnswerer : BERT系モデルを使用して、与えられた文章の内容に基づいて質問に答える。
「Task Library」は、Android ・iOS の両方でクロスプラットフォームで動作します。以下は、Java・Swift の「BertQA」モデルを使用した推論の例です。
// 初期化
BertQuestionAnswerer answerer = BertQuestionAnswerer.createFromFile(androidContext, modelFile);
// 質問応答
List answers = answerer.answer(context, question);// 初期化
let mobileBertAnswerer = TFLBertQuestionAnswerer.mobilebertQuestionAnswerer(modelPath: modelPath)
// 質問応答
let answers = mobileBertAnswerer.answer(context: context, question: question)◎ カスタマイズ推論API
既存の「Task Library」でサポートされていないユースケースの場合は、「Task API Infrastructure」を活用して、同じリポジトリにある「Wordpiece」や「Sentencepiece」のトークナイザーなどの一般的なNLPユーティリティを使って、独自のC++/Android/iOS推論APIを作成することもできます。
7. おわりに
今回の記事では、「TensorFlow Lite」でNLPタスクが新たにサポートされたことを紹介しました。「TensorFlow Lite」の最新アップデートにより、開発者は簡単にNLPモデルをオンデバイスで作成、変換、実行できるようになりました。今後もより便利なツールを提供し、研究から本番までのオンデバイスNLPモデルの開発を加速させていきます。皆様からのフィードバックや、新しいNLPツールやユーティリティへのご提案をお待ちしています。