Model（從 Keras 導入）導致 TensorflowJS 中的 Memory 泄漏

Question

I trained an image segmentation model with tf.keras in Python, saved it and reloaded it in with tensorflow.js (to use it in a web app).

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中，我從body-pix加載我的 model（Unet 與 MobileNet 主干）和基於 MobileNet 的分段 model（比較兩個模型）：

<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 導入有關？ 或者 memory 泄漏還有什么問題？

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

否則不會處理中間結果。

Model（從 Keras 導入）導致 TensorflowJS 中的 Memory 泄漏

問題描述

1 個解決方案

解決方案1
0 已采納 2021-05-12 15:39:02

Model（從 Keras 導入）導致 TensorflowJS 中的 Memory 泄漏

問題描述

1 個解決方案

解決方案1 0 已采納 2021-05-12 15:39:02

解決方案1
0 已采納 2021-05-12 15:39:02