停止 TensorFlow 数据集“from_generator”的正确方法？

Question

我想使用使用from_generator 构建的 TensorFlow 数据集来访问格式化文件。大多数情况下，除了我不知道如何在生成器用完数据时停止 Dataset 迭代器（当您超出范围时，生成器永远返回空列表）之外，其他一切都有效。

我的实际代码非常复杂，但我可以用这个简短的程序模拟这种情况：

import tensorflow as tf

def make_batch_generator_fn(batch_size=10, dset_size=100):
    feats, targs = range(dset_size), range(1, dset_size + 1)

    def batch_generator_fn():
        start_idx, stop_idx = 0, batch_size
        while True:
            # if stop_idx > dset_size: --- stop action?
            yield feats[start_idx: stop_idx], targs[start_idx: stop_idx]
            start_idx, stop_idx = start_idx + batch_size, stop_idx + batch_size

    return batch_generator_fn

def test(batch_size=10):
    dgen = make_batch_generator_fn(batch_size)
    features_shape, targets_shape = [None], [None]
    ds = tf.data.Dataset.from_generator(
        dgen, (tf.int32, tf.int32),
        (tf.TensorShape(features_shape), tf.TensorShape(targets_shape))
    )
    feats, targs = ds.make_one_shot_iterator().get_next()

    with tf.Session() as sess:
        counter = 0
        try:
            while True:
                f, t = sess.run([feats, targs])
                print(f, t)
                counter += 1
                if counter > 15:
                    break
        except tf.errors.OutOfRangeError:
            print('end of dataset at counter = {}'.format(counter))

if __name__ == '__main__':
    test()

如果我提前知道记录的数量，我可以调整批次的数量，但我并不总是知道。我尝试在上面的 sn-p 中放置一些代码，其中有一个类似 stop action? 的注释行。特别是，我尝试提出IndexError，但TensorFlow 不喜欢这样，即使我在我的执行代码中明确地catch 它。我也尝试提出tf.errors.OutOfRangeError，但我不确定如何实例化它：构造函数需要三个参数 - 'node_def'、'op' 和 'message'，我不太确定要使用什么'node_def' 和 'op' 一般。

如果您有任何关于此问题的想法或意见，我将不胜感激。谢谢！

Answer 1

满足停止条件时返回：

def make_batch_generator_fn(batch_size=10, dset_size=100):
    feats, targs = range(dset_size), range(1, dset_size + 1)

    def batch_generator_fn():
        start_idx, stop_idx = 0, batch_size
        while True:
            if stop_idx > dset_size:
                return
            else:
                yield feats[start_idx: stop_idx], targs[start_idx: stop_idx]
                start_idx, stop_idx = start_idx + batch_size, stop_idx + batch_size

    return batch_generator_fn

这符合Python 3 documentation:中指定的行为

在生成器函数中，return 语句指示生成器已完成并将引发 StopIteration。返回的值（如果有）用作构造 StopIteration 的参数，并成为 StopIteration.value 属性。

Answer 2

它适用于以下几行：

dataset_size = your dataset size
batch_size = your batch size
dataset = your tf.data.Dataset
steps_per_epoch = dataset_size // batch_size

for data, _ in zip(dataset, range(steps_per_epoch)):
    # your train_step

迭代结束后将停止。