Webアプリで軽快にAIモデルを動作させる (#3 IndexedDB 編)
この記事について
TensorFlow.jsを使ったAIモデルを組込んだWebアプリの性能改善についての実践記事です。なるべく自己完結型に書くようにしています。
- TensorFlow.js を拡張性を意識して Model-View-Presenterデザインパターンの中に取り込む
- ブラウザ上のAI推論タスクを Web Worker でUIスレッドと分離する
- AIモデルを IndexedDB でブラウザにキャッシュする (本記事)
サンプルアプリのソースコードは github に公開しますのでご参照ください。
想定するユースケースとビジネス背景
本シリーズが対象にしているのは、画像分類のようなAI推論を
サーバではなくブラウザ上で実行するWebアプリ
です。
推論をクライアントで完結させると、画像をサーバへ送らないためプライバシー面で有利で、推論用サーバの運用コストも不要になるという利点があります。
その反面、学習済みモデルのバイナリ（本記事の例では約17MB）をブラウザに配布する必要があり、これが初回以降の表示性能・通信コストの足かせになります。
特に、モバイル回線や従量課金のネットワークで繰り返し利用されるサービスでは、画面遷移のたびに十数MBの再ダウンロードが発生すると、待機時間の増加による離脱や通信コストの増大に直結します。
本記事で扱うキャッシュは、こうした
「2回目以降の体験」を改善し、通信量を削減する
ための施策です。
解決したい課題
初回の記事 で、TensorFlow.js で学習済AIモデルを配布してWebアプリを作る場合に気にすべき性能課題を３つ挙げましたが、この記事ではその中の
課題2. 初回ロードで、いつもモデルファイルをダウンロードするので、待機時間が長く、ネットワーク負荷が高い。
を改善したいと思います。まず、現状のサンプルアプリケーションの初回アクセス時のネットワーク負荷をChrome開発者機能のNetworkタブで見てみると以下のようになりました。
図1:
初回アクセス時のNetworkタブ。
model.json
（41.9 kB）と
group1-shard1of5.bin
〜
group1-shard5of5.bin
（各約4.2 MB ＋ 111 kB）がモデルのバイナリで、全体約20 MBの通信のうち約17
MBがモデルの転送に費やされている。
model.json
と
group1-***.bin
が我々の作ったWebモデルのバイナリで、数えてみると全20MBの通信のうち、17MBがモデルの転送に使われています。初回であれば仕方ないですが、画面の更新のたびにモデルのダウンロードに帯域を奪われるのは無駄ですので、
一度取得したモデルのバイナリファイルは、次回以降は取得しなくてよいようにブラウザにキャッシュしたいわけです。
本記事のゴールを定量的に言い換えると、 「初回ロードで発生する約17MBのモデル転送を、2回目以降のロードでは0にする（ブラウザのキャッシュから読み込む）」 ことです。
キャッシュ機構として IndexedDB を使う理由
ブラウザのキャッシュ機構としては、今回扱う IndexedDB 以外にも、 session storage や
local storage
がありますが、これをAIモデルのキャッシュに利用しづらい明確な理由があります。
まず上記の web storage
は主要なブラウザでの上限は約5MBと言われており、AIモデルのバイナリ（約17MB）を保持するには十分ではありません。
より致命的な点として、web storage は web worker スレッドでは利用できません。
そのため、UI描画とAI処理を別スレッドに分ける戦略（
前回記事
）を取る場合、 IndexedDB がほぼ一択になります。
以上を踏まえて、 IndexedDB を使って AIモデルの初回ロード時のシーケンスを以下のように変更することを行います。
図2: IndexedDB を導入した初回ロード時のシーケンス図。App初期化時にまず IndexedDB を参照し、モデルが無い場合に限ってサーバから取得・保存する。
ポイントは、
- App初期化の時、まずは IndexedDB からモデルのバイナリを取得するように試みる
- IndexedDBに該当のバイナリファイルが存在しない場合に、初めてWebサーバからモデルを取得する
ように実装するということです。
TensorFlow.js API による IndexedDB 操作
IndexedDB
は、 ブラウザのキャッシュ機構の中でも、非同期的なアクセスができ、Web
Workerスレッドからも参照可能な key-value ストアです。
容量制限も実質無く、主要な全てのWebブラウザでサポートされています。
TensorFlow.js は IndexedDB を内部でサポートしており、IndexedDB の
API仕様には深入りすることなく、IndexedDB のキャッシュ機構を実現できます。
なお、本記事で読み込むモデルは
tf.loadGraphModel
が返す
tf.GraphModel
（計算グラフ形式の推論専用モデル。学習はできないが、その分ロード・推論が高速）を前提としています。
以下、本記事で使う TensorFlow.js の API を挙げていきます。
IndexedDB へのキャッシュ書き込み
IndexedDB に web モデルを保存するには、
// mobilenet web モデルを Webサーバ localhost:8000 に配置した場合
const modelUrl = 'http://localhost:8000/models/mobilenet_1.0.0/model.json';
const model = await tf.loadGraphModel(modelUrl);
のようにしてロードした
model
に対して
const modelName = 'mobilenet_1.0.0'
const saveResults = await model.save(`indexeddb://${modelName}`);
を実行します。これを実行した後で、 Chrome の Applicatons タブの IndexedDB の中身を見ると、
-
tensorflowjsデータベース-
model_info_storeオブジェクトストア: モデルのメタ情報が格納される -
models_storeオブジェクトストア: モデルの重みなどのパラメータが格納される
-
が生成されていることがわかります。
図3:
model_info_store
オブジェクトストア。モデル本体は持たず、
modelTopologyBytes
（32,227 バイト）や
weightDataBytes
（16,888,228 バイト ≒ 約16.1 MiB）といったメタ情報のみが格納されている。
図4:
models_store
オブジェクトストア。
weightData
として約16.1 MiB（16,888,228 バイト）の
ArrayBuffer
が保存されており、図1でサーバから転送していたモデルの重みがそのままブラウザ内にキャッシュされていることが確認できる。
TensorFlow.js の仕様で、
indexeddb://...
を指定した
model.save
が実行されると、 上記の
tensorflowjs
データベースと 2つのオブジェクトストアが勝手に作成されます。
上の例では、オブジェクトストアには キーが
mobilenet_1.0.0
のオブジェクトが１つ保存されているのが確認できます。
TensorFlow.js では
model.save(indexeddb://${modelName})
を実行した時の
modelName
がオブジェクトストアのキーになります。
IndexedDB からモデル読み込み
IndexedDB からモデルをロードするには、単に
const loadedModel = await tf.loadGraphModel(`indexeddb://${modelName}`);
とするだけです。ただし、このとき IndexedDBの
tensorflowjs
データベースに、
modelName
と同じキーのモデルファイルが保存されている必要があります。そうでなければ例外がスローされます。
IndexedDB に保存済みのモデルの一覧
const list = await tf.io.listModels();
で、IndexedDBに保存されているモデルを取得することができます。この
list
はモデルのパスをキーにもつ連想配列で、単にモデルのリストアップをするだけならキーだけ参照すれば十分です。
console.log(Object.keys(list));
// [ 'indexeddb://mobilenet_1.0.0' ]
IndexedDB からのキャッシュの削除
最後に削除です。
await tf.io.removeModel('indexeddb://mobilenet_1.0.0')
実装
以上のAPIで実装したいことは簡単にできてしまいますので、この節で実装の説明をします。
ただし、実用上、注意しなければいけない問題が２つあります。
-
モデル更新の問題
- 新しいモデルファイルを更新配布したい場合、Webサーバを更新しても、クライアント側で IndexedDB にキャッシュされたモデルを利用され続けてしまうことが起こりえます。そこで、サーバとクライアントでモデルのバージョン管理を行う仕組みが必要になります。
-
不要なモデルの削除
- バージョン管理できたとして、古いバージョンのモデルファイルがローカルに残り続けるとユーザのストレージをどんどん圧迫してしまうため、不要なモデルは削除する必要があります。
このことに注意して、前回の記事で Web Worker 側で実装された AI ロードの実装クラスを修正します。
モデル初期化クラス
AiLoader
実装クラスを以下のように修正します。
// app/model/ai-loader/impl/AiLoader.ts
import {inject, injectable} from 'inversify';
import {IAiLoader} from '../IAiLoader';
import * as tf from '@tensorflow/tfjs';
import {IImageClassifier} from '../../image-classifier/IImageClassifier';
import {TYPES} from '../../../../config/dependency/types';
import {REQUIRED_MODELS} from '../../../../config/required-models';
const now = () => new Date().getTime();
@injectable()
export class AiLoader implements IAiLoader {
constructor(
@inject(TYPES.ImageClassifier) private imageClassifier: IImageClassifier,
) {}
async load(): Promise<void> {
// TensorFlow.js のセットアップを行う
const t00 = now();
await tf.ready();
console.info('tfjs backend=', tf.getBackend());
const t01 = now();
console.info('TensorFlow.js initialized. Elapsed:', t01 - t00, 'msec');
// 不必要なモデルを IndexedDB から削除する
const localModels = await tf.io.listModels();
const requiredIdbModelUrl = REQUIRED_MODELS.map(modelName => `indexeddb://${modelName}`);
for (const modelUrl of Object.keys(localModels)) {
if (requiredIdbModelUrl.includes(modelUrl)) continue;
await tf.io.removeModel(modelUrl);
console.info('model in indexedDB removed:', modelUrl);
}
// MobileNetV1 モデルのセットアップを行う
await this.imageClassifier.load();
}
}
アプリのロード時に
load
メソッドが実行されますが、その時に
REQUIRED_MODELS
リストに無いモデルは全て IndexedDB から削除します。
これによって、サーバ側で新しいモデルが配布されて
REQUIRED_MODELS
が更新されると、ローカルにキャッシュされた古いモデルは削除させることができます。
推論クラス
推論実装クラス
ImageClassifier
も
load
処理を以下のように修正します
// app/model/image-classifier/impl/ImageClassifier.ts
// (中略)
@injectable()
export class ImageClassifier implements IImageClassifier {
private mobilenet!: tf.GraphModel;
modelName = 'mobilenet_1.0.0';
async load(): Promise<void> {
const t00 = now();
// MobileNetV1 モデルのセットアップを行う
const localModels = await tf.io.listModels();
const idbModelUrl = `indexeddb://${this.modelName}`;
const webModelUrl = `/models/${this.modelName}/web/model.json`;
// indexedDBにモデルが保存されていれば、indexedDBから読み込む。
// そうでなければ webモデルを読み込み、indexedDBに保存する。
if (idbModelUrl in localModels) {
this.mobilenet = await tf.loadGraphModel(idbModelUrl);
console.info('model loaded from indexedDB:', idbModelUrl)
} else {
this.mobilenet = await tf.loadGraphModel(webModelUrl);
console.info('web model loaded:', webModelUrl)
await this.mobilenet.save(idbModelUrl);
console.info('model saved in indexedDB:', idbModelUrl)
}
// warm-up する（初回推論時の遅延を避けるため、ダミー入力で一度推論を回し、
// 演算カーネルのコンパイルやGPUテクスチャの確保を事前に済ませておく）
const resultTf = this.mobilenet.predict(tf.zeros([1, 224, 224, 3])) as tf.Tensor;
resultTf.dataSync();
resultTf.dispose();
const t01 = now();
console.info('MobileNetV1 initialized. Elapsed:', t01 - t00, 'msec');
}
// (中略)
}
コードコメントに書いてある通りですが、
tf.io.listModels
で取得した IndexedDB のキーの中に自分のバイナリがあれば IndexedDB
から読み込み、無ければ Webサーバから読み込んで IndexedDB に保存します。
この実装により、初回アクセス時には図1のとおり約17MBのモデルがサーバから転送・保存されますが、2回目以降のアクセスでは同じバイナリを
IndexedDB から読み込むため、
モデルに関するサーバ通信が0バイトになります
。
ネットワーク転送は回線の細い環境ほどロード時間のボトルネックになりやすいため、繰り返し利用されるシーンでの待機時間短縮と通信量削減の効果が期待できます。
実装は以上で終わりです。TensorFlow.js で用意されているAPIだけで簡単に実装できました。
結論
本記事では、WebアプリにおけるAIモデルの初回ロード性能を改善するため、TensorFlow.jsとIndexedDBを用いたキャッシュ戦略を解説しました。
冒頭で提示した課題「初回ロードで発生する約17MBのモデル転送を、2回目以降は0にする」は、以下の実装によって達成されました。
- 初回アクセス時 : Webサーバからモデルをダウンロードし、そのバイナリをIndexedDBに保存する。
- 2回目以降のアクセス時 : IndexedDBにキャッシュされたモデルを直接読み込むことで、サーバへのネットワークリクエストを完全に不要にする。
このアプローチにより、ユーザーは2回目以降のアクセスでモデルのダウンロードを待つ必要がなくなり、特にモバイル回線など通信環境が限定的な状況でのUXを大幅に改善します。また、事業者側にとってもデータ転送コストの削減に繋がります。
実装にあたっては、TensorFlow.jsが提供する
model.save('indexeddb://...')
や
tf.loadGraphModel('indexeddb://...')
といった高レベルAPIを利用することで、IndexedDBの複雑な詳細を意識することなく容易にキャッシュ機構を導入できました。ただし、キャッシュを実運用する上では、本記事で示したような
モデルのバージョン管理
と
不要なキャッシュの削除処理
を組み込むことが不可欠です。
以上のことから、TensorFlow.jsを利用するWebアプリにおいて、IndexedDBによるモデルキャッシュは、パフォーマンスとユーザー体験を向上させるための非常に効果的かつ実践的な手法であると言えます。