模型（从 Keras 导入）导致 TensorflowJS 内存泄漏

Question

我在 Python 中用 tf.keras 训练了一个图像分割模型，保存它并用 tensorflow.js 重新加载它（在网络应用程序中使用它）。

Python（转移模型）：

import tensorflow as tf
import tensorflowjs as tfjs

model = tf.keras.models.load_model('model')
tfjs.converters.save_keras_model(model, 'tfjs_model/')

在 Javascript 中，我加载了我的模型（Unet with MobileNet backbone）AND 一个 MobileNet基于 body-pix 的分割模型（比较两个模型）：

<head>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-converter"></script>
    <script src="https://cdn.jsdelivr.net/npm/@tensorflow-models/body-pix"></script>
    ...
</head>
<body>
    ...
    const model_lib = await bodyPix.load(); 
    // NumBytesInGPU: 5349984 (every log)

    const model_own = await tf.loadLayersModel('mobilenet_test/model.json');
    // NumBytesInGPU: 36930448 (1st log), 53707664 (5th log)
</body>

这样做效果很好，所有内容都可以无误地加载。然而，当尝试根据视频进行预测时，tf.memory() 会增加直到应用程序崩溃，而 body-pix 模型运行平稳。

async function body_segment() {
    const frame = document.getElementById("camera");
    const canvas = document.getElementById("body_pix");
    const draw = canvas.getContext("2d");

    // const model = await bodyPix.load(); 
    // NumBytesInGPU: 5349984

    const model = await tf.loadLayersModel('mobile_net.json');
    // NumBytesInGPU: 36930448 (1st log), 53707664 (5th log)


    const runPrediction = function(input) {
        return tf.tidy(() => {
            const asFloat = tf.cast(input, 'float32');
            const asBatch = tf.expandDims(asFloat, 0);
            const results = model.predict(asBatch);
            // Normally do something additional here, but removed due to debug reasons
            return results
        });
    }

    const resized = function(input) {
        return tf.tidy(() => {
            let imageTensor = tf.browser.fromPixels(input);
            return tf.image.resizeBilinear(imageTensor, [512, 512]);
        })
    }

    let ctr = 0;
    while (ctr < 10) {
        console.log("memory", tf.memory());



        // !!!!! THIS FUNCTION CAUSES THE MEMORY LEAK, BUT WHY ?????
        const result = await runPrediction(resized(video)); 
        // const result = await model.segmentPersonParts(frame);

        // do something with prediction here ...

        result.dispose(); // remove results from prediction to clean the memory

        ctr+=1;
        await tf.nextFrame();
    }    
}

我尝试使用与 body-pix 文件中完全相同的代码。此外，我一直使用 tidy 函数，所以实际上，它应该对所有内容进行垃圾回收。

是否与 Keras 导入有关？或者内存泄漏还有什么问题？

Answer 1

而不是：

const result = await runPrediction(resized(video)); 
// do smt
result.dispose();

使用

const res = await resized(video);
const result = await runPrediction(res);
res.dispose();
// do smt
result.dispose();

否则不会处理中间结果

模型（从 Keras 导入）导致 TensorflowJS 内存泄漏

Model (imported from Keras) causes Memory leak in TensorflowJS

javascript

memory-leaks

image-segmentation

keras

tensorflow.js