2

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?

Zarkonnen
  • 22,200
  • 14
  • 65
  • 81
bellisk
  • 149
  • 1
  • 5
  • Your batch size seems too big to b handled by this amount of gpu memory. Try decreasing the batch size. – asymptote May 22 '22 at 19:24
  • @asymptote Thanks for the suggestion. I've tried adjusting the batch size (up to 32 and down to 1) and it unfortunately hasn't changed the GPU memory usage at all. – bellisk May 22 '22 at 19:41

0 Answers0