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