几次 decodeJpeg 调用后内存不足

Question

我想使用 tfjs-node 使用 Tensorflow 处理几个图像（实际上相当多）。

不幸的是，在我的 Raspberry Pi 上，我很快就遇到了内存不足的错误。 这是 Typescript 中的 sn-p：

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 和 https://github.com/yhwang/node-red-contrib-tf-model 的 arm 库，它们的行为相同。

然后我在 Windows 上尝试了这个非常 sn-p 的测试，我得到了非常惊人的结果。内存消耗迅速增加到 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 块的末尾显式处理。 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)

Answer 2

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

我认为当setInterval 循环执行新操作时，之前执行的旧操作未完成，并且可能会创建太多堆叠的未完成操作，其中有大量内存已被旧操作且仍处于活动状态操作。

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

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