python - 在 TensorFlow 中读取 `tf.train.shuffle_batch` 文件时 `TFRecord` 崩溃

Question

我正在尝试使用 tf.train.shuffle_batch 来使用 TensorFlow 1.0 使用 TFRecord 文件中的批量数据。相关功能有:

def tfrecord_to_graph_ops(filenames_list):
    file_queue = tf.train.string_input_producer(filenames_list)
    reader = tf.TFRecordReader()
    _, tfrecord = reader.read(file_queue)

    tfrecord_features = tf.parse_single_example(
        tfrecord,
        features={'targets': tf.FixedLenFeature([], tf.string)}
    )
    ## if no reshaping: `ValueError: All shapes must be fully defined` in
    ## `tf.train.shuffle_batch`
    targets = tf.decode_raw(tfrecord_features['targets'], tf.uint8)
    ## if using `strided_slice`, always get the first record
    # targets = tf.cast(
    #     tf.strided_slice(targets, [0], [1]),
    #     tf.int32
    # )
    ## error on shapes being fully defined
    # targets = tf.reshape(targets, [])
    ## get us: Invalid argument: Shape mismatch in tuple component 0.
    ## Expected [1], got [1000]
    targets.set_shape([1])
    return targets


def batch_generator(filenames_list, batch_size=BATCH_SIZE):
    targets = tfrecord_to_graph_ops(filenames_list)
    targets_batch = tf.train.shuffle_batch(
        [targets],
        batch_size=batch_size,
        capacity=(20 * batch_size),
        min_after_dequeue=(2 * batch_size)
    )
    targets_batch = tf.one_hot(
        indices=targets_batch, depth=10, on_value=1, off_value=0
    )
    return targets_batch


def examine_batches(targets_batch):
    with tf.Session() as sess:
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        for _ in range(10):
            targets = sess.run([targets_batch])
            print(targets)
        coord.request_stop()
        coord.join(threads)

代码通过examine_batches()进入，并收到batch_generator()的输出。 batch_generator() 调用 tfrecord_to_graph_ops()，我相信问题出在该函数中。

我正在打电话

targets = tf.decode_raw(tfrecord_features['targets'], tf.uint8)

位于 1,000 字节(数字 0-9)的文件上。如果我在 session 中对此调用 eval() ，它会显示所有 1,000 个元素。但如果我尝试将其放入批处理生成器中，它就会崩溃。

如果我不 reshape 目标，当tf.train.shuffle_batch时，我会收到类似ValueError:所有形状必须完全定义的错误> 被调用。如果我调用 targets.set_shape([1])，让人想起 Google 的 CIFAR-10 example code ，我收到类似 Invalid argument: Shape Mismatch in tuple component 0. Expected [1], got [1000] in tf.train.shuffle_batch 的错误。我还尝试使用 tf.strided_slice 来剪切一大块原始数据 - 这不会崩溃，但会导致一遍又一遍地获取第一个事件。

执行此操作的正确方法是什么？要从 TFRecord 文件中提取批处理？

注意，我可以手动编写一个函数来切碎原始字节数据并进行某种批处理 - 如果我使用 feed_dict 方法将数据获取到图表中，这尤其容易 - 但我我正在尝试学习如何使用 TensorFlow 的 TFRecord 文件以及如何使用其内置批处理函数。

谢谢!

Answer 1

Allen Lavoie 在评论中指出了正确的解决方案。重要的缺失部分是 enqueue_many=True 作为 tf.train.shuffle_batch() 的参数。编写这些函数的正确方法是:

def tfrecord_to_graph_ops(filenames_list):
    file_queue = tf.train.string_input_producer(filenames_list)
    reader = tf.TFRecordReader()
    _, tfrecord = reader.read(file_queue)

    tfrecord_features = tf.parse_single_example(
        tfrecord,
        features={'targets': tf.FixedLenFeature([], tf.string)}
    )
    targets = tf.decode_raw(tfrecord_features['targets'], tf.uint8)
    targets = tf.reshape(targets, [-1])
    return targets

def batch_generator(filenames_list, batch_size=BATCH_SIZE):
    targets = tfrecord_to_graph_ops(filenames_list)
    targets_batch = tf.train.shuffle_batch(
        [targets],
        batch_size=batch_size,
        capacity=(20 * batch_size),
        min_after_dequeue=(2 * batch_size),
        enqueue_many=True
    )
    return targets_batch