How to avoid excessive memory usage while training multiple models in Tensorflow

Question

I'm currently writing a piece of code that aims to explain how applying varying dropout rates influence the performance of a generic CNN model across several datasets.

I've set it up so, for each dataset, I train 10 different models (with 10 distinct dropout rates) a total of 3 times and register the accuracy for each run and dropout. Hopefully this dataframe explains what I've said a little bit better:

The code looks something like this:

for i, dataset in tqdm(enumerate(datasets)):
    dataset_path = pathlib.Path(args.input_folder) / dataset

    ds_train, ds_test, ds_validation = loader.get_image_data_generators(dataset_path, BATCH_SIZE)
    CLASS_NAMES = list(ds_train.class_indices.keys())

    INPUT_SHAPE = ds_train.next()[0].shape[1:]
    OUTPUT_SIZE = ds_train.next()[1].shape[1]

    ds_train_size, ds_test_size, ds_validation_size = loader.get_split_sizes(dataset_path)

    performance = {'dataset': [], 'dropout_rate': []}

    # Pre-fill dictionary with dataset and dropout labels.
    performance['dataset'] = [dataset for i in range(DROPOUT_STEPS)]
    performance['dropout_rate'] = [i/DROPOUT_STEPS for i in range(DROPOUT_STEPS)]

    for run_i in range(RUNS):
        performance[f'run_{run_i}_acc'] = []

        for i in range(DROPOUT_STEPS):

            # Compute dropout rate.
            dropout_rate = i / DROPOUT_STEPS

            # Initialize model
            model = models.Sequential()
            model.add(layers.Conv2D(32, (3, 3), activation='relu', input_shape=INPUT_SHAPE))

            model.add(layers.MaxPooling2D((2, 2)))
            model.add(layers.Conv2D(64, (3, 3), activation='relu'))

            model.add(layers.MaxPooling2D((2, 2)))
            model.add(layers.Conv2D(64, (3, 3), activation='relu'))

            model.add(layers.Flatten())
            model.add(layers.Dense(64, activation='relu'))
            model.add(layers.Dropout(dropout_rate))
            model.add(layers.Dense(OUTPUT_SIZE, activation='softmax'))

            model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

            model.fit(
                x=ds_train,
                epochs=EPOCHS, 
                steps_per_epoch=math.ceil(ds_train_size / BATCH_SIZE), 
                validation_data=ds_validation, 
                validation_steps=math.ceil(ds_validation_size / BATCH_SIZE)
            )

            test_loss, test_accuracy = model.evaluate(ds_test, steps=math.ceil(ds_test_size / BATCH_SIZE))
            performance[f'run_{run_i}_acc'].append(test_accuracy)
            print(f'✔️  Dropout rate {dropout_rate} resulted on {test_accuracy}')

    df = pd.DataFrame(performance)
    print(df)

    df.to_pickle(f'output/performance/{dataset}-perf.pkl')

In some (smaller) datasets, this runs smoothly. In larger datasets, my computer's memory usage slowly rises and in some occasions, halts the whole process by the second run, complaining there's not enough free memory.

How would I go about optimizing this code, avoiding the excessive memory usage? Is Tensorflow saving any temporary files while iterating between runs or even dropout steps? If so, how do I reset the memory on each loop cycle?

Thank you for your help.

Answer 1

Clear memory with tf.keras.backend.clear_session() after each model trains. Keras documentation states the following:

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.