num_workers 的作用与具体体现
num_workers 通常是用于控制数据加载和预处理的并行工作进程数量。主要体现在在进行一次epoch训练之前,加载一个批次图片所需的时间(若batchsize设置的比较大,如64,将num_workers设为0,就会发现进行一次epoch迭代训练之前所花费的时间很长;若batchsize设置的比较小,如8,则不是很明显)
如何选择合适的num_workers
num_workers和batchsize一样,在硬件允许的条件下,理论上越大越好。但实际应该设置为多少,目前都没有一个确定的说法。在此给出一个参考代码,通过设置的batchsize、图片尺寸与数据集的制作与加载,来查看哪个num_worker最快
代码如下(示例):
from dataset import trainIC_dataset
from torch.utils.data import DataLoader
import logging
import time
logger = logging.getLogger()
cropsize = (448,448)
image_path = 'lastdata/20231016'
ds = trainIC_dataset(image_path, paste_dir='', size=cropsize, mode='train')
dsval = trainIC_dataset(image_path, paste_dir='', size=cropsize, mode='val')
print(f"num of CPU: {mp.cpu_count()}")
for num_workers in range(2, mp.cpu_count(), 2):
dl = DataLoader(ds,
batch_size=128,
shuffle=True,
num_workers=num_workers,
pin_memory=False,
drop_last=False
)
dlval = DataLoader(dsval,
batch_size=128,
shuffle=True,
pin_memory=False,
# sampler = sampler_val,
num_workers=1,
drop_last=False)
start = time.time()
for epoch in range(1, 3):
for i, data in enumerate(dl, 0):
pass
end = time.time()
print("Finish with:{} second, num_workers={}".format(end - start, num_workers))
上述代码是测试训练集的num_worker,先得到cpu的最大核心数,再把num_worker从2、4、6、8…一直到最大核心数,来记录加载数据集的时间。(这里将验证集的num_workers 设置为1是为了确保可重现性和避免潜在的并发问题; 当然也可以为了效率设置其他值)。
经实验发现,并不是num_worker越大,它的时间越少;并且,batchsize或图像尺寸的改变,也会改变其时间最少的num_worker值(如当batchsize=64,时间最少的num_worker为12;当batchsize=128时,时间最少的num_worker为8)。当然,这个时间仅供参考。