def multiprocess_generator(self, max_queue_size=32, num_processes=8): # Re-shuffle file list if self.shuffle and cfg.NUM_TRAINERS > 1: np.random.RandomState(self.shuffle_seed).shuffle(self.all_lines) num_lines = len(self.all_lines) // cfg.NUM_TRAINERS self.lines = self.all_lines[num_lines * cfg.TRAINER_ID:num_lines * (cfg.TRAINER_ID + 1)] self.shuffle_seed += 1 elif self.shuffle: np.random.shuffle(self.lines) # Create multiple sharding generators according to num_processes for multiple processes generators = [] for pid in range(num_processes): generators.append(self.sharding_generator(pid, num_processes)) try: enqueuer = GeneratorEnqueuer(generators) enqueuer.start(max_queue_size=max_queue_size, workers=num_processes) while True: generator_out = None while enqueuer.is_running(): if not enqueuer.queue.empty(): generator_out = enqueuer.queue.get(timeout=5) break else: time.sleep(0.01) if generator_out is None: break yield generator_out finally: if enqueuer is not None: enqueuer.stop()
def reader(): cnt = 0 try: enqueuer = GeneratorEnqueuer( infinite_reader(), use_multiprocessing=use_multiprocess_reader) enqueuer.start(max_queue_size=max_queue, workers=num_workers) generator_out = None while True: while enqueuer.is_running(): if not enqueuer.queue.empty(): generator_out = enqueuer.queue.get() break else: time.sleep(0.02) yield generator_out cnt += 1 if cnt >= total_iter: enqueuer.stop() return generator_out = None except Exception as e: print("Exception occured in reader: {}".format(str(e))) finally: if enqueuer: enqueuer.stop()
def reader(): try: enqueuer = GeneratorEnqueuer( batches, use_multiprocessing=use_multiprocessing) enqueuer.start(max_queue_size=max_queue, workers=num_workers) generator_out = None for i in range(total_step): while enqueuer.is_running(): if not enqueuer.queue.empty(): generator_out = enqueuer.queue.get() break else: time.sleep(0.02) yield generator_out generator_out = None enqueuer.stop() finally: if enqueuer is not None: enqueuer.stop()
def reader(): cnt = 0 try: enqueuer = GeneratorEnqueuer( infinite_reader(), use_multiprocessing=use_multiprocess_reader) enqueuer.start(max_queue_size=max_queue, workers=num_workers) generator_out = None while True: while enqueuer.is_running(): if not enqueuer.queue.empty(): generator_out = enqueuer.queue.get() break else: time.sleep(0.02) yield generator_out cnt += 1 if cnt >= total_iter: enqueuer.stop() return generator_out = None finally: if enqueuer is not None: enqueuer.stop()