OOM 在第二轮交叉验证

Question

我需要什么帮助/我想知道什么 我正在使用 keras API 执行交叉验证，并将执行一轮 CV 的所有代码放入单个函数中。第一轮 CV 成功，但在第二轮，我在尝试构建下一个模型时遇到 OOM 错误。

为什么会这样？
如何通过单个 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}")

到目前为止我尝试过的 我在 CV 回合开始时清理 keras 后端，我还在每回合创建一个新的分发策略范围。我已经尝试过 [1,2,4,8] 的批量大小。对于所有批次大小，它都可以正常进行一轮，但在下一轮开始时会出现 OOM。

如果……就好了 如果可以访问对内存管理的较低级别控制，那就太好了。这可能是复杂的层次。就像，最简单的情况是释放与特定图形相关的所有设备内存的函数。在 TF1 中，我会在每轮 CV 中创建一个新会话，这不是问题。

环境信息 （如果适用）

操作系统：ubuntu18.04
Python版本：3.8
Docker: tensorflow/tensorflow:2.3.1-gpu

Answer 1

答案是朋友发现的。如果在 run_fold_training 函数之外创建了对图 ops/variables 的引用，那么 clear_session 将无法完全工作。解决方案是确保在 clear_session 之后创建整个新图。例如。不要重复使用优化器等

OOM 在第二轮交叉验证

OOM in second round of cross-validation

gpu

out-of-memory

cross-validation

keras

tensorflow