CODEWITHSUNDEEP
pythontensorflow
Maybe
Maybe

Reputation: 2279

Clear the graph and free the GPU memory in Tensorflow 2

I'm training multiple models sequentially, which will be memory-consuming if I keep all models without any cleanup. However, I am not aware of any way to the graph and free the GPU memory in Tensorflow 2.x.

What I've tried but not working

tf.keras.backend.clear_session does not work in my case as I've defined some custom layers

tf.compat.v1.reset_default_graph does not work either.

Upvotes: 1

Views: 5177

Answers (1)

user11530462
user11530462

Reputation:

Few workarounds to avoid the memory growth. Use either one

1.

del model 

tf.keras.backend.clear_session()

gc.collect()

  1. Enable allow_growth (e.g. by adding TF_FORCE_GPU_ALLOW_GROWTH=true to the environment). This will prevent TF from allocating all of the GPU memory on first use, and instead "grow" its memory footprint over time.

  2. Enable the new CUDA malloc async allocator by adding TF_GPU_ALLOCATOR=cuda_malloc_async to the environment.

Upvotes: 1

Related Questions