PyTorch中的DataLoader和DataLoaderIter源码解读

蔚落 2022-02-09 10:25 447阅读 0赞

首先简单介绍一下DataLoader，它是PyTorch中数据读取的一个重要接口，该接口定义在dataloader.py中，只要是用PyTorch来训练模型基本都会用到该接口（除非用户重写…），该接口的目的：将自定义的Dataset根据batch size大小、是否shuffle等封装成一个Batch Size大小的Tensor，用于后面的训练。  
  
官方对DataLoader的说明是：“数据加载由数据集和采样器组成，基于python的单、多进程的iterators来处理数据。”关于iterator和iterable的区别和概念请自行查阅，在实现中的差别就是iterators有\_\_iter\_\_和\_\_next\_\_方法，而iterable只有\_\_iter\_\_方法。

dataset(Dataset): 传入的数据集
        batch_size(int, optional): 每个batch有多少个样本
    
        shuffle(bool, optional): 在每个epoch开始的时候，对数据进行重新排序
    
        sampler(Sampler, optional): 自定义从数据集中取样本的策略，如果指定这个参数，那么shuffle必须为False
    
        batch_sampler(Sampler, optional): 与sampler类似，但是一次只返回一个batch的indices（索引），需要注意的是，一旦指定了这个参数，那么batch_size,shuffle,sampler,drop_last就不能再制定了（互斥——Mutually exclusive）
    
        num_workers (int, optional): 这个参数决定了有几个进程来处理data loading。0意味着所有的数据都会被load进主进程。（默认为0）
    
        collate_fn (callable, optional): 将一个list的sample组成一个mini-batch的函数
    
        pin_memory (bool, optional)： 如果设置为True，那么data loader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中.
    
        drop_last (bool, optional): 如果设置为True：这个是对最后的未完成的batch来说的，比如你的batch_size设置为64，而一个epoch只有100个样本，那么训练的时候后面的36个就被扔掉了，如果为False（默认），那么会继续正常执行，只是最后的batch_size会小一点。
    
        timeout(numeric, optional): 如果是正数，表明等待从worker进程中收集一个batch等待的时间，若超出设定的时间还没有收集到，那就不收集这个内容了。这个numeric应总是大于等于0。默认为0
    
        worker_init_fn (callable, optional): 每个worker初始化函数 If not None, this will be called on each
        worker subprocess with the worker id (an int in [0, num_workers - 1]) as
        input, after seeding and before data loading. (default: None)

源码中DataLoader的参数解释：

class DataLoader(object):
        r"""
        Data loader. Combines a dataset and a sampler, and provides
        single- or multi-process iterators over the dataset.
    
        Arguments:
            dataset (Dataset): dataset from which to load the data.
            batch_size (int, optional): how many samples per batch to load
                (default: ``1``).
            shuffle (bool, optional): set to ``True`` to have the data reshuffled
                at every epoch (default: ``False``).
            sampler (Sampler, optional): defines the strategy to draw samples from
                the dataset. If specified, ``shuffle`` must be False.
            batch_sampler (Sampler, optional): like sampler, but returns a batch of
                indices at a time. Mutually exclusive with :attr:`batch_size`,
                :attr:`shuffle`, :attr:`sampler`, and :attr:`drop_last`.
            num_workers (int, optional): how many subprocesses to use for data
                loading. 0 means that the data will be loaded in the main process.
                (default: ``0``)
            collate_fn (callable, optional): merges a list of samples to form a mini-batch.
            pin_memory (bool, optional): If ``True``, the data loader will copy tensors
                into CUDA pinned memory before returning them.
            drop_last (bool, optional): set to ``True`` to drop the last incomplete batch,
                if the dataset size is not divisible by the batch size. If ``False`` and
                the size of dataset is not divisible by the batch size, then the last batch
                will be smaller. (default: ``False``)
            timeout (numeric, optional): if positive, the timeout value for collecting a batch
                from workers. Should always be non-negative. (default: ``0``)
            worker_init_fn (callable, optional): If not ``None``, this will be called on each
                worker subprocess with the worker id (an int in ``[0, num_workers - 1]``) as
                input, after seeding and before data loading. (default: ``None``)
    
        .. note:: By default, each worker will have its PyTorch seed set to
                  ``base_seed + worker_id``, where ``base_seed`` is a long generated
                  by main process using its RNG. However, seeds for other libraies
                  may be duplicated upon initializing workers (w.g., NumPy), causing
                  each worker to return identical random numbers. (See
                  :ref:`dataloader-workers-random-seed` section in FAQ.) You may
                  use :func:`torch.initial_seed()` to access the PyTorch seed for
                  each worker in :attr:`worker_init_fn`, and use it to set other
                  seeds before data loading.
    
        .. warning:: If ``spawn`` start method is used, :attr:`worker_init_fn` cannot be an
                     unpicklable object, e.g., a lambda function.
        """

当代码运行到要从torch.utils.data.DataLoader类生成的对象中取数据的时候，比如：

train_data=torch.utils.data.DataLoader(...) 
    for i, (input, target) in enumerate(train_data): 
    ...

就会调用DataLoader类的\_\_iter\_\_方法，\_\_iter\_\_方法就一行代码：return DataLoaderIter(self)，输入正是DataLoader类的属性。因此当调用\_\_iter\_\_方法的时候就牵扯到另外一个类：DataLoaderIter

DataLoaderIter类源码如下。self.index\_queue = multiprocessing.SimpleQueue()中的multiprocessing是Python中的多进程管理包，而threading则是Python中的多线程管理包，二者很大一部分的接口用法类似。还是照例先看看\_\_init\_\_，前面部分都是一些赋值操作，比较特殊的是self.sample\_iter = iter(self.batch\_sampler)，得到的self.sample\_iter可以通过next(self.sample\_iter)来获取batch size个数据的index。self.rcvd\_idx表示读取到的一个batch数据的index，初始化为0，该值在迭代读取数据的时候会用到。if self.num\_workers语句是针对多进程或单进程的情况进行初始化，如果不是设置为多进程读取数据，那么就不需要这些初始化操作，后面会介绍单进程数据读取。在if语句中通过multiprocessing.SimpleQueue()类创建了一个简单的队列对象。multiprocessing.Process类就是构造进程的类，这里根据设定的进程数来启动，然后赋值给self.workers。接下来的一个for循环就通过调用start方法依次启动self.workers中的进程。接下来关于self.pin\_memory的判断语句，该判断语句内部主要是实现了多线程操作。self.pin\_memory的含义在前面已经介绍过了，当为True的时候，就会把数据拷到CUDA中。self.data\_queue = queue.Queue()是通过Python的queue模块初始化得到一个先进先出的队列（queue模块也可以初始化得到先进后出的队列，需要用queue.LifoQueue()初始化），queue模块主要应用在多线程读取数据中。在threading.Thread的args参数中，第一个参数in\_data就是一个进程的数据，一个进程中不同线程的数据也是通过队列来维护的，这里采用的是Python的queue模块来初始化得到一个队列：queue.Queue()。初始化结束后，就会调用\_\_next\_\_方法，接下来介绍。   
**总的来说，如果设置为多进程读取数据，那么就会采用队列的方式来读，如果不是采用多进程来读取数据，那就采用普通方式来读**。

class DataLoaderIter(object):
        "Iterates once over the DataLoader's dataset, as specified by the sampler"
    
        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)
    
            if self.num_workers > 0:
                self.worker_init_fn = loader.worker_init_fn
                self.index_queue = multiprocessing.SimpleQueue()
                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 = {}
    
                base_seed = torch.LongTensor(1).random_()[0]
                self.workers = [
                    multiprocessing.Process(
                        target=_worker_loop,
                        args=(self.dataset, self.index_queue, 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()
                    self.worker_manager_thread = threading.Thread(
                        target=_worker_manager_loop,
                        args=(self.worker_result_queue, self.data_queue, self.done_event, self.pin_memory,
                              torch.cuda.current_device()))
                    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(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()

DataLoaderIter类的\_\_next\_\_方法如下，包含3个if语句和1个while语句。   
第一个if语句是用来处理self.num\_workers等于0的情况，也就是不采用多进程进行数据读取，可以看出在这个if语句中先通过indices = next(self.sample\_iter)获取长度为batch size的列表：indices，这个列表的每个值表示一个batch中每个数据的index，每执行一次next操作都会读取一批长度为batch size的indices列表。然后通过self.collate\_fn函数将batch size个tuple（每个tuple长度为2，其中第一个值是数据，Tensor类型，第二个值是标签，int类型）封装成一个list，这个list长度为2，两个值都是Tensor，一个是batch size个数据组成的FloatTensor，另一个是batch size个标签组成的LongTensor。所以简单讲self.collate\_fn函数就是将batch size个分散的Tensor封装成一个Tensor。batch = pin\_memory\_batch(batch)中pin\_memory\_batch函数的作用就是将输入batch的每个Tensor都拷贝到CUDA中，该函数后面会详细介绍。   
第二个if语句是判断当前想要读取的batch的index(self.rcvd\_idx)是否之前已经读出来过(已读出来的index和batch数据保存在self.reorder\_dict字典中，可以结合最后的while语句一起看，因为self.reorder\_dict字典的更新是在最后的while语句中），如果之前已经读取过了，就根据这个index从reorder\_dict字典中弹出对应的数据。最后返回batch数据的时候是 return self.\_process\_next\_batch(batch)，该方法后面会详细介绍。主要做是获取下一个batch的数据index信息。   
第三个if语句，self.batches\_outstanding的值在前面初始中调用self.\_put\_indices()方法时修改了，所以假设你的进程数self.num\_workers设置为3，那么这里self.batches\_outstanding就是3\*2=6，可具体看self.\_put\_indices()方法。   
最后的while循环就是真正用来从队列中读取数据的操作，最主要的就是idx, batch = self.\_get\_batch()，通过调用\_get\_batch()方法来读取，后面有介绍，简单讲就是调用了队列的get方法得到下一个batch的数据，得到的batch一般是长度为2的列表，列表的两个值都是Tensor，分别表示数据（是一个batch的）和标签。\_get\_batch()方法除了返回batch数据外，还得到另一个输出：idx，这个输出表示batch的index，这个if idx != self.rcvd\_idx条件语句表示如果你读取到的batch的index不等于当前想要的index:selg,rcvd\_idx，那么就将读取到的数据保存在字典self.reorder\_dict中：self.reorder\_dict\[idx\] = batch，然后继续读取数据，直到读取到的数据的index等于self.rcvd\_idx。

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)
    
            if self.batches_outstanding == 0:
                self._shutdown_workers()
                raise StopIteration
    
            while True:
                assert (not self.shutdown and self.batches_outstanding > 0)
                idx, batch = self._get_batch()
                self.batches_outstanding -= 1
                if idx != self.rcvd_idx:
                    # store out-of-order samples
                    self.reorder_dict[idx] = batch
                    continue
                return self._process_next_batch(batch)

**pin\_memory\_batch函数**不是定义在DataLoader类或DataLoaderIter类中。该函数主要是对batch中的Tensor执行batch.pin\_memory()操作，这里的很多条件语句只是用来判断batch的类型，假如batch是一个列表，列表中的每个值是Tensor，那么就会执行 elif isinstance(batch, collections.Sequence):这个条件，从而遍历该列表中的每个Tensor，然后执行第一个条件语句的内容： return batch.pin\_memory()

def pin_memory_batch(batch):
        if torch.is_tensor(batch):
            return batch.pin_memory()
        elif isinstance(batch, string_classes):
            return batch
        elif isinstance(batch, collections.Mapping):
            return {k: pin_memory_batch(sample) for k, sample in batch.items()}
        elif isinstance(batch, collections.Sequence):
            return [pin_memory_batch(sample) for sample in batch]
        else:
            return batch

**DataloaderIter类的\_get\_batch方法**。主要根据是否设置了超时时间来操作，如果超过指定的超时时间后没有从队列中读到数据就报错，如果不设置超时时间且一致没有从队列中读到数据，那么就会一直卡着且不报错，这部分是PyTorch后来修的一个bug。

def _get_batch(self):
            if self.timeout > 0:
                try:
                    return self.data_queue.get(True, self.timeout)
                except queue.Empty:
                    raise RuntimeError('DataLoader timed out after {} seconds'.format(self.timeout))
            else:
                return self.data_queue.get()

**DataLoaderIter类的\_process\_next\_batch方法**。首先对self.rcvd\_idx进行加一，也就是更新下下一个要读取的batch数据的index。然后调用\_put\_indices()方法获取下一个batch的每个数据的index。

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

**DataLoaderIter类的\_put\_indices方法。**该方法主要实现从self.sample\_iter中读取下一个batch数据中每个数据的index：indices = next(self.sample\_iter, None)，注意这里的index和前面idx是不一样的，这里的index是一个batch中每个数据的index，idx是一个batch的index；然后将读取到的index通过调用queue对象的put方法压到队列self.index\_queue中：self.index\_queue.put((self.send\_idx, indices))

def _put_indices(self):
            assert self.batches_outstanding < 2 * self.num_workers
            indices = next(self.sample_iter, None)
            if indices is None:
                return
            self.index_queue.put((self.send_idx, indices))
            self.batches_outstanding += 1
            self.send_idx += 1

参考：[https://blog.csdn.net/u014380165/article/details/79058479][https_blog.csdn.net_u014380165_article_details_79058479]

[https_blog.csdn.net_u014380165_article_details_79058479]: https://blog.csdn.net/u014380165/article/details/79058479