Memory leak for TFP distribution
fotisdr opened this issue · comments
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