How to fix memory leak issue in Keras using K.clear_session()?

user_4512432 0
Viewed 1246

I have a network which I'm training by feeding batches of my data, and I'm using model.train_on_batch() to do this. If I run only this training portion, I see my network trains just fine at 3% RAM utilisation for 40+ epochs (as yet) and each epoch has about 2000 iterations. When I try to do validation after each epoch (which also happens in batches), there is a very bad memory leak leading to 90% RAM utilisation and my code getting killed. So I've tried a couple of things in the past few days, it seems model.predict() in a loop causes memory leak open issue at Tensorflow GitHub. I tried predict_on_batch(), same behaviour. model(inputs, training=False) seemed slow down the memory leak, instead of the abrupt jump from 3% - 7% - 13% - 40% - 80% - 90% (intervals of 60sec), it increased as 1% per minute. But at a point it reached 90% too. The only thing left for me to try out of this github thread is using K.clear_session().

I tried reading up documentation for K.clear_session() and some SO posts, all suggest using it when you are creating multiple models, which I am not doing. So my question is if I have a single model being trained and evaluated in loops, where should I use K.clear_session(), after each epoch and reload saved model before each epoch? Will that be correct?

Apart from this, I also get the topological sort error another open issue, so I wonder if it's because I'm training in loops because my code otherwise has no loops and this is somehow causing the memory leak too, and K.clear_session() would help somehow?

Minimal example of my code structure:

from tensorflow.keras.models import Model
K = tf.keras.backend

def myModel():
    **some architecture**

ip = Input(shape=(h, w, 3))
op = myModel(ip)
model = Model(ip, op)
model.compile(optimizer=Adam(lr=1e-6), loss=custom_mean_squared_error)

for e in range(numEpochs):
    for batch in range(0, num_train_batches):
        x = readImages()
        y = readLabels()
        loss = model.train_on_batch(x, y)
        

    for batch in range(0, num_val_batches):
        x = readImages()
        y = model.predict(x)
        val_loss = K.get_value(custom_mean_squared_error(x,y))
        # save predictions

# plot training vs validation loss

Tensorflow-gpu-1.14, python3.6. Would appreciate suggestions if I'm doing something wrong too.

keras tensorflow python
1 Answers

This seems to work for me, though makes the process slower:

from tensorflow.keras.models import Model
K = tf.keras.backend

def myModel():
    **some architecture**

ip = Input(shape=(h, w, 3))
op = myModel(ip)
model = Model(ip, op)
model.compile(optimizer=Adam(lr=1e-6), loss=custom_mse)

for e in range(numEpochs):
    for batch in range(0, num_train_batches):
        x = readImages()
        y = readLabels()
        # define appropriate flags for first loop
        model = tf.keras.models.load_model(model_path,custom_objects={ 'custom_mse': custom_mse } )
        loss = model.train_on_batch(x, y)

    model.save(model_path)

    for batch in range(0, num_val_batches):
        x = readImages()
        model = tf.keras.models.load_model(model_path,custom_objects={ 'custom_mse': custom_mse } )
        y = model.predict(x)
        K.clear_session()
        val_loss = K.get_value(custom_mean_squared_error(x,y))
        # save predictions

# plot training vs validation loss
user_45124320
Related