def get_batch(num_workers, **kwargs): try: enqueuer = GeneratorEnqueuer(generator(**kwargs), use_multiprocessing=True) enqueuer.start(max_queue_size=24, workers=num_workers) generator_output = None while True: while enqueuer.is_running(): if not enqueuer.queue.empty(): generator_output = enqueuer.queue.get() break else: time.sleep(0.01) yield generator_output generator_output = None finally: if enqueuer is not None: enqueuer.stop()
def get_batch(num_workers, **kwargs): try: enqueuer = GeneratorEnqueuer(generator(**kwargs), use_multiprocessing=True) print('Generator use 10 batches for buffering, this may take a while, you can tune this yourself.') enqueuer.start(max_queue_size=10, workers=num_workers) generator_output = None while True: while enqueuer.is_running(): if not enqueuer.queue.empty(): generator_output = enqueuer.queue.get() break else: time.sleep(0.01) yield generator_output generator_output = None finally: if enqueuer is not None: enqueuer.stop()