python - 使用 tensorflow 的数据集管道，我如何*命名* `map` 操作的结果？

Question

我有下面的 map 函数(可运行示例)，它输入一个 string 并输出一个 string 和一个 integer。

在 tf.data.Dataset.from_tensor_slices 中，我将原始输入命名为 'filenames'。但是当我从 map 函数 map_element_counts 返回值时，我只能返回一个元组(返回字典会产生异常)。

有没有办法命名从我的 map_element_counts 函数返回的 2 个元素？

import tensorflow as tf

filelist = ['fileA_6', 'fileB_10', 'fileC_7']

def map_element_counts(fname):
  # perform operations outside of tensorflow
  return 'test', 10

ds = tf.data.Dataset.from_tensor_slices({'filenames': filelist})
ds = ds.map(map_func=lambda x: tf.py_func(
  func=map_element_counts, inp=[x['filenames']], Tout=[tf.string, tf.int64]
))
element = ds.make_one_shot_iterator().get_next()

with tf.Session() as sess:
  print(sess.run(element))

结果:

(b'test', 10)

期望的结果:

{'elementA': b'test', 'elementB': 10)

---

添加的细节:

当我执行 return {'elementA': 'test', 'elementB': 10} 时，我得到了这个异常:

tensorflow.python.framework.errors_impl.UnimplementedError: Unsupported object type dict

Answer 1

为了后代，我提出了这个问题的最终解决方案。下面的代码是一个复制/粘贴示例，可在该问题解决的最复杂条件下工作(请注意，其他两个答案不是可复制/粘贴的代码示例):

代码的目标是:

• 获取(大)文件列表并将其分成 block (文件名/索引对)
• 使用映射操作处理每个 block (生成器在这里不是可行的解决方案，请参阅:https://github.com/tensorflow/tensorflow/issues/16343)
• 从只需要 1 个文件/ block 作为输入的映射操作中输出多个样本。
• 在整个过程中维护元素命名

Tensorflow 1.5/Python 3.x 的复制/粘贴工作示例

import tensorflow as tf
import numpy as np

files = [b'testA', b'testB', b'testC']

def mymap1(x):
  result_tensors = tf.py_func(func=mymap2, inp=[x], Tout=[tf.string, tf.int64])
  return {'filename': result_tensors[0], 'value': result_tensors[1]}

def mymap2(x):
  return np.array([x, x, x]), np.array([10, 20, 30])

def myflatmap(named_elements):
  return tf.data.Dataset.zip({
    'filename': tf.data.Dataset.from_tensor_slices(named_elements['filename']),
    'value': tf.data.Dataset.from_tensor_slices(named_elements['value'])
  })

ds = tf.data.Dataset.from_tensor_slices(files)
ds = ds.map(map_func=mymap1)
ds = ds.flat_map(map_func=myflatmap)

element = ds.make_one_shot_iterator().get_next()

with tf.Session() as sess:
  for _ in range(9):
    print(sess.run(element))

输出:

{'filename': b'testA', 'value': 10}
{'filename': b'testA', 'value': 20}
{'filename': b'testA', 'value': 30}
{'filename': b'testB', 'value': 10}
{'filename': b'testB', 'value': 20}
{'filename': b'testB', 'value': 30}
{'filename': b'testC', 'value': 10}
{'filename': b'testC', 'value': 20}
{'filename': b'testC', 'value': 30}