Why is creating a model taking so much space in VRAM?

Issue

I would like to know what is loaded in VRAM when creating a new model. Even in simple cases like:

import tensorflow as tf

model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

memory usage increases from 0 to 4000+ MiB the moment I run the model definition command. Only thing that comes to my mind is that various CUDA libraries are loaded in memory, but I am not sure about that. Any thoughts?

Thank you in advance!

Solution

Tensorflow has a habit of taking as much space as it can incase it needs it later on. You can try the following to allow dynamic memory growth which should fix your issue.

physical_devices = tf.config.experimental.list_physical_devices('GPU')
for i in physical_devices:
    tf.config.experimental.set_memory_growth(i, True)

Answered By – Sean

This Answer collected from stackoverflow, is licensed under cc by-sa 2.5 , cc by-sa 3.0 and cc by-sa 4.0

Leave a Reply

(*) Required, Your email will not be published