单进程

  1. 当用于在进程之间共享数据的资源（eg. 共享内存，文件描述符）有限时，或当整个数据集很小并能完全加载到内存中时，单进程模式可能是我们的首选。
  2. 当调试代码时，单进程加载通常可以显示更多可读的错误跟踪。（用Pycharm调试要设置成单进程才能步进）

但是单进程下，Dataloader初始化的进程跟取数据的进程一样，所以数据加载可能会阻止计算。

多进程

为了避免加载数据时阻塞计算，需要设置多进程
在PyTorch中只需将参数num_workers设置为正整数即可执行多进程数据加载，若设置为0则是单进程数据加载

在设置多进程模式时，每次 DataLoader 创建 iterator 时（例如，当调用 enumerate(dataloader) 时），都会创建 num_workers 个工作进程。此时dataset, collate_fn, worker_init_fn 都会被传到每个worker中，而每个 worker 都用独立的进程。

  1. 对于map-style数据，主线程会用sampler产生indices，并将它们送到worker里。因此shuffle是在主线程做的
  2. 对于iterable-style数据，因为每个worker都有相同的 data 复制样本，并在各个进程里进行不同的操作，以防止每个进程输出的数据是重复的，所以一般会使用torch.utils.data.get_worker_info()来进行辅助处理。这里，torch.utils.data.get_worker_info()会返回 worker 进程的一些信息(如id, dataset, num_workers, seed)，如果在主线程的话返回 None

注意，通常不建议在多进程加载中返回 CUDA 张量，因为在使用 CUDA 和在多处理中共享 CUDA 张量时存在许多微妙之处（文档中提出：只要接收过程保留张量的副本，就需要发送过程来保留原始张量）。建议采用 pin_memory=True ，以将数据快速传输到支持 CUDA 的 GPU。简而言之，不建议在使用多线程的情况下返回 CUDA 的 Tensor。

锁页内存

主机中的内存，有两种存在方式，一是锁页，二是不锁页。锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。
当数据源自固定(页面锁定）内存时，主机在GPU上的数据副本运行速度会快得多。 CPU Tensors 和存储器暴露了一个pin_memory()方法，该方法返回对象的一个副本，并将数据放入固定区域。

而显卡中的显存全部是锁页内存！当计算机的内存充足的时候，可以设置pin_memory=True。设置pin_memory=True，则意味着生成的 Tensor 数据最开始是属于内存中的锁页内存，这样将内存的 Tensor 转义到 GPU 的显存就会更快一些。同时，由于 pin_memory 的作用是将张量返回之前将其复制到 CUDA 固定的内存中，所以只有在 CUDA 环境支持下才有用。

Pytorch中原生的pin_memory方法支持大部分python数据类型的处理：

  1. def pin_memory(data):
  2.     if isinstance(data, torch.Tensor):
  3.         return data.pin_memory()
  4.     elif isinstance(data, string_classes):
  5.         return data
  6.     elif isinstance(data, container_abcs.Mapping):
  7.         return {k: pin_memory(sample) for k, sample in data.items()}
  8.     elif isinstance(data, tuple) and hasattr(data, '_fields'):  # namedtuple
  9.         return type(data)(*(pin_memory(sample) for sample in data))
  10.     elif isinstance(data, container_abcs.Sequence):
  11.         return [pin_memory(sample) for sample in data]
  12.     elif hasattr(data, "pin_memory"):
  13.         return data.pin_memory()
  14.     else:
  15.         return data

默认情况下，如果上面固定逻辑对一个属于自定义类型（custom type）对batch（如果有一个collate_fn返回自定义批处理类型的批处理，则会发生），或者如果该批处理的每个元素都是custom type，则该固定逻辑将无法识别它们，它会直接返回该批处理（或那些元素）而无需固定内存。而如果要为自定义批处理或数据类型启用内存固定，我们需使用pin_memory()在自定义类型上定义一个方法，如下：

  1. class SimpleCustomBatch:
  2.     # 自定义一个类，该类不能被PyTorch原生的pin_memory方法所支持
  4.     def __init__(self, data):
  5.         transposed_data = list(zip(*data))
  6.         self.inp = torch.stack(transposed_data[0], 0)
  7.         self.tgt = torch.stack(transposed_data[1], 0)
  9.     # custom memory pinning method on custom type
  10.     def pin_memory(self):
  11.         self.inp = self.inp.pin_memory()
  12.         self.tgt = self.tgt.pin_memory()
  13.         return self
  15. def collate_wrapper(batch):
  16.     return SimpleCustomBatch(batch)
  18. inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
  19. tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
  20. dataset = TensorDataset(inps, tgts)
  22. loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
  23.                     pin_memory=True)
  25. for batch_ndx, sample in enumerate(loader):
  26.     print(sample.inp.is_pinned())  # True
  27.     print(sample.tgt.is_pinned())  # True

预取（Prefetch）

Dataloader通过指定 prefetch_factor （默认为 2）来进行数据的预取。

  1. class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter):
  2.     def __init__(self, loader):
  3.         ...
  4.         self._reset(loader, first_iter=True)
  6.     def _reset(self, loader, first_iter=False):
  7.         ...
  8.         # prime the prefetch loop
  9.         for _ in range(self._prefetch_factor * self._num_workers):
  10.             self._try_put_index()

通过源码可以看到，prefetch 功能仅适用于多进程加载中（下面也会有多进程 dataloader 的部分代码分析）。