申请tf.py_func
内部ds.map
作品。
dummy_file.txt:
10
这里是脚本:
import tensorflow as tf
filelist = ['dummy_file.txt', 'dummy_file.txt', 'dummy_file.txt']
def py_func(input):
# perform operations outside of tensorflow
parsed_txt_file = int(input)
return 'test', parsed_txt_file
def map_element_counts(fname):
# let tensorflow read the text file
file_string = tf.read_file(fname['filenames'])
# then use python function on the extracted string
a, b = tf.py_func(
func=py_func, inp=[file_string], Tout=[tf.string, tf.int64]
)
return {'elementA': a, 'elementB': b, 'file': fname['filenames']}
ds = tf.data.Dataset.from_tensor_slices({'filenames': filelist})
ds = ds.map(map_element_counts)
element = ds.make_one_shot_iterator().get_next()
with tf.Session() as sess:
print(sess.run(element))
print(sess.run(element))
print(sess.run(element))
输出:
{'file': b'dummy_file.txt', 'elementA': b'test', 'elementB': 10}
{'file': b'dummy_file.txt', 'elementA': b'test', 'elementB': 10}
{'file': b'dummy_file.txt', 'elementA': b'test', 'elementB': 10}