时间戳

首页 > 解决方案 > 模型（从 Keras 导入）导致 TensorflowJS 中的内存泄漏

javascript - 模型（从 Keras 导入）导致 TensorflowJS 中的内存泄漏

问题描述

我用 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中，我加载了我的模型（带有 MobileNet 主干的 Unet）和来自body-pix的基于 MobileNet 的分割模型（比较两个模型）：

<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 导入有关？或者还有什么可能是内存泄漏的问题？

标签： javascriptkerasmemory-leaksimage-segmentationtensorflow.js

解决方案

代替：

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();

否则不会处理中间结果。

推荐阅读