参考链接
https://www.cnblogs.com/hesse-summer/p/11343870.html
https://stackoverflow.com/questions/53998282/how-does-the-number-of-workers-parameter-in-pytorch-dataloader-actually-work
注:将数据转移到GPU上并不是
DataLoader的job;
最好不要在Dataset类中将数据放到GPU上(Dataset的__getitem__()
方法),因为此时处理的数据不是batch,batch数据是通过DataLoader加载进来的;若提前将数据放到GPU上,会加大开销-incurring a lot of overhead
第一部分
train_loader = DataLoader(dataset=train_data, batch_size=train_bs, shuffle=True, num_worker=4)
- dataloader一次性创建
num_worker
个worker,(也可以说dataloader一次性创建num_worker
个工作进程,worker也是普通的工作进程),并用batch_sampler
将指定batch分配给指定worker,worker将它负责的batch加载进RAM。 - 然后,dataloader从RAM中找本轮迭代要用的batch,如果找到了,就使用。如果没找到,就要
num_worker
个worker继续加载batch到内存,直到dataloader在RAM中找到目标batch。一般情况下都是能找到的,因为batch_sampler
指定batch时当然优先指定本轮要用的batch。 - 如果
num_worker
设为0,意味着每一轮迭代时,dataloader不再有自主加载数据到RAM这一步骤(因为没有worker了),而是在RAM中找batch,找不到时再加载相应的batch。缺点当然是速度更慢。
综上,我对
num_worker
的理解为:创建多线程,提前加载未来会用到的batch数据
第二部分
参考链接–https://tianws.github.io/skill/2019/08/27/gpu-volatile/
参考代码–pytorch的dataloader.py的第830行self._shutdown_workers()
此外,使用data_prefetcher的方法可以预先读取
下一个
batch的数据。
DataLoader的num_worker
和data_prefetcher
的区别:
- PyTorch 默认的 DataLoader 会创建一些 worker 线程来预读取新的数据,但是只有等到这些线程的数据全部都被清空(即预读取的batch数据全部被主进程使用),这些线程才会读下一批数据。
- 使用 data_prefetcher,可以保证线程不会等待,每个线程都总有至少一个数据在加载(模型的forward和预读取下一个batch同时进行)