Company

Webアプリで軽快にAIモデルを動作させる (#1 TensorFlow.js 導入編)

記事の目的

当社ではこれまで、AIモデルをブラウザで動作するように最適化したWebアプリを数多く開発・リリースしてきました。

サーバでのAI推論に対して、 クライアントサイドでのAI推論(オンデバイスAI) を選択するのには、明確なビジネスメリットとコンプライアンス上の強みがあります。

通常、画像解析やリアルタイム推論をサーバーサイドで行う場合、大量のリクエストを処理するために高価なGPUサーバー(クラウドGPUサーバの維持費は、少なくとも 月額約5万〜10万円以上の価格帯となる)を常時維持、あるいは負荷に応じてスケールアウトさせる必要があります。
クライアント推論方式を採用すると、計算リソースはユーザーの端末(PCやスマートフォン)にオフロードされるため、サーバー側の推論計算コストを実質0に削減できます。アクセス数が数万、数百万と急増した際でも、サーバー維持費の増加を最小限(静的ファイルの配信コストのみ)に抑えられます。

また、「カメラ映像」や「個人情報・機密書類のデータ」をAIで解析する場合、データを一切外部サーバーに送信せず、ユーザーのブラウザ内だけで処理を完結できることは極めて重要です。
これにより、厳格なデータ移転・保管規制に対するコンプライアンス要件もクリアしやすくなります。通信上のデータ傍受や、サーバー側での情報漏洩リスク自体を設計段階から排除でき、ユーザーの信頼獲得にも大きく寄与します。

しかし、これらのメリットの裏返しとして、計算リソースがユーザー環境に完全に依存するため、サーバー方式とは質の異なる「パフォーマンスとUXの課題」が発生します。特に以下の課題は、Web技術のもとでクライアント推論を行う場合は避けて通れないものです。

クライアント推論特有のパフォーマンス課題

  1. AI推論処理によるUIスレッドの占有(描画遅延)
    ブラウザのJavaScriptは基本的にシングルスレッド(メインスレッド)で動作します。AI初期化や推論処理は極めて高負荷なCPU/GPU計算を伴うため、メインスレッドが完全に占有され、画面上のアニメーションがカクついたり、ボタン操作が効かなくなったりします。
  2. モデルファイルのダウンロード負荷
    AIモデルはファイルサイズが大きいため、初回ロード時のネットワーク負荷が高くなります。本サンプルで使用する軽量な画像分類モデル「MobileNetV1 (1.0 224)」であっても、**ファイルサイズは約17MB(Webモデル変換後も同様の容量)**に達します。モバイル回線や低速なネットワーク環境では、初期起動までに10秒以上の待機時間が発生するリスクがあります。
  3. 実行環境による推論速度の極端な格差
    クライアントPCのグラフィックス機能(WebGL)が利用可能な環境(GPU推論)と、非対応でCPU推論しか行えない環境では、処理速度に数十倍もの差が生じます。

これらの課題を解決するため、当社では Web Worker (別スレッドへのAI処理分離)、 IndexedDB (ブラウザ内でのモデルキャッシュによるロード削減)、 WebAssembly (WASM) バックエンド (CPU推論時の高速化)などの Web API や拡張仕様を調査していきました。

シリーズを通して扱うサンプルアプリ

以下のようなアプリをサンプルとして扱いたいと思います。

本記事で実装する画像分類サンプルアプリの動作デモ

仕様は文章で書くと、以下のようなものです。

  • 入力:
    • ローカルに存在する画像ファイルを一つ選択し、アップロードする。
  • 出力:
    • 入力画像ファイルに対して、以下の推論処理を行う。
      • 画像分類タスクの事前学習モデル (MobileNetV1) で推論を行う。
    • 推論結果のラベルとスコアを、画面に表示する。

シーケンス図

アプリの処理フローを把握するために、シーケンス図を以下に示します。

クライアント推論における初期状態のシーケンス図

注意したいのは、Webブラウザでは、通常はスレッド一つで処理が行われるということです。AI初期化や推論処理は高負荷なため、UI描画をしばしば阻害しやすくなってしまいます。

なおJavaScript には Promise 等の擬似的な非同期処理の仕組みがあるのはご存じの方が多いと思います。確かに通信の応答待ちなどWebブラウザが計算タスクを持たない場合には待ち時間中にUI処理はブロックされません。ところがAI初期化や推論など、Webブラウザ側で本当にやるべき計算がある場合には、Promise等の非同期制御だけでは結局UI処理が阻害されてしまうことに注意してください。

つまり、ユーザがWebブラウザを開いてから操作可能な画面が帰ってくるまでの間、ユーザはブラウザの画面が固まっている状態を体験してしまいます。

この問題は、別の記事で説明する Web Worker を用いたWorkerスレッドの分離で解決しますが、本記事はひとまずアプリの骨組みを作るところまで進めましょう。

ソフトウェア・アーキテクチャ

まずアプリの骨組みとなる、クラス設計を行います。ここでは、いわゆる Model-View-Presenter モデルを採用し、画面表示と描画ロジック、そしてAI処理を完全に分離することを目指します。

依存関係を明示したクラス図は以下のようになります。

MVPパターンとDIを適用したクラス図

それぞれの役割は、

要素 役割
View Presenterからの描画指示を受けて、描画する。また、UIイベントが発生したら処理を Presenter に委譲する。
Presenter 描画ロジックを実行し、View に描画指示を行う。
Model AiLoader : AIの初期化処理を行う。 ImageClassifier : 画像データを受け取り、推論を行い、推論結果を返す。

と定めます。ポイントは、TensorFlow.js に依存する実装はすべて Model の領域で閉じ込めるようすることです。
これにより、AIの推論処理と、UI処理が完全に分離され、AI部分の改善に対しても改修範囲をModel内部に局所化することができます。

フレームワーク・ライブラリ

実装に関連するライブラリ選定をしておきます。

まずフロントエンドのプロダクト開発で型付き言語は必須なので Typescript を前提にします。
また、この記事は特定のUIフレームワークに依存したものにしたくないため、React、Vue.js、Angular などのライフサイクルを司るフレームワークは利用しません。
ただし、先に示したソフトウェア・アーキテクチャを実現する上で、 Dependency Injection(DI: 依存性の注入) を活用することは、ソフトウェア拡張性維持のために非常に有効です。インターフェース(仕様)に対して具体的な実装クラスを疎結合に注入するため、軽量なDIコンテナライブラリである InversifyJS ( inversify.js ) reflect-metadata を利用します。

以上を踏まえて、 package.json は以下のようにします。
(ビルドには、 esbuild を利用します)

{
  "name": "tfjs-web-sandbox",
  "version": "1.0.0",
  "description": "",
  "main": "index.js",
  "scripts": {
    "build": "esbuild --bundle ./src/index.ts --outfile=./www/index.js",
    "start": "yarn build --servedir=./www"
  },
  "author": "",
  "license": "ISC",
  "devDependencies": {
    "@typescript-eslint/eslint-plugin": "^5.8.1",
    "@typescript-eslint/parser": "^5.8.1",
    "esbuild": "^0.14.9",
    "eslint": "^8.6.0",
    "eslint-config-google": "^0.14.0",
    "typescript": "^4.4.4"
  },
  "dependencies": {
    "@tensorflow/tfjs": "^3.12.0",
    "inversify": "^6.0.1",
    "reflect-metadata": "^0.1.13"
  }
}

学習済のfrozen model(.pb)をWebモデルに変換する

実装に入る前に、今回利用するモデルをTensorFlow.jsで取込可能なWebモデルに変換する必要があります。

この解説では、TensorFlow frozen model形式のpbファイルを、事前に作成済とします。
ここでは、画像を1000クラスに分類する MobileNetV1 事前学習モデルを、 ここ から入手しました。

pbファイルを tensorflowjs_converter でWebモデルに変換します。

pip install tensorflowjs

(2023/8更新:converterは tensorflowjs にバンドルされるようになりました)

インストールができたら、 pbファイルを指定してWebモデルを生成します。
./frozen/mobilenet_v1_1.0_224_frozen.pb このパスのpbファイルを配置してから、

$ tensorflowjs_converter \
  --input_format=tf_frozen_model \
  --output_node_names='MobilenetV1/Predictions/Reshape_1' \
  ./frozen/mobilenet_v1_1.0_224_frozen.pb ./web   

を実行します。成功すると、 model.json ファイルと、Webモデルのバイナリファイルが ./web 配下に配置されます。

--output_node_names オプションには、モデルの出力ノード名称を指定する必要がありますので、入力モデルのアーキテクチャを確認して正しく設定します。

生成されたWebモデルファイルは、コンテンツサーバ上に以下のように公開します。

<DocumentRoot>
└── /models/mobilenet
    ├── group1-shard1of5.bin
    ├── group1-shard2of5.bin
    ├── group1-shard3of5.bin
    ├── group1-shard4of5.bin
    ├── group1-shard5of5.bin
    └── model.json

AIの実装 (Model部分)

それでは、AI部分の実装を始めていきましょう。

画像分類クラス: ImageClassifier

まず画像分類クラスから始めます。仕様を決めるインターフェースを作ってから、具体的な実装を作る、というふうにインターフェースと実装を意識して分離します。
別記事で説明する Web Worker の実装の際に、この分離が効きます。

インターフェース定義

まずインターフェースを、以下のように決めます。

// file path: app/model/image-classifier/IImageClassifier.ts

export type ImageClassifierOutput = { classId: number, score: number };

export interface IImageClassifier {
  load(): Promise<void>
  classify(input: ImageData): Promise<ImageClassifierOutput>
}

つまり、 ImageClassifier は、Webモデルファイルをロードする機能( load )と、 入力となる画像データを受け取って推論結果を返す機能( classify )を持つこととします。

classify の引数について、なぜ入力画像を HTMLImageElement 等のDOM要素ではなく、 ImageData 形式にしているのか、気になった方もあるかもしれません。
その理由は、Web WorkerスレッドにAI処理を分離した際に、Web WorkerスレッドからDOMが読めないからなのですが、別記事で詳しく説明します。

実装

さてインターフェースを、以下のように実装します。

// file path: app/model/image-classifier/impl/ImageClassifier.ts

import {injectable} from 'inversify';
import * as tf from '@tensorflow/tfjs';
import {IImageClassifier, ImageClassifierOutput} from '../IImageClassifier';


@injectable()
export class ImageClassifier implements IImageClassifier {
  private mobilenet!: tf.GraphModel;

  async load(): Promise<void> {
    // MobileNetV1 モデルのセットアップを行う
    this.mobilenet = await tf.loadGraphModel(
        '/models/mobilenet/web/model.json',
    );

    // warm-up する
    const resultTf = this.mobilenet.predict(tf.zeros([1, 224, 224, 3])) as tf.Tensor;
    resultTf.dataSync();
    resultTf.dispose();
  }

  async classify(input: ImageData): Promise<ImageClassifierOutput> {
    // 入力画像をモデル入力サイズ [1, 224, 224, 3] に変換する
    const inputTf = tf.tidy(() => {
      return tf.browser
          .fromPixels(input)
          .resizeBilinear([224, 224])
          .expandDims(0);
    });

    // 推論を行い、確率ベクトル [1, 1000] を取得する
    const resultTf = this.mobilenet.predict(inputTf) as tf.Tensor;

    // 最大確率をもつラベルIDを取得する
    const argmaxTf = tf.tidy(() => resultTf.squeeze().argMax());

    // スコア値を取得する
    const classId = argmaxTf.dataSync()[0];
    const score = resultTf.dataSync()[classId];

    // 使い終わったtensorを開放する
    inputTf.dispose();
    resultTf.dispose();
    argmaxTf.dispose();

    return {classId, score};
  }
}

まず load メソッドでは、WebモデルのURLを指定してモデルを初期化し、空のテンソルを渡して推論処理を初回実行し、モデルを完全にロードさせます。これを「Warm-up(ウォームアップ)」と呼び、初回推論時の実行時遅延を事前に消化しておく有用な手法となります。
WebモデルのURLは、先程 tensorflowjs_converter で生成したWebモデルのjsonのURLを指定します。

推論メソッド classify では、TensorFlow.js 固有のメモリ管理や同期処理 API の特性について以下の点に注意します。

  • tf.tidy() による自動的な一時テンソルのメモリ解放
    モデルファイルの入力形式に合わせて、画像データをテンソルに変換する必要があります。今回利用した MobileNetV1 は入力テンソルの形式が [1, 224, 224, 3] に指定されているので、それに合わせて入力の ImageData を変換します。
    このテンソル変換処理を tf.tidy() で囲っています。TensorFlow.jsでは、JavaScriptの通常のガベージコレクション(GC)が、GPU(WebGL)上のメモリ(WebGLテクスチャ等)を自動解放できません。 tf.tidy() を使用すると、そのスコープ内で生成された一時的な中間テンソル(この例では、 fromPixelsresizeBilinearexpandDims など、各ステップで内部生成される中間状態のテンソル)を自動的に検知して一括解放し、GPUメモリリークを防ぎます。

  • 明示的なメモリ解放 ( dispose )
    tf.tidy() の戻り値として返されたテンソル( inputTf など)や、スコープ外( predict メソッドの実行など)で新しく生成されたテンソルは、使い終わった後に明示的に .dispose() メソッドを呼び出してメモリから解放する必要があります。

  • 出力テンソルからのクラスID抽出
    このモデルファイルの出力テンソル形式は [1, 1000] です。各クラス(1000種類)に対する予測確率(スコア)が含まれているため、 argMax (最大確率をもつ要素のインデックスを取得するAPI)を使用し、最も確からしいクラスIDとそのスコア(確率)を取り出しています。

AI実行環境の初期化: AiLoader

個別のモデルファイルを初期化する前に、TensorFlow.js の初期化 tf.ready() を実行し、最適なバックエンドを設定する必要があります。
そこで、 ImageClassifier とは別の初期化用のサービスクラスを作ります。

インターフェース定義と、実装を分離して以下のようにします。

// file path: app/model/ai-loader/IAiLoader.ts

export interface IAiLoader {
  load(): Promise<void>;
}
// file path: 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';

@injectable()
export class AiLoader implements IAiLoader {
  constructor(
    @inject(TYPES.ImageClassifier) private imageClassifier: IImageClassifier,
  ) {}

  async load(): Promise<void> {
    // TensorFlow.js のセットアップを行う
    await tf.ready();
    console.info('tfjs backend=', tf.getBackend());

    // MobileNetV1 モデルのセットアップを行う
    await this.imageClassifier.load();

    // 複数モデルを使う場合は、ここにloadを追記していく...
  }
}

load メソッドでは、2つの重要な処理を行っています。

  • tf.ready() によるバックエンドの非同期初期化と選定
    tf.ready() は、TensorFlow.js が現在のプラットフォーム(WebGL、WebAssembly、CPU など)に適した計算バックエンドを自動選定し、内部エンジンのロードと非同期初期化を完了するのを待機するAPIです。初期化が完了したあと、現在アクティブなバックエンド名を tf.getBackend() で取得してコンソールに出力します。選定の仕様は、おおよそ以下の通りです。
    • ハードウェアアクセラレータが実行できるWebブラウザ環境の場合は webgl バックエンドが選定される
    • 上記以外の場合は cpu バックエンドが選定される
  • アプリで利用するAIモデルを個別にセットアップします。

TensorFlow.js のバックエンドについては、WebAssembly 導入時に触れますので、別記事で詳しく整理したいと思います。

以上で AI部分の実装は完了しました。

UIの実装 (View, Presenter)

あとは、Presenter と View を実装するだけですが、このUI関連部分の実装はこのシリーズの解説の目的からはずれてくるので、細かくは説明しません。

参考までに、Presenter の実装は以下のように実装しました。

  • アプリ初期化のとき initialize メソッドが呼ばれ、AIモデルがロードされる
  • ファイルアップロード時に processFile メソッドが呼ばれ、 画像の推論処理が行われる

この2点だけ、注目いただければ良いと思います。

import {inject, injectable} from 'inversify';
import {TYPES} from '../../../config/dependency/types';
import {IAiLoader} from '../../model/ai-loader/IAiLoader';
import {IImageClassifier} from '../../model/image-classifier/IImageClassifier';
import {IIndexView} from '../../view/index/IIndexView';

@injectable()
export class IndexPresenter {
  constructor(
    @inject(TYPES.AiLoader) private aiLoader: IAiLoader,
    @inject(TYPES.ImageClassifier) private imageClassifier: IImageClassifier,
    @inject(TYPES.IndexView) private indexView: IIndexView,
  ) {
    this.indexView.setPresenter(this);
  }

  async initialize(): Promise<void> {
    this.indexView.showLoading(true);
    await this.aiLoader.load();
    this.indexView.showLoading(false);
  }

  async processFile(imageData: ImageData): Promise<void> {
    this.indexView.showLoading(true);
    const res = await this.imageClassifier.classify(imageData);
    this.indexView.updateResult(`class id: ${res.classId}, score: ${res.score})`);
    this.indexView.showLoading(false);
  }
}

本記事のまとめと次回の展望

本記事では、Webブラウザ上で動作するオンデバイスAI(クライアント推論)の第一歩として、TensorFlow.jsを用いた画像分類アプリの骨組みを実装しました。

本記事で達成したこと

  • クリーンなアーキテクチャの確立 : MVP(Model-View-Presenter)パターンと DI(依存性の注入)パターンを組み合わせることで、TensorFlow.jsの初期化・メモリ管理・推論ロジックをModel層( ImageClassifierAiLoader )に完全に隔離しました。これにより、UI描画(View)や制御(Presenter)がAIエンジンの具体的な実装に依存しない保守性の高い設計となっています。
  • Webモデル変換ワークフローの整理 : Pythonの tensorflowjs_converter を活用し、既存の pb(frozen model)ファイルをWebブラウザからダウンロードしてパースできるJSON分割形式(Web Model)に変換・公開する手順を明確化しました。

未解決の課題

冒頭で挙げた「クライアント推論におけるパフォーマンス課題」のうち、本記事の「初期骨組み」段階では、以下の部分が未解決のままとなっています。

  • 課題1:メインスレッド占有によるUIフリーズ(カクつき)
    推論処理( dataSync() )やモデルのWarm-up、大きな画像データのロード計算がメインスレッド上で行われるため、推論実行中や初期化時にローディングアニメーションがカクついたり、ボタンの応答がなくなったりするUX上の不満があります。
  • 課題2:約17MBのモデルロードによる初期遅延とネットワーク負荷
    ページを開くたびに17MBを超えるモデルファイルをコンテンツサーバーからダウンロードしてロードするため、初回ロード完了まで長い待機時間が発生します。
  • 課題3:CPUバックエンド時の著しい低パフォーマンス
    WebGL(GPU)がサポートされていない古い端末や特定のブラウザ、あるいはサーバーサイド等では、自動的に cpu バックエンドにフォールバックされます。この場合、1推論に数百ms〜1.5秒以上の時間がかかり、実用的なインタラクションが得られません。

これらの未解決課題を解決するために、本記事で構築した「インターフェースと実装を分離した疎結合設計」が大きな威力を発揮します。UI側のPresenterやViewに一切変更を加えずに、Model層の実装を最適化されたものに差し替えることで、段階的にアプリをプロフェッショナルな品質へと引き上げていきます。

  1. 次回 (#2 Web Worker編) :
    AI推論プロセス全体を、バックグラウンドの別スレッド( Web Worker )へ分離します。これにより、メインスレッド(UIスレッド)のブロッキングを完全に解消し、100%滑らかに動き続けるローディングアニメーションと軽快なUXを実現します。
  2. 次々回 (#3 IndexedDBキャッシュ編) :
    ブラウザ上のキーバリューストアである IndexedDB を活用し、一度ダウンロードしたモデルファイルをブラウザ内にキャッシュします。2回目以降のアクセスではサーバー通信をスキップし、即座にローカルからモデルを復元して起動速度を爆発的に向上させます。
  3. さらにその先 (#4 WASMバックエンド編) :
    CPUフォールバック環境における実行速度向上を目指し、最適化された C++ コードをブラウザで実行する WebAssembly (WASM) バックエンド (SIMD / Multithreading対応)を導入し、WebGL非対応端末でのパフォーマンスを定量的に底上げします。

参考文献

TensorFlow.js 公式のAPIドキュメント

tensorflowjs_converter 公式ドキュメント (github)

InversifyJS (DIコンテナ) 公式ドキュメント

MobileNet系事前学習モデルのアーカイブ