Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save reubenjohn/7ab84fdb282d74366e0d4994a5e4ee09 to your computer and use it in GitHub Desktop.
Save reubenjohn/7ab84fdb282d74366e0d4994a5e4ee09 to your computer and use it in GitHub Desktop.
Advancing tensorflow dataset iterator in python multiprocessing Queue
from multiprocessing import Process, Queue
import numpy as np
import tensorflow as tf
def store(next_m, queue):
with tf.Session() as sess:
while True:
queue.put(sess.run(next_m))
if __name__ == '__main__':
pqueue = Queue()
a1 = np.arange(1000)
m = tf.data.Dataset.from_tensor_slices(a1).repeat().batch(1)
iter_m = m.make_one_shot_iterator()
m_init_ops = iter_m.make_initializer(m)
next_m = iter_m.get_next()
pp_process = Process(target=store, args=(next_m, pqueue,))
pp_process.daemon = True
pp_process.start() # <- Fork before starting this session!
with tf.Session() as sess:
for i in range(100000):
print(pqueue.get())
@wontleave
Copy link

output
This is the output I have from running your code with Visual Studio 2019.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment