如何通过tensorflow的tf.data API加载pickle文件

时间:2018-06-15 06:22:22

标签: tensorflow pickle

我的数据存储在磁盘上的多个pickle文件中。我想使用tensorflow的tf.data.Dataset将我的数据加载到训练管道中。我的代码是:

def _parse_file(path):
    image, label = *load pickle file*
    return image, label
paths = glob.glob('*.pkl')
print(len(paths))
dataset = tf.data.Dataset.from_tensor_slices(paths)
dataset = dataset.map(_parse_file)
iterator = dataset.make_one_shot_iterator()

问题是我不知道如何实现_parse_file功能。该函数的参数path具有张量类型。我试过了

def _parse_file(path):
    with tf.Session() as s:
        p = s.run(path)
        image, label = pickle.load(open(p, 'rb'))
    return image, label

并收到错误消息:

InvalidArgumentError (see above for traceback): You must feed a value for placeholder tensor 'arg0' with dtype string
     [[Node: arg0 = Placeholder[dtype=DT_STRING, shape=<unknown>, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

在互联网上进行一些搜索后,我仍然不知道该怎么做。我将非常感谢任何向我提供暗示的人。

2 个答案:

答案 0 :(得分:3)

我自己解决了这个问题。我应该使用此doc中的tf.py_func

答案 1 :(得分:-1)

tf.py_func 该函数用于解决该问题,也作为文档中的提示。