Memory leak for TFP distribution

[fotisdr]

Hi,

I'm facing a memory leak issue with TFP distributions. Specifically, I have a custom function that samples from a Categorical distribution and returns the output:

def compute_categorical_distribution(classes_pred,dtype=np.int8,validate_args=True):
  # Define the distribution
  distr = tfp.distributions.Categorical(probs=classes_pred, validate_args=validate_args)
  # Sample from the distribution
  output = distr.sample().numpy().astype(dtype)
del distr
K.clear_session()
gc.collect()
return output

I was expecting that, when I call this function, I would get the output array back and the memory allocation of the distribution would be freed up. However, this does not seem to be the case, as for every consecutive call of this function the memory usage keeps increasing, indicating that a new tfp distribution is created each time and remains in the memory. Is there a way to clear up the memory allocation of the distribution? As you can see above, I tried del distr, tensorflow.keras.backend.clear_session() and gc.collect() but none of these seems to make any difference. I am using tfp version 0.19.0.

Thanks! Fotis

[BNerot]

Hi,

I met the same problem for tfp.distributions.Normal.

In my case, the problem arises due to conflicts between keras and tensorflow in their way to manage random seeds.

I had:

tf.keras.utils.set_random_seed(42)
distr.sample(seed=None)

I would suggest for a dirty fix:

# tf.keras.utils.set_random_seed(42)
distr.sample(seed=42)

I am - apparently - running TF eagerly.

tensorflow==2.19.0
tensorflow-probability==0.25.0
keras==3.10

I don't have the time to investigate further for the moment. A good start may be to have a look at get_seed in tensorflow.python.framework.random_seed.

Hope it helps some people.