Hub: Potential memory leak loading hub modules

[edugp]

It looks like loading a module multiple times leaks memory:

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text
import gc
import os
import psutil


process = psutil.Process(os.getpid())
memory = [process.memory_info().rss / (1024.0 ** 3)]
for i in range(10):
    embedder = hub.load('https://tfhub.dev/google/universal-sentence-encoder-large/5')
    del embedder
    gc.collect()
    memory.append(process.memory_info().rss / (1024.0 ** 3))

I am using:

tensorflow-gpu==2.0.0
tensorflow-hub==0.7.0
tensorflow-text==2.0.1

This is also an issue with tf.keras models using hub layers. Here is an example of a tf.keras model using a hub layer leaking memory:

def embedder_model():
    text = tf.keras.layers.Input(shape=[], dtype=tf.string)
    embedding_layer = hub.KerasLayer(MODULE_URL, input_shape=[], dtype=tf.string, trainable=True)
    embeddings = embedding_layer(text)
    return tf.keras.models.Model(inputs=text, outputs=embeddings)

process = psutil.Process(os.getpid())
memory = [process.memory_info().rss / (1024.0 ** 3)]
for i in range(10):
    model = embedder_model()
    del model
    gc.collect()
    memory.append(process.memory_info().rss / (1024.0 ** 3))

In comparison, a tf.keras model that does not use hub, only holds a significant amount of memory on the first load:

def dummy_model():
    inputs = tf.keras.layers.Input(shape=(10,))
    dense_layer = tf.keras.layers.Dense(10000)
    outputs = dense_layer(inputs)
    model = tf.keras.Model(inputs=inputs, outputs=outputs)

process = psutil.Process(os.getpid())
memory = [process.memory_info().rss / (1024.0 ** 3)]
for i in range(10):
    model = dummy_model()
    del model
    gc.collect()
    memory.append(process.memory_info().rss / (1024.0 ** 3))

[edugp]

It looks like this issue is partially fixed in 2.1.0, using tf.keras.backend.clear_session():

memory = [process.memory_info().rss / (1024.0 ** 3)]
for i in range(10):
    embedder = hub.load('https://tfhub.dev/google/universal-sentence-encoder-large/5')
    tf.keras.backend.clear_session()
    del embedder
    gc.collect()
    memory.append(process.memory_info().rss / (1024.0 ** 3))

It still leaks memory, but way less.
Also, tf.keras.backend.clear_session() is not ideal. In my case, I have multiple models loaded in memory at once, so that would remove all of them.
Does anyone understand the underlying issue here?
And also, does anyone know if there is any way to only clear a part of the graph rather than the whole thing with tf.keras.backend.clear_session()?

[arnoegw]

Thank you, @edugp, for your updated report!

hub.load() is really just a wrapper to call tf.saved_model.load() after downloading the SavedModel, so the core issue is with TensorFlow, not TF Hub, and your reports have been forwarded accordingly.