Out of memory after a couple of decodeJpeg calls

Viewed 357

I want to process a couple of images (actually rather a lot) with Tensorflow using tfjs-node.

Unfortunately, on my Raspberry Pi I quickly experience an out-of-memory error. This is a snippet in 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)

This outputs something like

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)

I tried the arm libraries from https://github.com/Qengineering/TensorFlow-Raspberry-Pi and https://github.com/yhwang/node-red-contrib-tf-model, both behave the same.

Then I tried this very snippet test on Windows, and I got very astonishing results. Memory consumption raises quickly up to 8GB and then it fluctuates unreproducibly between about 800MB and 8GB. Of course it works since the system has a lot more memory than the RPi.

But what can I do with the RPi? Can I control memory management of Tensorflow?

2 Answers

Is setInterval to blame ?

There can be memory issue related to how setInterval is used but generally it is related to setInterval callback using a closure keeping a reference to certain objects. Having said that, the garbage collector GC in javascript engines have improved so well that they can detect these unreachable codes and remove most of them. setInterval is clearly not the cause of the issue here. A for or while loop would have caused the issue.

Why is the memory filled up ?

The main reason is because too many tensors are created.

const tfimage = tfnode.node.decodeJpeg(imageBuffer)

Tensors are immutable. Whenever there is a new assignement, a new tensor is created. Consequently the memory grows as the number of tensors increases. Things work differently with in-build js objects. The following code though takes some memory - for creating the array - will not cause a memory leak. Because the GC knows that for each iteration of the while loop the previous variable v will no longer be used and thus garbage collected.

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

How to solve the issue

Each tensor created needs to be explicitly disposed at the end of the while block. tf.dispose will help disposing one tensor.

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)

In case many tensors need to be disposed, they can be used in a block of 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)

Actually I have not worked with TensorFlow before but I have an idea about the setInterval part.

I think that when the setInterval loop executes the new operation, the old operation which was executed before isn't completed and might create too many stacked uncompleted operations with plenty of memory that is already being used by old and still active operations.

So you may want to do that in a sequential way like using a traditional while or for loop to not block the event loop like the setImmediate method.

I am not sure about my solution but give it a chance.

Related