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()