文章目录

  • 1 基本流程
  • 2 参数介绍
  • 关于worker_init_fn的问题
  • pin_memory

1 基本流程

首先熟知，pytorch 的数据加载到模型的操作顺序是这样的：

  1. 创建一个 Dataset 对象
  2. 创建一个 DataLoader 对象
  3. 循环这个 DataLoader 对象，将img, label加载到模型中进行训练
dataset = MyDataset()
dataloader = DataLoader(dataset)
num_epoches = 100
for epoch in range(num_epoches):
    for img, label in dataloader:
        ....

2 参数介绍

def __init__(self, dataset: Dataset[T_co], batch_size: Optional[int] = 1,
                 shuffle: bool = False, sampler: Optional[Sampler[int]] = None,
                 batch_sampler: Optional[Sampler[Sequence[int]]] = None,
                 num_workers: int = 0, collate_fn: _collate_fn_t = None,
                 pin_memory: bool = False, drop_last: bool = False,
                 timeout: float = 0, worker_init_fn: _worker_init_fn_t = None,
                 multiprocessing_context=None, generator=None,
                 *, prefetch_factor: int = 2,
                 persistent_workers: bool = False):
  • 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初始化的时候需要执行的函数。如果不是None, 则会以worker id[0, num_workers - 1]的每个子进程调用（在sedding后，数据加载前）

关于worker_init_fn的问题

如果程序一开始指定了各种seed，为了得到确定的结果。如

random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)

此时，在生成Dataloader时，如果要设置num_workers的数量大于0，使用多进程。则需要传入

def worker_init_fn(worker_id):
        random.seed(args.seed + worker_id)

pin_memory

pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。

主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。而显卡中的显存全部是锁页内存！

当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。

参考链接