node.js - 几次 decodeJpeg 调用后内存不足

Question

我想使用 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 中的 arm 库和 https://github.com/yhwang/node-red-contrib-tf-model , 两者的行为相同。

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

但是我可以用 RPi 做什么？我可以控制 Tensorflow 的内存管理吗？

Answer 1

是 setInterval 的罪魁祸首吗？

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

为什么内存满了？

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

const tfimage = tfnode.node.decodeJpeg(imageBuffer)

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

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

如何解决问题

创建的每个张量都需要在 while block 的末尾显式处理。 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

的 block 中使用它们

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)