Pytorch中num_workers
Question
一直很迷,在给Dataloader
设置worker
数量(num_worker
)时,到底设置多少合适?这个worker
到底怎么工作的?如果将num_worker
设为0(也是默认值),就没有worker
了吗?
worker
的使用场景:
from torch.utils.data import DataLoader train_loader = DataLoader(dataset=train_data, batch_size=train_bs, shuffle=True, num_worker=4) valid_loader = DataLoader(dataset=valid_data, batch_size=valid_bs, num_worker=4)
Answer
-
每每轮到
dataloader
加载数据时:for epoch in range(start_epoch, end_epoch): for i, data in enumerate(trainloader):
-
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
设置得大,好处是寻batch
速度快,因为下一轮迭代的batch
很可能在上一轮/上上一轮…迭代时已经加载好了。坏处是内存开销大,也加重了CPU
负担(worker
加载数据到RAM
的进程是CPU
复制的嘛)。num_workers
的经验设置值是自己电脑/服务器的CPU
核心数,如果CPU
很强、RAM
也很充足,就可以设置得更大些。 -
如果
num_worker
设为0,意味着每一轮迭代时,dataloader
不再有自主加载数据到RAM
这一步骤(因为没有worker
了),而是在RAM
中找batch
,找不到时再加载相应的batch
。缺点当然是速度更慢。
转自 https://blog.csdn.net/weixin_43593330/article/details/107483671
同时可参考 https://zhuanlan.zhihu.com/p/69250939