Clearing cache of traced/jitted functions

TensorFlow General Discussion
,
1

Hi all,

Is there a public way to clear/invalidate the traced functions? In order to a) retrace it (with a changed Python state) and b) clear up the memory

The use-case is that a function is sometimes (Python) state based and therefore one would need to retrace it, basically invalidating the old trace and removing it.

I am aware that there are often ways around such a design, but sometimes there are None: if you build a model that depends on Python variable and run this a couple of times. Manually keeping track of changes is one thing (this is up to the user, clearly) but we cannot simply check for re-tracing in general (since this would mean that we would also need to handle the re-tracing logic of Tensors).

How could we go about clearing the traced functions?

(I managed to “hack” around and “get something working” but accessing the private interface, so that’s probably not the way forward)

3

Hi @jonas-eschle,

Sorry for the delay in response.
To clear or invalidate traced functions in TensorFlow, you can use the @tf.function decorator with input_signature to manage retracing and you can create new function instances or call get_concrete_function() with updated inputs for clearing all retraced function. In addition, manually delete unnecessary variables with del, use gc.collect() to trigger garbage collection, and call tf.keras.backend.clear_session() to free up resources in Keras. Kindly refer this documentation about controlling retracing for more information.

Hope this helps.Thank You.

1 Like