Memory leak with tf.py_function in eager mode using TF 2.0

[workingloong]

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): MacOS 10.14.6
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.0.0
• Python version: 3.6.3 64bit
• Bazel version (if compiling from source): N/A
• GCC/Compiler version (if compiling from source): N/A
• CUDA/cuDNN version: N/A
• GPU model and memory: N/A

Describe the current behavior
I customize a Keras layer and use tf.py_function in the call of the layer. The used memory keeps increasing linearly with the number of calling the layer with inputs in eager mode. What's

Describe the expected behavior
The memory should keep stable.

Code to reproduce the issue

import psutil
import numpy as np
import tensorflow as tf


class TestLayer(tf.keras.layers.Layer):
    def __init__(
        self,
        output_dim,
        **kwargs
    ):
        super(TestLayer, self).__init__(**kwargs)
        self.output_dim = output_dim

    def build(self, input_shape):
        self.built = True

    def call(self, inputs):
        batch_embedding = tf.py_function(
            self.mock_output, inp=[inputs], Tout=tf.float64,
        )
        return batch_embedding

    def mock_output(self, inputs):
        shape = inputs.shape.as_list()
        batch_size = shape[0]
        return tf.constant(np.random.random((batch_size,self.output_dim)))


test_layer = TestLayer(1000)

for i in range(1000):
    test_layer.call(tf.constant(np.random.randint(0,100,(256,10))))
    if i % 100 == 0:
        used_mem = psutil.virtual_memory().used
        print('used memory: {} Mb'.format(used_mem / 1024 / 1024))

other info
We encounter this issue when developing a customized embedding layer in ElasticDL and resolving the ElasticDL issue 1567

[skydoorkai]

Found some issues in _internal_py_func in script_ops.py

Every time the py_function is called in the eager mode, a new EagerFunc is created, and a new token is added into _py_funcs.
Also, this func is appended into graph._py_funcs_used_in_graph.

Looks like this logic is for graph mode, and has a side effect in the eager mode.

_py_funcs and graph._py_funcs_used_in_graph will continously increasing in the eager mode.

This only accounts for a very small portion of the memory leak in the above sample code.
There must be other reasons for the huge memory leak.

[ravikyram]

@workingloong
I have tried in colab with TF 2.0, 2.1.0-rc0. With eager execution i can see huge memory leak (gist attached).By disabling eager execution i am seeing very small portion of memory leak (gist attached).Please, confirm is this the expected behavior?.Thanks!

[workingloong]

@workingloong
I have tried in colab with TF 2.0, 2.1.0-rc0. With eager execution i can see huge memory leak (gist attached).By disabling eager execution i am seeing very small portion of memory leak (gist attached).Please, confirm is this the expected behavior?.Thanks!

Yes, there is a very portion of memory leak in graph mode by using @tf.function to decorate the call in the customized layer. However, i can not convert the inputs to a numpy array and use python to process the array because the inputs is a Tensor not an EagerTensor in graph mode. So i hope i can use the tf.py_functionwith eager execution.

@tf.function
def call(self, inputs):
    batch_embedding = tf.py_function(
        self.mock_output, inp=[inputs], Tout=tf.float64,
    )
    return batch_embedding

[loretoparisi]

@skydoorkai we have found a similar issue when using

results = tf.py_function(
                    func=self.safe_load,
                    inp=[audio_descriptor, offset, duration, sample_rate, dtype],
                    Tout=(tf.float32, tf.bool)),
                waveform, error = results[0]

where self._safe_load loads a wave file. This caused a leak of ~24MB (the wave file size) at every prediction, afterwards calling

tf.reset_default_graph()
tf.keras.backend.clear_session()

we only see a leak of ~600-700KB at every call.

@ravikyram thanks we have tried tf.compat.v1.disable_eager_execution() but did not work.

[TheCrazyT]

I had a similar problem on kaggle.
TLDR: try tf_nightly-2.3.0.dev20200520 (or newer)

I'm using TFRecordDataset and run_eagerly on compile (sadly I am currently forced to use eargerly because of another problem that happens in a custom loss function).
The memory usage gets increased after every epoch, so i can not execute the amount of epochs i want because the kernel dies because of "out of memory".
(Raising slowly from 6 GB to over 16 GB with each epoch)
I am not shure if it is a problem about tensorflow or keras.

Garbage collection, clear_session, event setting every variable to None etc. did not help.
Somehow the memory was still used.

Somehow it seems to be better with tf_nightly-2.3.0.dev20200520.
So maybe the bug was fixed?