分配器用完了 memory - 如何从 Z074DD699710DA0EC1EB45F13Z31 数据集中清除 GPU memory

Question

假设一个 Numpy 数组X_train的形状为(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稍作修改后），由于 memory 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

It seems that the first time, it finds 100% free GPU memory so all works fine, but the subsequent times, the GPU memory is already almost full and hence the error.

据我了解，似乎只需从旧的 train_dataset 中清除 GPU train_dataset就足以解决问题，但我在 Z074DD699710DA0EC1EB45F13B3178E 中找不到任何方法来实现这一点目前重新分配数据集的唯一方法是杀死 Python kernel 并从头重新运行所有内容。

Is there a way to avoid re-starting the Python kernel from scratch and instead free the GPU memory so that the new dataset can be loaded into it?

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

Answer 1

尝试对总 GPU memory 设置硬限制，如下所示

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