Skip to content

Commit

Permalink
removed threadpoolctl
Browse files Browse the repository at this point in the history
FabianIsensee committed Sep 16, 2019
1 parent 52c4c77 commit 5fdf8d9
Showing 1 changed file with 21 additions and 23 deletions.
44 changes: 21 additions & 23 deletions batchgenerators/dataloading/multi_threaded_augmenter.py
Original file line number Diff line number Diff line change
@@ -28,37 +28,35 @@
from queue import Empty, Full
import traceback
from time import sleep, time
from threadpoolctl import threadpool_limits


def producer(queue, data_loader, transform, thread_id, seed, abort_event):
try:
with threadpool_limits(limits=1, user_api="blas"):
np.random.seed(seed)
data_loader.set_thread_id(thread_id)
item = None

while True:
# check if abort event was set
if not abort_event.is_set():
np.random.seed(seed)
data_loader.set_thread_id(thread_id)
item = None

if item is None:
while True:
# check if abort event was set
if not abort_event.is_set():

try:
item = next(data_loader)
if transform is not None:
item = transform(**item)
except StopIteration:
item = "end"
if item is None:

try:
queue.put(item, timeout=2)
item = None
except Full:
# queue was full because items in it were not consumed. Try again.
pass
else:
break
item = next(data_loader)
if transform is not None:
item = transform(**item)
except StopIteration:
item = "end"

try:
queue.put(item, timeout=2)
item = None
except Full:
# queue was full because items in it were not consumed. Try again.
pass
else:
break

except KeyboardInterrupt:
# drain queue, then give 'end', set abort flag and reraise KeyboardInterrupt

0 comments on commit 5fdf8d9

Please sign in to comment.