def next_func(data_queue, event):
while True:
if event.is_set():
batch, label = data_queue.get(block=True)
batch = dataloader._as_in_context(batch, context.cpu())
label = dataloader._as_in_context(label, context.cpu())
label = label.reshape((label.shape[0], ))
self._data_buffer.put((batch, label))
def _same_process_iter():
for batch in self._batch_sampler:
if isinstance(batch[0], (list, tuple)):
rets = [self._batchify_fn([self._dataset[idx] for idx in shard])
for shard in batch]
if self._pin_memory:
rets = [_as_in_context(ret, context.cpu_pinned()) for ret in rets]
yield rets
else:
ret = self._batchify_fn([self._dataset[idx] for idx in batch])
if self._pin_memory:
ret = _as_in_context(ret, context.cpu_pinned())
yield ret
def same_process_iter():
for batch in self._batch_sampler:
ret = self._batchify_fn(
[self._dataset[idx] for idx in batch])
if self._pin_memory:
ret = _as_in_context(
ret, context.cpu_pinned(self._pin_device_id))
yield ret
def _same_process_iter():
for batch in self._batch_sampler:
if isinstance(batch[0], (list, tuple)):
rets = [
self._batchify_fn(
[self._dataset[idx] for idx in shard])
for shard in batch
]
if self._pin_memory:
rets = [
_as_in_context(ret, context.cpu_pinned())
for ret in rets
]
yield rets
else:
ret = self._batchify_fn(
[self._dataset[idx] for idx in batch])
if self._pin_memory:
ret = _as_in_context(ret, context.cpu_pinned())
yield ret
def __next__(self):
self._push_next()
if self._rcvd_idx == self._sent_idx:
assert not self._data_buffer, 'Data buffer should be empty at this moment'
raise StopIteration
assert self._rcvd_idx < self._sent_idx, 'rcvd_idx must be smaller than sent_idx'
assert self._rcvd_idx in self._data_buffer, 'fatal error with _push_next, rcvd_idx missing'
ret = self._data_buffer.pop(self._rcvd_idx)
batch = pickle.loads(ret.get()) if self._dataset is None else ret.get()
if self._pin_memory:
batch = _as_in_context(batch, context.cpu_pinned())
self._rcvd_idx += 1
return batch
def _same_process_iter():
urls = []
dataset = [self._dataset[i] for i in iter(self._file_sampler)]
for i, url in enumerate(dataset):
urls.append(url)
if i < len(dataset) - 1:
if len(urls) < self._circle_length:
continue
if self._circle_length == 1:
urls = urls[0]
dataset, batch_sampler = _dataset_worker_fn(urls, self._dataset_fn,
self._batch_sampler_fn)
for batch in batch_sampler:
ret = self._batchify_fn([dataset[idx] for idx in batch])
if self._pin_memory:
ret = _as_in_context(ret, context.cpu_pinned())
yield ret
urls = []
def __next__(self):
self._iters += 1
if not self._sample_times or self._iters < self._sample_times:
self._push_next()
if self._rcvd_idx == self._sent_idx:
assert not self._data_buffer, "Data buffer should be empty at this moment"
raise StopIteration
assert self._rcvd_idx < self._sent_idx, "rcvd_idx must be smaller than sent_idx"
assert self._rcvd_idx in self._data_buffer, "fatal error with _push_next, rcvd_idx missing"
ret = self._data_buffer.pop(self._rcvd_idx)
try:
if self._dataset is None:
batch = pickle.loads(ret.get(self._timeout))
else:
batch = ret.get(self._timeout)
if self._pin_memory:
batch = _as_in_context(batch,
context.cpu_pinned(self._pin_device_id))
batch = batch[0] if len(batch) == 1 else batch
self._rcvd_idx += 1
return batch
except multiprocessing.context.TimeoutError:
msg = '''Worker timed out after {} seconds. This might be caused by \n
- Slow transform. Please increase timeout to allow slower data loading in each worker.
'''.format(self._timeout)
if not isinstance(self._worker_pool,
multiprocessing.pool.ThreadPool):
msg += '''- Insufficient shared_memory if `timeout` is large enough.
Please consider reduce `num_workers` or increase shared_memory in system.
'''
print(msg)
raise
except Exception:
self._worker_pool.terminate()
raise