def __init__(
self,
num_workers=1,
pin_memory=False,
pinned_device=torch.device('cuda', 0),
worker_init_fn=None,
multiprocessing_context=None,
async_sleep=0.01,
):
if num_workers < 1:
raise ValueError('num_workers option should be positive')
if multiprocessing_context is None:
multiprocessing_context = multiprocessing
self.num_workers = num_workers
self._pin_memory = pin_memory
self._pinned_device = pinned_device
self._async_sleep = async_sleep
self._worker_pids_set = False
self._shutdown = False
self._workers_done_event = multiprocessing_context.Event()
self._task_id = 0
self._worker_is_active = []
self._task_queues = []
self._worker_result_queues = []
self._workers = []
self._available_worker_queue = queue.Queue()
for i in range(self.num_workers):
task_queue = multiprocessing_context.Queue()
result_queue = multiprocessing_context.Queue()
w = multiprocessing_context.Process(
target=_worker_loop,
args=(
task_queue,
result_queue,
self._workers_done_event,
worker_init_fn,
i,
),
)
w.daemon = True
# NB: Process.start() actually take some time as it needs to
# start a process and pass the arguments over via a pipe.
# Therefore, we only add a worker to self._workers list after
# it started, so that we do not call .join() if program dies
# before it starts, and __del__ tries to join but will get:
# AssertionError: can only join a started process.
w.start()
self._task_queues.append(task_queue)
self._worker_result_queues.append(result_queue)
self._workers.append(w)
self._worker_is_active.append(True)
self._available_worker_queue.put(i)
if self._pin_memory:
self._pin_memory_thread_done_event = threading.Event()
self._results_queues = [
queue.Queue() for _ in range(self.num_workers)
]
pin_memory_thread = threading.Thread(
target=_pin_memory_loop,
args=(
self._worker_result_queues,
self._results_queues,
self._pinned_device,
self._pin_memory_thread_done_event,
self._async_sleep,
),
)
pin_memory_thread.daemon = True
pin_memory_thread.start()
# Similar to workers (see comment above), we only register
# pin_memory_thread once it is started.
self._pin_memory_thread = pin_memory_thread
else:
self._results_queues = self._worker_result_queues
signal_handling._set_worker_pids(id(self),
tuple(w.pid for w in self._workers))
signal_handling._set_SIGCHLD_handler()
self._worker_pids_set = True
self.__multiprocessing_context = multiprocessing_context
self.__initialized = True
def __init__(self, loader):
self.random_vars = loader.random_vars
self.dataset = loader.dataset
self.collate_fn = loader.collate_fn
self.batch_sampler = loader.batch_sampler
self.num_workers = loader.num_workers
self.pin_memory = loader.pin_memory and torch.cuda.is_available()
self.timeout = loader.timeout
self.done_event = threading.Event()
self.sample_iter = iter(self.batch_sampler)
base_seed = torch.LongTensor(1).random_().item()
if self.num_workers > 0:
self.worker_init_fn = loader.worker_init_fn
self.index_queues = [
multiprocessing.Queue() for _ in range(self.num_workers)
]
self.worker_queue_idx = 0
self.worker_result_queue = multiprocessing.SimpleQueue()
self.batches_outstanding = 0
self.worker_pids_set = False
self.shutdown = False
self.send_idx = 0
self.rcvd_idx = 0
self.reorder_dict = {}
self.workers = [
multiprocessing.Process(
target=_worker_loop,
args=(self.dataset, self.index_queues[i],
self.worker_result_queue, self.collate_fn,
base_seed + i, self.worker_init_fn, i))
for i in range(self.num_workers)
]
if self.pin_memory or self.timeout > 0:
self.data_queue = queue.Queue()
if self.pin_memory:
maybe_device_id = torch.cuda.current_device()
else:
# do not initialize cuda context if not necessary
maybe_device_id = None
self.worker_manager_thread = threading.Thread(
target=_pin_memory_loop,
args=(self.worker_result_queue, self.data_queue,
maybe_device_id, self.done_event))
self.worker_manager_thread.daemon = True
self.worker_manager_thread.start()
else:
self.data_queue = self.worker_result_queue
for w in self.workers:
w.daemon = True # ensure that the worker exits on process exit
w.start()
# 이부분 _update_worker_pids 를 _set_worker_pids로 변경
_set_worker_pids(id(self), tuple(w.pid for w in self.workers))
_set_SIGCHLD_handler()
self.worker_pids_set = True
# prime the prefetch loop
for _ in range(2 * self.num_workers):
self._put_indices()
def __init__(self, streamer):
self.data_reader = streamer.data_reader
self.num_workers = streamer.num_workers
self.pin_memory = streamer.pin_memory and torch.cuda.is_available()
self.timeout = streamer.timeout
base_seed = torch.LongTensor(1).random_().item()
if self.num_workers > 0:
self.worker_init_fn = streamer.worker_init_fn
self.worker_result_queue = multiprocessing.Queue()
self.batch_queue = multiprocessing.Queue()
self.batches_outstanding = 0
self.worker_pids_set = False
self.shutdown = False
self.send_idx = 0
self.done_event = multiprocessing.Event()
self.worker_done_events = [
multiprocessing.Event() for _i in range(self.num_workers)
]
self.workers = []
for i in range(self.num_workers):
w = multiprocessing.Process(
target=_worker_loop,
args=(
self.data_reader,
self.batch_queue,
self.worker_result_queue,
self.done_event,
self.worker_done_events[i],
base_seed + i,
self.worker_init_fn,
i,
),
)
w.daemon = True # ensure that the worker exits on process exit
# Process.start() actually take some time as it needs to start a
# process and pass the arguments over via a pipe. Therefore, we
# only add a worker to self.workers list after it started, so
# that we do not call .join() if program dies before it starts,
# and __del__ tries to join it but will get:
# AssertionError: can only join a started process.
w.start()
self.workers.append(w)
self.num_live_workers = self.num_workers
if self.pin_memory:
self.data_queue = queue.Queue()
self.pin_memory_thread = threading.Thread(
target=_pin_memory_loop,
args=(
self.worker_result_queue,
self.data_queue,
self.done_event,
self.pin_memory,
torch.cuda.current_device(),
),
)
self.pin_memory_thread.daemon = True
self.pin_memory_thread.start()
else:
self.data_queue = self.worker_result_queue
_set_worker_pids(id(self), tuple(w.pid for w in self.workers))
_set_SIGCHLD_handler()
self.worker_pids_set = True
# prime the prefetch loop
for _ in range(2 * self.num_workers):
self._put_indices()
else:
# No workers
self.data_reader_iter = iter(self.data_reader)
def __init__(self, streamer):
self.data_reader = streamer.data_reader
self.num_workers = streamer.num_workers
self.pin_memory = streamer.pin_memory and torch.cuda.is_available()
self.timeout = streamer.timeout
base_seed = torch.LongTensor(1).random_().item()
if self.num_workers > 0:
self.worker_init_fn = streamer.worker_init_fn
self.worker_result_queue = multiprocessing.Queue()
self.batch_queue = multiprocessing.Queue()
self.batches_outstanding = 0
self.worker_pids_set = False
self.shutdown = False
self.send_idx = 0
self.done_event = multiprocessing.Event()
self.worker_done_events = [
multiprocessing.Event() for _i in range(self.num_workers)
]
self.workers = []
for i in range(self.num_workers):
w = multiprocessing.Process(
target=_worker_loop,
args=(
self.data_reader,
self.batch_queue,
self.worker_result_queue,
self.done_event,
self.worker_done_events[i],
base_seed + i,
self.worker_init_fn,
i,
),
)
w.daemon = True # ensure that the worker exits on process exit
# Process.start() actually take some time as it needs to start a
# process and pass the arguments over via a pipe. Therefore, we
# only add a worker to self.workers list after it started, so
# that we do not call .join() if program dies before it starts,
# and __del__ tries to join it but will get:
# AssertionError: can only join a started process.
w.start()
self.workers.append(w)
self.num_live_workers = self.num_workers
if self.pin_memory:
self.data_queue = queue.Queue()
self.pin_memory_thread = threading.Thread(
target=_pin_memory_loop,
args=(
self.worker_result_queue,
self.data_queue,
self.done_event,
self.pin_memory,
torch.cuda.current_device(),
),
)
self.pin_memory_thread.daemon = True
self.pin_memory_thread.start()
else:
self.data_queue = self.worker_result_queue
_set_worker_pids(id(self), tuple(w.pid for w in self.workers))
_set_SIGCHLD_handler()
self.worker_pids_set = True
# prime the prefetch loop
for _ in range(2 * self.num_workers):
self._put_indices()
else:
# No workers
self.data_reader_iter = iter(self.data_reader)