コード例 #1
0
 def __init__(self,
              loader,
              prefetch_size,
              batch_size,
              num_cores=1,
              devices=None,
              fused_mode=False):
     self._loader = loader
     self._prefetch_size = prefetch_size
     self._batch_size = batch_size
     self._num_cores = num_cores
     self._devices = list(devices) if devices else None
     self._fused_mode = fused_mode
     self._batch_number = 0
     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
     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()
コード例 #2
0
ファイル: parallel_loader.py プロジェクト: zcain117/xla
 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()
コード例 #3
0
 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)