Creating Models In A Loop Makes Keras Increasingly Slower

when I train a model multiple times, the training iterations slow down, even if all the relevant quantities are created inside a for loop (and should therefore be overwritten each

Solution 1:

You might want to try tf.keras.backend.clear_session(), as covered in the documentation, where it explains what happens to memory consumption when models are created in a loop:

If you are creating many models in a loop, this global state will consume an increasing amount of memory over time, and you may want to clear it. Calling clear_session() releases the global state: this helps avoid clutter from old models and layers, especially when memory is limited.

Without clear_session(), each iteration of this loop will slightly increase the size of the global state managed by Keras

for _ in range(100):
  model = tf.keras.Sequential([tf.keras.layers.Dense(10) for _ in range(10)])

With clear_session() called at the beginning, Keras starts with a blank state at each iteration and memory consumption is constant over time.

for _ in range(100):
  tf.keras.backend.clear_session()
  model = tf.keras.Sequential([tf.keras.layers.Dense(10) for _ in range(10)])

Solution 2:

I don't have an exact reason what's the main cause of it. It looks like they're some kind of memory leak probably. In your code, I ran it with a bit of modification. Try this:

import numpy as np
import tensorflow as tf
import time, gc

tf.config.run_functions_eagerly(False)

n_samples = 300000
n_features = 100
n_targets = 5
batch_size = 100
x = np.array(range(n_samples * n_features), dtype=np.float64).reshape((n_samples, n_features))
y = np.array(range(n_samples * n_targets), dtype=np.float64).reshape((n_samples, n_targets))

defget_model():
    inputs = tf.keras.Input(shape=(n_features,), name='input')
    outputs = tf.keras.layers.Dense(n_features, name='dense_1', activation=tf.keras.activations.relu)(inputs)
    outputs = tf.keras.layers.Dense(n_features, name='dense_2', activation=tf.keras.activations.relu)(outputs)
    outputs = tf.keras.layers.Dense(n_features, name='dense_3', activation=tf.keras.activations.relu)(outputs)
    outputs = tf.keras.layers.Dense(n_features, name='dense_4', activation=tf.keras.activations.relu)(outputs)
    outputs = tf.keras.layers.Dense(n_targets, name='output', activation=tf.keras.activations.linear)(outputs)
    model = tf.keras.Model(inputs=inputs, outputs=outputs)
    return model

    
for t_idx inrange(7):
    tf.keras.backend.clear_session()
    gc.collect()
    
    dataset = [x, y]
    dataset = tf.data.Dataset.from_tensor_slices(tuple(dataset)).shuffle(n_samples).repeat().batch(batch_size=batch_size).prefetch(tf.data.experimental.AUTOTUNE)
    data_iterator = iter(dataset)
    
    model = get_model()
    trainable_variables = list(model.trainable_variables)
    adam_opt = tf.optimizers.Adam(learning_rate=0.001)
    
    @tf.functiondefloss(batch):
        x_, y_ = batch
        y_pred_ = model(x_)
        return tf.keras.losses.MSE(y_pred_, y_)

    @tf.functiondefoptimization_step():
        batch = next(data_iterator)
        deff(): return loss(batch)
        adam_opt.minimize(f, var_list=trainable_variables)

    iterations = 50000
    loop_start = time.time()
    
    for idx inrange(iterations):
        optimization_step()

    print(f'Elapsed: {time.time() - loop_start}')
    
    del model, dataset

When using tf.config.run_functions_eagerly(False), it gave as follows

Elapsed: 123.39986753463745Elapsed: 121.13122296333313Elapsed: 118.44610977172852Elapsed: 116.83040761947632Elapsed: 118.46350479125977Elapsed: 118.59502696990967Elapsed: 120.34098505973816

And without it gave as follows

Elapsed: 114.46430349349976Elapsed: 124.23725700378418Elapsed: 126.14825916290283Elapsed: 127.5985496044159Elapsed: 126.79593586921692Elapsed: 124.1206603050232Elapsed: 125.85739850997925

Though I don't think that is the main issue. There can be something more that needs a deep look.

[My OS: Windows 10, TF: 2.4.1, RTX:2070.]