python - 加载 tfrecord 文件时出现 InvalidArgumentError

Question

我开始加载和保存 tfrecord 文件，以编写输入函数。我已经设置了以下测试，但收到 InvalidArgumentError。我已经使用 save() 方法保存了 tfrecord 文件，并尝试使用 load() 方法将其加载回来。

import tensorflow as tf

def save(filename):
    writer = tf.python_io.TFRecordWriter(filename)
    for i in range(0,10):
        features = {
            'x': tf.train.Feature(float_list=tf.train.FloatList(value=[1.0])),
            'y': tf.train.Feature(int64_list=tf.train.Int64List(value=[i]))
        }
        example = tf.train.Example(features=tf.train.Features(feature=features))
        writer.write(example.SerializeToString())
    writer.close()

def load(filename_pattern):
    def parse(single_file):
        features = {
            'x': tf.FixedLenFeature([1], dtype=tf.float32),
            'y': tf.FixedLenFeature([1], dtype=tf.int64)
        }
        parsed_features = tf.parse_single_example(single_file, features)
        return parsed_features
    dataset = tf.data.TFRecordDataset.list_files(filename_pattern)
    dataset = dataset.map(parse)
    return dataset

def main():

    with tf.Session() as sess:
        dataset = load('test.tfrecords')
        iterator = dataset.make_one_shot_iterator()
        next_element = iterator.get_next()

        for i in range(5):
            value = sess.run(next_element)
            print(value)


if __name__ == '__main__':
    main()

完整的错误是:

2018-02-21 11:05:12.311589: I tensorflow/core/platform/cpu_feature_guard.cc:137] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.2 AVX
2018-02-21 11:05:12.378303: W tensorflow/core/framework/op_kernel.cc:1198] Invalid argument: Could not parse example input, value: './test.tfrecords'
2018-02-21 11:05:12.378375: W tensorflow/core/framework/op_kernel.cc:1198] Invalid argument: Could not parse example input, value: './test.tfrecords'
     [[Node: ParseSingleExample/ParseExample/ParseExample = ParseExample[Ndense=2, Nsparse=0, Tdense=[DT_FLOAT, DT_INT64], dense_shapes=[[1], [1]], sparse_types=[]](ParseSingleExample/ExpandDims, ParseSingleExample/ParseExample/ParseExample/names, ParseSingleExample/ParseExample/ParseExample/dense_keys_0, ParseSingleExample/ParseExample/ParseExample/dense_keys_1, ParseSingleExample/ParseExample/Const, ParseSingleExample/ParseExample/Const_1)]]
Traceback (most recent call last):
  File "tf_records_save.py", line 39, in <module>
    main()
  File "tf_records_save.py", line 34, in main
    value = sess.run(next_element)
  File "/Users/bm14368/Desktop/Project/tensorflow/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 895, in run
    run_metadata_ptr)
  File "/Users/bm14368/Desktop/Project/tensorflow/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1128, in _run
    feed_dict_tensor, options, run_metadata)
  File "/Users/bm14368/Desktop/Project/tensorflow/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1344, in _do_run
    options, run_metadata)
  File "/Users/bm14368/Desktop/Project/tensorflow/lib/python2.7/site-packages/tensorflow/python/client/session.py", line 1363, in _do_call
    raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Could not parse example input, value: './test.tfrecords'
     [[Node: ParseSingleExample/ParseExample/ParseExample = ParseExample[Ndense=2, Nsparse=0, Tdense=[DT_FLOAT, DT_INT64], dense_shapes=[[1], [1]], sparse_types=[]](ParseSingleExample/ExpandDims, ParseSingleExample/ParseExample/ParseExample/names, ParseSingleExample/ParseExample/ParseExample/dense_keys_0, ParseSingleExample/ParseExample/ParseExample/dense_keys_1, ParseSingleExample/ParseExample/Const, ParseSingleExample/ParseExample/Const_1)]]
     [[Node: IteratorGetNext = IteratorGetNext[output_shapes=[[1], [1]], output_types=[DT_FLOAT, DT_INT64], _device="/job:localhost/replica:0/task:0/device:CPU:0"](OneShotIterator)]]

Answer 1

可以重写加载函数以加载一个文件而不是文件模式。使用构造函数tf.data.TFRecordDataset(filename)。那么它应该可以工作。看来 tf.data.TFRecordDataset.list_files(filename_pattern) 无法处理单个文件名。

def load(filename):
    def parse(single_file):
        features = {
            'x': tf.FixedLenFeature([1], dtype=tf.float32),
            'y': tf.FixedLenFeature([1], dtype=tf.int64)
        }
        parsed_features = tf.parse_single_example(single_file, features)
        return parsed_features
    dataset = tf.data.TFRecordDataset(filename)
    dataset = dataset.map(parse)
    return dataset