Tensorflow 第二轮交叉验证中的OOM

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

• 为什么会这样
• 如何从单个python进程正确地执行这种类型的CV
• 有没有办法完全刷新GPU/TPU内存来控制内存碎片之类的事情

将tensorflow导入为tf
def运行折叠训练（k折叠、数量折叠、批量大小）：
#清晰图形
tf.keras.backend.clear_session（）
#尝试获取tpu或gpu
尝试：
tpu=tf.distribute.cluster_resolver.TPUClusterResolver（）
打印（'Device:'，tpu.master（））
tf.config.experimental\u连接到集群（tpu）
tf.tpu.experimental.initialize\u\tpu\u系统（tpu）
策略=tf.distribute.experimental.TPUStrategy（tpu）
除：
strategy=tf.distribute.get_strategy（）
打印（'Number of replica:'，strategy.num\u replicas\u in\u sync）
使用strategy.scope（）：
#生成k-fold数据集
ds=构建数据集（）
列\u ds=ds.enumerate（）.filter(
lambda i，ds，num_fold=num_fold，k_fold=k_fold:i%num_fold！=k_fold）.map(
lambda i，ds:ds）。批次（批次大小）
test_ds=ds.enumerate（）.filter(
lambda i，ds，num_fold=num_fold，k_fold=k_fold:i%num_fold==k_fold）.map(
lambda i，ds:ds）。批次（批次大小）
#制作、训练、评估模型
模型=我的模型（**模型）
model.compile（**compile_-kwargs）
模型拟合（系列，历次=25）
结果=模型。评估（测试，返回值=真）
返回结果[“分数”]
num_folds=5
批量大小=8
cv_损失=总和（[对范围内的k（num_fold）进行[run_fold_training（k，num_fold，batch_size）]）/num_fold
打印（f“最终{num_folds}-折叠交叉验证分数为：{cv_loss}”）

我到目前为止所做的尝试 在CV回合开始时，我将清除keras后端，并在每回合中创建一个新的分发策略范围。我已经尝试过批量大小为[1,2,4,8]。对于所有批量大小，它只执行一轮，但在下一轮开始时提供OOM

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

环境信息 （如适用）

• 操作系统：ubuntu 18.04
• Python版本：3.8
• Docker:

  tensorflow/tensorflow:2.3.1-gpu

---

答案被一位朋友发现了。如果在

run\u fold\u training

功能之外创建了对图形操作/变量的引用，则

clear\u会话将无法完全工作。解决方案是确保在
清除\u会话之后创建整个新图形。例如，不要重复使用优化器等