def __init__(self,
loader,
batch_size,
devices,
prefetch_size=4,
batchdim=0):
self._loader = loader
self._prefetch_size = prefetch_size
self._batch_size = batch_size
self._devices = list(devices)
self._device_indices = dict()
for i, device in enumerate(self._devices):
self._device_indices[device] = i
self._batch_number = 0
self._batchdim = batchdim
self._done = False
self._lock = threading.Lock()
self._loader_queue = kq.Queue(maxsize=self._prefetch_size)
self._queue = kq.KeydQueue(maxsize=self._prefetch_size)
self._worker_count = 0
self._data = None
self._device_slices = None
thread = threading.Thread(target=self._loader_worker)
thread.daemon = True
thread.start()
for _ in range(0, prefetch_size):
thread = threading.Thread(target=self._worker)
thread.daemon = True
thread.start()
def __init__(self, device, loader_prefetch_size, device_prefetch_size):
self.device = device
self.batch_number = 0
self.loader_queue = kq.Queue(maxsize=loader_prefetch_size)
self.queue = kq.Queue(maxsize=device_prefetch_size)