Two more steps were required to achieve this. First, the map function needs to return a numpy array, not a list.
Then you can use flat_map
combined with Dataset().from_tensor_slices()
to flatten them. The code below now produces the desired result:
Tested in Tensorflow 1.5 (copy/paste runnable example)
import tensorflow as tf
import numpy as np
input = [10, 20, 30]
def my_map_func(i):
return np.array([i, i + 1, i + 2])
ds = tf.data.Dataset.from_tensor_slices(input)
ds = ds.map(map_func=lambda input: tf.py_func(
func=my_map_func, inp=[input], Tout=[tf.int64]
))
ds = ds.flat_map(lambda x: tf.data.Dataset().from_tensor_slices(x))
element = ds.make_one_shot_iterator().get_next()
with tf.Session() as sess:
for _ in range(9):
print(sess.run(element))
Here is a method of doing this if you have multiple variables to return, in this example I input a string (such as a filename) and output multiples of both strings and integers. In this case I repeat the string for each of the integers of [10, 20, 30].
Copy/paste runnable example:
import tensorflow as tf
import numpy as np
input = [b'testA', b'testB', b'testC']
def my_map_func(input):
return np.array([input, input, input]), np.array([10, 20, 30])
ds = tf.data.Dataset.from_tensor_slices(input)
ds = ds.map(map_func=lambda input: tf.py_func(
func=my_map_func, inp=[input], Tout=[tf.string, tf.int64]))
ds = ds.flat_map(lambda mystr, myint: tf.data.Dataset().zip((
tf.data.Dataset().from_tensor_slices(mystr),
tf.data.Dataset().from_tensor_slices(myint))
))
element = ds.make_one_shot_iterator().get_next()
with tf.Session() as sess:
for _ in range(9):
print(sess.run(element))
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…