「ヘイ Friko」を聞き分けるモデルを、画面を固めずにブラウザで動かす

Metronome.fm に、マスコットの名前で「ヘイ Friko」と呼びかけると、声でメトロノームを操作する機能が始まる合言葉を作っています。合言葉を聞き分けるモデルは、自分で学習させた小さなものです（学習の話は「ヘイ Friko」の呼びかけを聞き分けるモデルを、自分で学習させたに書きました）。

この記事では、そのモデルをブラウザで動かした話を書きます。合言葉は、マイクの音を端末の中でずっと聞き分けていて、音はどこにも送りません。聞き分けの処理をブラウザの中で続けることになるので、その処理が画面を止めないように気をつけました。

3 段のモデルを、onnxruntime-web で動かす

モデルは、合言葉の聞き分けを作るオープンソースの仕組み openWakeWord の形で、3 段に分かれています。

1. 音（16kHz）をメルスペクトログラムにする
2. Google の音声埋め込みモデルで、96 次元の埋め込みにする
3. 埋め込みの直近 16 個を、自分で学習させた分類器に入れて、合言葉の確率を出す

openWakeWord はブラウザでの推論に対応していないので、openWakeWord が流れてくる音を処理する部分のコードを読んで、同じ手順を TypeScript で書きました。音が 80ms（1,280 サンプル）分たまるごとに、直前の 1,760 サンプル（新しい 1,280 サンプルと、その前の 480 サンプル）からメルスペクトログラムを 8 フレーム作り、openWakeWord と同じく ÷10 + 2 で値を整えます。そのうえで、直近 76 フレームから埋め込みを 1 つ作ります。

// 要点だけにまとめたもの（run・tensor・rowsOf は、記事のために短くした名前）
private async processChunk(chunk: Float32Array): Promise<number> {
  // 直前の 480 サンプルを残して、新しい 1,280 サンプルを後ろに足す
  buffers.raw.copyWithin(0, CHUNK);
  buffers.raw.set(chunk, MEL_CONTEXT);

  const spec = await run(sessions.mel, { input: tensor(buffers.raw) });
  for (const row of rowsOf(spec, MEL_BINS)) buffers.mel.push(row.map((v) => v / 10 + 2));
  buffers.mel.splice(0, buffers.mel.length - MEL_WINDOW);

  const embedding = await run(sessions.embedding, { input_1: tensor(buffers.mel, [1, 76, 32, 1]) });
  buffers.features.push(embedding);
  buffers.features.shift();

  return (await run(sessions.classifier, { input: tensor(buffers.features, [1, 16, 96]) }))[0];
}

比べる基準として、openWakeWord と同じ手順で確率を出すスクリプトを Python でも書き、同じ録音で確率を比べました。305 秒・3,817 回の判定で、確率の差は最大 0.000009 でした。1 回の判定にかかる時間は 2.7ms で、80ms に 1 回なので、処理に使うのは時間の 3% ほどです。

マイクの音は 44.1kHz や 48kHz で届くので、AudioWorklet で 16kHz に下げます。窓をかけた sinc の低域通過フィルター（63 タップ）を通してから、サンプルの間を補いながら間引いています。これを確かめるために、録音をいったん 44.1kHz に上げ、AudioWorklet で 16kHz に下げ直して比べました。合言葉とみなした回数（確率が境目の 0.5 以上になった回数）は、Python の 33 回に対して 31 回でした。違いは、確率が境目すれすれの数回で結果が分かれた程度です。

WASM は自分のサイトから配信する

onnxruntime-web は、WASM のファイル（11MB、gzip で 2.7MB）を読み込んで動きます。Metronome.fm は CSP で外の CDN から読み込ませていないので、ビルドのときに node_modules から public/ort/ にコピーして、サイトから配信しています。合言葉をオンにしたときにだけ読み込みます。

ここで 2 つつまずきました。

• ort.env.wasm.wasmPaths に場所を文字列で渡すと、WASM だけでなく、それを読み込む役の ort-wasm-simd-threaded.mjs もそこから読みにいきます。最初は WASM だけをコピーしていて、404 になりました
• 一度 404 になると、ファイルを置いたあとも同じページの中では読み込めませんでした。ブラウザは、同じページの中で失敗した動的 import の結果を使い回すためです。ページを読み込み直すと直りました

CSP には、WebAssembly のコンパイルだけを許す 'wasm-unsafe-eval' を加えました。onnxruntime-web が中で複数のスレッドを使うには crossOriginIsolated が要りますが、サイトはそうしていないので、1 本のスレッドで動かしています（numThreads = 1）。

メインスレッドで動かすと、準備のときに止まった

最初は、判定をメインスレッドで動かしていました。マイクの音は AudioWorklet から postMessage でメインスレッドに届き、そこで 3 段のモデルを動かします。メトロノームの音は AudioWorklet の中で鳴らしているので、判定が音を遅らせることはありません。

それでも、画面のほうは止まっていました。PerformanceObserver で 50ms を超えた処理（long task）を数えると、合言葉をオンにした直後に、約 50ms の処理が 2 回ありました。onnxruntime-web がセッション（モデルを読み込んで動かせる状態にしたもの）を準備する部分です。判定しているあいだは、80ms に 3ms ほどの処理がすきまに収まっていました。

ただ、これは速い Mac での数字です。古いスマホでは同じ処理に数倍かかることがあるので、準備のときに画面が目に見えて止まり、判定のあいだもメインスレッドを使い続けるおそれがあります。

音は AudioWorklet から Worker へ直接渡す

そこで、モデルの準備も含めて、聞き分けの処理をすべて Web Worker に移しました。onnxruntime-web にも、推論を Worker で動かす設定（proxy）があります。ただ、それだと音がいったんメインスレッドを通るので、自分で Worker を作り、AudioWorklet から Worker へ音を MessagePort で直接渡すことにしました。

// AudioWorklet と Worker を、メインスレッドを通らない MessagePort でつなぐ
const channel = new MessageChannel();
node.port.postMessage({ port: channel.port1 }, [channel.port1]);
worker.postMessage(
  { type: "init", audio: channel.port2, classifierFile: CLASSIFIER_FILE },
  [channel.port2],
);

AudioWorklet は、受け取った port1 に 16kHz の音を約 30ms ずつ送ります。Worker は port2 で音を受け取り、80ms 分たまるごとに判定します。モデルの準備と判定に加えて、確率が境目の 0.5 以上なら合言葉とみなすことと、一度聞き取ったら 2 秒は次を数えないことも、すべて Worker の中で行います。メインスレッドに返すのは、準備ができた・聞き取った・エラー、の知らせだけです。Worker のコードは、AudioWorklet のコードと同じく esbuild で 1 つのファイルにまとめていて、onnxruntime-web の JavaScript を含めて 50KB でした（WASM のファイルは別に読み込みます）。

聞き取りは、移す前と変わりませんでした。合言葉を言い終えてから、操作の頼みの聞き取りが始まるまでは約 0.4 秒です。