Open3

TypeScript+TensorFlow.jsでメモリリークする問題(WebGL: CONTEXT_LOST_WEBGL: loseContext: context lost)

TypeScript
TensorFlow.js
PINTOPINTO

はじめに

タイトルに TypeScript と記載したがほとんど関係ない。問題の本質は、Javascript と TensorFlow.js (TFJS) の dynamic tensor と NonMaxSuppression の実装。

前提

  • 入力テンソルに dynamic tensor を含む場合や NonMaxSuppression を含むモデルを TensorFlow.js で実行すると様々な問題が重なってメモリリークしてアプリケーションが最終的に Abort する。おそらく、音声処理系のモデルはほぼ確実にこの問題を引くと思う。

  • NonMaxSuppression を含むモデルは model.execute() では実行できないにも関わらず、TensorFlow.js が無意味に This model execution did not contain any nodes with control flow or dynamic output shapes. You can use model.execute() instead. を大量に出力する

  • model.execute() の代わりに await model.executeAsync() を使用しないと推論を実行できない

問題点. 1

  • 特に WebGL バックエンドを使用しているとき、割り当てられた領域が GC で回収されず早いタイミングで WebGL: CONTEXT_LOST_WEBGL: loseContext: context lost というエラーが発生して Abort する

対策

  • TensorFlow.js でテンソル操作を行うときは tf.tidy() で必ずラップする
  • 念の為、生成された tf オブジェクトは try-catch-finally の構文の finally句 などで必ず dispose() が実行されるようにしておくと安全
    test.ts
    // dispose() 対象オブジェクトのストック用リスト
const tensorsToCleanup: tf.Tensor[] = [];

// imageData を取得してくる何らかの処理
//  :
// imageData を取得してくる何らかの処理

// 前処理
// tf.tidy() によるラップ
// const でいちいち受けなくても `.` でつなげてしまえば良いが
// tensor1.reverse(-1) が any を返す問題があるので TypeScript ではやめたほうが良い
// TensorFlow.js のトップレベル関数の `tf.xxx()` をあえて使用している
const tensors = tf.tidy(() => {
  const tensor1 = tf.browser.fromPixels(imageData, 3);
  const tensor2 = tf.reverse(tensor1, -1);
  const tensor3 = tf.expandDims(tensor2, 0);
  const tensor4 = tf.cast(tensor3, "float32");
  return [tensor1, tensor2, tensor3, tensor4];
});
// dispose() 対象オブジェクトのストック
// 上記以外にもテンソルの操作を行っている場合はすべて tensorsToCleanup へ push しておく
tensors.forEach((tensor, ) => {
  tensorsToCleanup.push(tensor);
});

// 推論実行 - ココで無意味なワーニングが大量出力される
// This model execution did not contain any nodes with control flow or
// dynamic output shapes. You can use model.execute() instead.
const results = (await model.executeAsync(tensors[3])) as tf.Tensor;

// results を使用した何かの後処理
//  :
// results を使用した何かの後処理

// ストックされた dispose() 対象オブジェクトを強制的に dispose()
tensorsToCleanup.forEach((t) => t.dispose());

問題点. 2

  • TensorFlow.js が推論１回ごとに意味の無いワーニングメッセージを console へ出力することがあり、ワーニングメッセージが console に１行出力されるたびにメモリリークする
    // ココ
const results = (await model.executeAsync(tensors[3])) as tf.Tensor;
  • 特に dynamic tensor と await model.executeAsync() を併用して推論を実行しているときに発生する問題
  • 長時間稼働していると This model execution did not contain any nodes with control flow or dynamic output shapes. You can use model.execute() instead. というどうでもいいワーニングが数百万件consoleに出力される
  • 一方で、TensorFlow.js のAPIにはワーニングメッセージを抑止するインタフェースが実装されていない
  • したがって、TensorFlow.js のせいで徐々にメモリリークしていき、やがて自爆してAbortする

対策

  • TensorFlow.js が console へ無意味に大量出力するワーニングメッセージだけを hook するスクリプトを別ファイルで書く

  • 主処理を記述している側のロジックの import 文で処理の最初に import して console.warn へのワーニングメッセージ出力を抑止する

  • なお、デバッグ目的のために自ら記述した console.log や console.info, console.warn, console.error でも、１回実行されるだけでメモリリークしていく https://stackoverflow.com/questions/54204313/does-console-log-increase-memory-on-nodejs-server

  • ChatGPT と壁打ちしてようやく下記のメモリリーク回避手段を得た

    consoleWarnSuppression.ts
    // 既存の console.warn を退避しておく
const originalWarn = console.warn;
console.warn = (...args: unknown[]): void => {
    // 引数を連結して検索しやすくする
    const joinedMsg = args.map((v) =>
        typeof v === "object" ? JSON.stringify(v) : String(v)
    ).join(" ");
    // TensorFlow.js の警告で抑制したい文字列をチェック（実際の文言に合わせて要調整）
    if (joinedMsg.includes("This model execution")) {
        return;
    }
    // その他のメッセージは従来通り表示
    originalWarn(...args);
};
    test.ts
    import * as tf from "@tensorflow/tfjs";
import { loadGraphModel } from "@tensorflow/tfjs-converter";

// TensorFlow.js が強制的に大量に出力する警告メッセージの無効化
import "./consoleWarnSuppression";
PINTOPINTO
  • 参考： console.log, console.info, console.warn, console.error などのコンソール出力でメモリリークするのは平常営業

https://stackoverflow.com/questions/54204313/does-console-log-increase-memory-on-nodejs-server

  • 参考： TensorFlow.js のissueが何も解決していないのに強制クローズされている問題

https://github.com/tensorflow/tfjs/issues/6249

  • 参考：Warning を抑止するプルリクエストが発行されているが、Googleの中の人たち自身が作ったモノのバグの本質を理解していない。世紀末。

https://github.com/tensorflow/tfjs/pull/6309