几次 decodeJpeg 调用后内存不足

Out of memory after a couple of decodeJpeg calls

我想使用 tfjs-node 使用 Tensorflow 处理几张图像(实际上相当多)。

不幸的是,在我的 Raspberry Pi 上,我很快遇到了内存不足的错误。 这是 Typescript 中的一个片段:

import * as tfnode from '@tensorflow/tfjs-node'
import * as fs from 'fs'
import * as path from 'path'
import * as process from 'process'

let counter = 0

setInterval(() => {
    const imageBuffer = fs.readFileSync(path.join(__dirname, 'samples', 'image.jpg'))

    const tfimage = tfnode.node.decodeJpeg(imageBuffer)

    console.log(`${++counter} - ${process.memoryUsage().rss / 1024 / 1024} MB`)
}, 100)

这输出类似

1 - 216.390625 MB
2 - 357.78515625 MB
3 - 499.421875 MB
4 - 641.5703125 MB
5 - 782.9453125 MB
6 - 924.3203125 MB
7 - 1066.7265625 MB
8 - 1208.09765625 MB
9 - 1349.4765625 MB
10 - 1491.625 MB
11 - 1633.2578125 MB
2020-05-25 22:36:12.101809: W tensorflow/core/framework/op_kernel.cc:1651] OP_REQUIRES failed at cast_op.cc:109 : Resource exhausted: OOM when allocating tensor with shape[3024,4032,3] and type int32 on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu
/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:404
            throw ex;
            ^

Error: Invalid TF_Status: 8
Message: OOM when allocating tensor with shape[3024,4032,3] and type int32 on /job:localhost/replica:0/task:0/device:CPU:0 by allocator cpu
    at NodeJSKernelBackend.executeSingleOutput (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-node/dist/nodejs_kernel_backend.js:193:43)
    at NodeJSKernelBackend.cast (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-node/dist/nodejs_kernel_backend.js:1141:21)
    at engine_1.ENGINE.runKernelFunc.x (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/ops/array_ops.js:139:78)
    at /home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:542:55
    at /home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:388:22
    at Engine.scopedRun (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:398:23)
    at Engine.tidy (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:387:21)
    at kernelFunc (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:542:29)
    at /home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:553:27
    at Engine.scopedRun (/home/pi/develop/tftest/build/node_modules/@tensorflow/tfjs-core/dist/engine.js:398:23)

我尝试了 https://github.com/Qengineering/TensorFlow-Raspberry-Pi and https://github.com/yhwang/node-red-contrib-tf-model 的 arm 库,两者的行为相同。

然后我在 Windows 上尝试了这个非常片段的测试,我得到了非常惊人的结果。内存消耗迅速增加到 8GB,然后它在大约 800MB 和 8GB 之间不可重复地波动。当然可以,因为系统的内存比树莓派多很多。

但是我可以用树莓派做什么?我可以控制 Tensorflow 的内存管理吗?

实际上我之前没有使用过 TensorFlow,但我对 setInterval 部分有所了解。

我认为当 setInterval 循环执行新操作时,之前执行的旧操作未完成,可能会创建太多堆栈未完成的操作,其中大量内存已被使用旧的和仍然活跃的操作。

因此您可能希望按顺序执行此操作,例如使用传统的 while 或 for 循环,而不像 setImmediate 方法那样阻塞事件循环。

我不确定我的解决方案,但给它一个机会。

是setInterval惹的祸吗?

可能存在与如何使用 setInterval 有关的内存问题,但通常它与 setInterval 使用闭包保持对某些对象的引用的回调有关。话虽如此,javascript 引擎中的垃圾收集器 GC 已经改进得如此之好,以至于它们可以检测到这些无法访问的代码并删除其中的大部分。 setInterval 显然不是这里问题的原因。 forwhile 循环会导致此问题。

为什么内存被占满了?

主要原因是创建的张量太多

const tfimage = tfnode.node.decodeJpeg(imageBuffer)

张量是不可变的。每当有新的分配时,就会创建一个新的张量。因此,内存随着张量数量的增加而增长。内置 js 对象的工作方式不同。尽管以下代码需要一些内存(用于创建数组),但不会导致内存泄漏。因为 GC 知道对于 while 循环的每次迭代,前一个变量 v 将不再使用,因此将被垃圾收集。

while(true) {
    const v = Array.from({length: 1000000}, k => k+1)
}

如何解决问题

创建的每个张量都需要在 while 块的末尾显式处理。 tf.dispose 将有助于处理一个张量。

setInterval(() => {
    const imageBuffer = fs.readFileSync(path.join(__dirname, 'temp.png'))

    const tfimage = tfnode.node.decodeJpeg(imageBuffer)
    tfnode.dispose(tfimage)

    console.log(`${++counter} - ${process.memoryUsage().rss / 1024 / 1024} MB`)
}, 100)

如果需要处理很多张量,可以在tf.tidy

的一个块中使用它们
setInterval(() => {
    tfnode.tidy(() => {
        const imageBuffer = fs.readFileSync(path.join(__dirname, 'temp.png'))

        const tfimage = tfnode.node.decodeJpeg(imageBuffer)
        console.log(`${++counter} - ${process.memoryUsage().rss / 1024 / 1024} MB`)
    })
}, 100)