Блог

С чем мне нужна помощь / Что мне было интересно, я выполняю перекрестную проверку с использованием keras API и поместил весь код для выполнения одного раунда CV в одну функцию. Первый раунд CV работает, но затем, после второго раунда, я получаю ошибку ООМ при попытке построить следующую модель.

• Почему это происходит?
• Как мне правильно выполнить этот тип CV из одного процесса python?
• Есть ли способ полностью очистить память GPU / TPU для управления такими вещами, как фрагментация памяти?

 import tensorflow as tf

def run_fold_training(k_fold, num_folds, batch_size):
    #clear graph
    tf.keras.backend.clear_session()
    #try to get tpu or else gpu
    try:
        tpu = tf.distribute.cluster_resolver.TPUClusterResolver()
        print('Device:', tpu.master())
        tf.config.experimental_connect_to_cluster(tpu)
        tf.tpu.experimental.initialize_tpu_system(tpu)
        strategy = tf.distribute.experimental.TPUStrategy(tpu)
    except:
        strategy = tf.distribute.get_strategy()
    print('Number of replicas:', strategy.num_replicas_in_sync)
    with strategy.scope():
        # make k-fold dataset
        ds = build_dataset()
        train_ds = ds.enumerate().filter(
            lambda i, ds, num_folds=num_folds, k_fold=k_fold: i % num_folds != k_fold).map(
            lambda i, ds: ds).batch(batch_size)
        test_ds = ds.enumerate().filter(
            lambda i, ds, num_folds=num_folds, k_fold=k_fold: i % num_folds == k_fold).map(
            lambda i, ds: ds).batch(batch_size)
        # make, train, evaluate model
        model = MyModel(**model_kwargs)
        model.compile(**compile_kwargs)
        model.fit(train_ds, epochs=25)
        results = model.evaluate(test_ds, return_dict=True)

    return results["score"]

num_folds = 5
batch_size = 8
cv_loss = sum([run_fold_training(k, num_folds, batch_size) for k in range(num_folds)]) / num_folds
print(f"Final {num_folds}-fold cross validation score is: {cv_loss}")
  

Что я пробовал до сих пор
Я очищаю серверную часть keras в начале раунда CV, а также создаю новую область стратегии распространения для каждого раунда. Я уже пробовал пакеты размером [1,2,4,8]. Для всех пакетов он выполняет один раунд нормально, но выдает ООМ в начале следующего раунда.

Было бы неплохо, если…Было бы здорово, если бы был доступ к контролю более низкого уровня над управлением памятью. Это может быть в уровнях сложности. Например, простейшим случаем была бы функция, которая освобождает всю память устройства, связанную с определенным графиком. В TF1 я бы просто создал новый сеанс для каждого раунда CV, и это не было бы проблемой.

Информация об окружающей среде (если применимо)

• Операционная система: ubuntu 18.04
• Версия Python: 3.8
• Docker: tensorflow/tensorflow:2.3.1-gpu

Ответ был обнаружен другом. Если есть ссылки на графические операции / переменные, созданные вне run_fold_training функции, то clear_session они не будут работать полностью. Решение состоит в том, чтобы убедиться, что весь новый граф создается после clear_session . Например. не используйте оптимизаторы повторно и т. Д.