1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

if batch_sampler is None:
if sampler is None:
if shuffle:
sampler = RandomSampler(dataset)
##如果shuffer就随机

else:
sampler = SequentialSampler(dataset)
##否则顺序采样

batch_sampler = BatchSampler(sampler, batch_size, drop_last)
self.sampler = sampler
self.batch_sampler = batch_sampler

1
2
3
4
5
6
7
8
9
10
11
12
13

def __iter__(self):
batch = []
for idx in self.sampler:
batch.append(idx)
##遍历sampler获取数据，满batch_size就yield

if len(batch) == self.batch_size:
yield batch
batch = []
if len(batch) > 0 and not self.drop_last:
yield batch

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
# Our data model looks like this (queues are indicated with curly brackets):
#
#
main process
||
#
|
||
#
{index_queue}
||
#
|
||
#
worker processes
||
DATA
#
|
||
#
{worker_result_queue}
||
FLOW
#
|
||
#
pin_memory_thread of main process
||
DIRECTION
#
|
||
#
{data_queue}
||
#
|
||
#
data output
/
#
# P.S. `worker_result_queue` and `pin_memory_thread` part may be omitted if
#
`pin_memory=False`.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

self.index_queues = []
self.workers = []
for i in range(self.num_workers):
index_queue = multiprocessing.Queue() # 1.每个子进程一个队列放要处理的下标
index_queue.cancel_join_thread()
w = multiprocessing.Process(
target=_utils.worker._worker_loop, # 每个子进程循环执行的函数

args=(self.dataset, index_queue,
self.worker_result_queue, self.done_event, #2.self.worker_result_queue 多子进程公用要返回batch数据的队列

self.collate_fn, base_seed + i,
self.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.index_queues.append(index_queue)
self.workers.append(w)

1
2
3
4
5
6
7
8
9
10

def _process_next_batch(self, batch):
self.rcvd_idx += 1
self._put_indices()
## 先放下一批数据下标
if isinstance(batch, ExceptionWrapper):
raise batch.exc_type(batch.exc_msg)
return batch
## 再返回该批数据

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

def __next__(self):
if self.num_workers == 0:
# same-process loading
indices = next(self.sample_iter)
# may raise StopIteration
batch = self.collate_fn([self.dataset[i] for i in indices])
if self.pin_memory:
batch = pin_memory_batch(batch)
return batch
# check if the next sample has already been generated
if self.rcvd_idx in self.reorder_dict:
batch = self.reorder_dict.pop(self.rcvd_idx)
return self._process_next_batch(batch) ## 5. 之前以及取出来该下标数据，直接返回
if self.batches_outstanding == 0:
self._shutdown_workers()
raise StopIteration
while True:
## 1.直到取的数据下标正确才return
assert (not self.shutdown and self.batches_outstanding > 0)
idx, batch = self._get_batch()
## 2.从worker_result_queue中获取数据
self.batches_outstanding -= 1
if idx != self.rcvd_idx:
# store out-of-order samples
self.reorder_dict[idx] = batch
## 3.下标不对先存一下
continue
return self._process_next_batch(batch) ## 4.内部先放下一批数据下标再返回batch数据

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def _worker_loop(dataset, index_queue, data_queue, done_event, collate_fn, seed, init_fn, worker_id):
# See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the
# logic of this function.
try:
global _use_shared_memory
_use_shared_memory = True
# Intialize C side signal handlers for SIGBUS and SIGSEGV. Python signal
# module's handlers are executed after Python returns from C low-level
# handlers, likely when the same fatal signal happened again already.
# https://docs.python.org/3/library/signal.html Sec. 18.8.1.1
_set_worker_signal_handlers()
torch.set_num_threads(1)
random.seed(seed)
torch.manual_seed(seed)
data_queue.cancel_join_thread()
if init_fn is not None:
init_fn(worker_id)
watchdog = ManagerWatchdog()
while watchdog.is_alive():
try:
r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL) ##从index_queue中获取要处理的下标
except queue.Empty:
continue
if r is None:
# Received the final signal
assert done_event.is_set()
return
elif done_event.is_set():
# Done event is set. But I haven't received the final signal
# (None) yet. I will keep continuing until get it, and skip the
# processing steps.
continue
idx, batch_indices = r
try:
samples = collate_fn([dataset[i] for i in batch_indices]) ##1.根据下标取样本数据

except Exception:
# It is important that we don't store exc_info in a variable,
# see NOTE [ Python Traceback Reference Cycle Problem ]
data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
else: ## 2. 没有抛异常就将样本数据放入结果返回队列

data_queue.put((idx, samples))
del samples
except KeyboardInterrupt:
# Main process will raise KeyboardInterrupt anyways.
pass