Pytorch中DataLoader类的多线程实现方法分析

之前在改自定义的DataSet的时候，由于在getitem()里面写了太多操作，导致训练过程贼慢，于是考虑用多线程优化一下。查阅一些资料发现pytorch在DataLoader里面就有多线程的实现，只要在定义的时候将num_worker设置成大于0就可以了。遂想要探索一下pytorch具体的实现方法。

首先找到迭代器：

def __iter__(self):
    return _DataLoaderIter(self)

初始化：

def __init__(self, loader):
    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()

collate_fn：将数据整合成一个batch返回的方法，用户可以自定义
batch_sampler：自定义如何取样
pin_menory：是否将数据集拷贝到显卡上
done_event：事件管理标志
sample_iter：迭代器，所以batch_sampler应该类似于用户自定义的一个数据的列表，用来生成可迭代对象sample_iter。

下面是与多线程有关的一些定义：

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)]

worker_init_fn：用户定义的每个worker初始化的时候需要执行的函数。
index_queues：这里用到了multiprocessing，pytorch的multiprocessing是对python原生的multiprocessing的一个封装，不过好像基本没什么变化。这里定义一个队列，multiprocessing的Queue类（这个Queue的父类）提供了put()和get()方法，用来向队列中增加线程和移除线程并返回结果。Pytorch的封装另外提供了send()和recv()方法，用来接收和读取缓存，具体实现和作用这里暂且按下不表。通过阅读后面的代码发现，这个队列里面返回的是当前数据在数据集中的位置。

workers：创建用户定义数量的线程，首先来看这个_worker_loop

def _worker_loop(dataset, index_queue, data_queue, collate_fn, seed, init_fn, worker_id):
    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)

    if init_fn is not None:
        init_fn(worker_id)

    watchdog = ManagerWatchdog()

    while True:
        try:
            r = index_queue.get(timeout=MANAGER_STATUS_CHECK_INTERVAL)
        except queue.Empty:
            if watchdog.is_alive():
                continue
            else:
                break
        if r is None:
            break
        idx, batch_indices = r
        try:
            samples = collate_fn([dataset[i] for i in batch_indices])
        except Exception:
            data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
        else:
            data_queue.put((idx, samples))
            del samples

_set_worker_signal_handlers()：是一个C函数，作用可以通过字面理解。
torch.set_num_threads(1)：设置线程数为1
下面两句固定了python和pytorch产生的随机数。
watchdog：查看管理进程状态是否改变

接下来的代码就可以理解了：每个worker的工作其实就是从index序列中读取当前数据在数据集中的序号，然后将对应的数据从数据集中取出来，扔到collate_fn中形成一个batch，再把batch扔到数据序列中，完成一次工作循环。

一直都没有接触多线程，正好这次了解一下多线程的基本写法以及思想，且作为第一篇文章更出来吧。
后面还有一些代码，坑想起来再填~