Excessive GPU memory usage for Keras model.fit()

I am trying to use Keras to do some image analysis using the Xception model. When I run model.fit, GPU memory usage increases rapidly to 16 GB, despite my training data only being (7546, 299, 299, 3) in size, so about 1.9 GB.

I have tried using K.clear_session() between epochs, which helps a bit, but GPU memory usage still quickly reaches the maximum of 16 GB and stays there.

I am running this on a Jupyter notebook on Google Colab Pro+ using TensorFlow and Keras 2.8.0. Here is the code I use to create and run the network:

class EvalCallback(Callback):
    def on_epoch_end(self, epoch, logs=None):
        gc.collect()
        K.clear_session()

Xception_initial=Xception(include_top=False,
                 weights='imagenet',
                 input_shape=(299,299,3),pooling ='avg',
                 )

for layer in Xception_initial.layers:
    layer.trainable = True

x = Xception_initial.output
eval_callback = EvalCallback()
predicted = Dense(2,activation ='softmax')(x)
model_pretrain = Model(inputs = Xception_initial.input, outputs = predicted)
model_pretrain.compile(loss=keras.losses.categorical_crossentropy,
          optimizer=tf.keras.optimizers.Adam(lr = 0.0002),
          metrics=['accuracy'])
pretraining_Xception =model_pretrain.fit(x_train, Y_train,
                                              verbose=1,
                                              batch_size=16,  
                                              epochs=3,
                                              callbacks=[eval_callback])

Running this code does complete, but GPU memory stays maxed out, and I am unable to do any additional work using Keras, getting a dst tensor not initialized error which I believe is caused by the GPU being out of memory.

What can I do to fix this memory problem?

Topic keras tensorflow python

Category Data Science

About

Geeks Mental is a community that publishes articles and tutorials about Web, Android, Data Science, new techniques and Linux security.