分配器内存不足 - 如何从 TensorFlow 数据集中清除 GPU 内存？

Question

假设有一个形状为

X_train

的 Numpy 数组

(4559552, 13, 22)

，代码如下：

train_dataset = tf.data.Dataset \
    .from_tensor_slices((X_train, y_train)) \
    .shuffle(buffer_size=len(X_train) // 10) \
    .batch(batch_size)

一次就可以正常工作。当我重新运行它时（对

X_train

进行轻微修改后），它会由于 GPU 内存不足而触发

InternalError

：

2021-12-19 15:36:58.460497: W tensorflow/core/common_runtime/bfc_allocator.cc:457]
Allocator (GPU_0_bfc) ran out of memory trying to allocate 9.71GiB requested by op _EagerConst

似乎第一次，它发现 100% 可用的 GPU 内存，因此一切正常，但随后的时间，GPU 内存已经几乎满了，因此出现错误。

据我了解，似乎简单地从旧的

train_dataset

中清除GPU内存就足以解决问题，但我在TensorFlow中找不到任何方法来实现这一点。目前，重新分配数据集的唯一方法是终止 Python 内核并从头开始重新运行所有内容。

有没有办法避免从头开始重新启动Python内核，而是释放GPU内存以便可以将新数据集加载到其中？

数据集不需要完整的 GPU 内存，因此我会考虑切换到 TFRecord 解决方案，作为这里的非理想解决方案（因为它会带来额外的复杂性）。

Answer 1

尝试对总 GPU 内存设置硬限制，如此处

所示

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpus[0], True)

Answer 2

执行垃圾收集似乎可以在出现问题时解决问题，而无需重新启动内核。

import gc
gc.collect()