TensorFlow 数据集导致内核在迭代期间终止进程

Posted 2023-02-15

【中文标题】TensorFlow 数据集导致内核在迭代期间终止进程

【英文标题】：TensorFlow dataset causes kernel to kill the process during iteration

【发布时间】：2022-01-09 22:52:32

【问题描述】：

我想创建一个用于训练 TensorFlow 模型的数据管道。数据存储在非常大的 HDF5 文件中 (250+ GB)。

我编写了一个适用于较小输入文件的管道，但在消耗过多 RAM+swap 后最终被内核杀死（通过监控验证了这一点）。

import tensorflow as tf
import h5py

class TestGenerator:
    """
    Implements a generator that can be used by tf.data.Dataset.from_generator 
    to produce a dataset for any test data.
    """
    def __init__(self, src, dset):
        self.src = src
        self.dset = dset
        self.output_signature = (
            tf.TensorSpec(shape=(2,), dtype=tf.uint64)
        )

    def __call__(self):
        """This is needed for tf.data.Dataset.from_generator to work."""
        with h5py.File(self.src, 'r', swmr=True) as f:
            for sample in f[self.dset]:
                yield sample[0], sample[1]

gen = TestGenerator('h5file.h5', 'dset_path')

dataset = tf.data.Dataset.from_generator(
    gen, 
    output_signature=gen.output_signature
)

for sample in dataset:
    pass

一开始我以为这可能是h5py模块的问题，所以单独测试了一下：

with h5py.File('h5file.h5', 'r', swmr=True) as f:
    for sample in f['dset_path']:
        pass

这没有问题。由此得出结论，TensorFlow 是造成内存问题的原因。让我恼火的是，我假设 TensorFlow 可以即时获取所需的数据，因此可以避免内存问题。

代码经过测试，适用于较小的文件。我还测试了在迭代之前使用dataset.prefetch 的版本，但结果相同。

TensorFlow 是否会在后台加载整个数据集？

【参考方案1】：

如果你打开任务管理器，在导入 tensorflow 和创建模型时，它会显示 GPU 内存保留的极高值吗？

如果是这样，可能不是模型或数据的大小，而是 tensorflow 为所有模型训练保留尽可能多的内存这一事实。

就我个人而言，我的 3080 有 10 GB 的专用 GPU 内存，而 tensorflow 占用了 9.7 GB。

如果是这种情况，请查看https://www.tensorflow.org/guide/gpu#limiting_gpu_memory_growth 上的 set_memory_growth 方法。

使用它可以将我的专用 GPU 内存使用量从 9.7GB 减少到 3.2GB 和 4GB 之间。

编辑：我不确定它会在什么时候分配该内存，但是如果您尝试训练模型并密切关注 GPU 的任务管理器性能，您应该会发现它是否以这种方式运行！

【讨论】：

我尝试了所有建议，但不幸的是，它没有成功。我认为这是因为 TF 数据不使用 GPU 及其内存。我使用watch -n 1 free -m 监控系统上的 RAM 内存分配，很容易看到可用空间逐渐减少，直到几乎为零。发生这种情况时，进程会被内核杀死并释放内存。

啊，这听起来更像是您试图加载太多信息？我不确定，但你能进一步定义你提到的“RAM+swap”吗？因为你的 GPU 内存需要 1) 足够的内存来运行模型和 2) 足够的内存来获取数据，也就是说，从 RAM 到 GPU 内存的副本？

感谢您的帮助，但似乎我在底层 h5py 库中遇到了一个错误。我从上面的测试代码中省略的是我在数据集上使用了切片运算符，例如f['dataset'][start:stop] 选择数据子集。切片似乎导致将所有选定的数据复制到 RAM 中。不过，我通过在可迭代对象上使用 itertools.islice 找到了一种解决方法。我将保留此主题，因为它可能对其他人有所帮助。