训练代码中num_worker的设置


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)。当然,这个时间仅供参考

猜你喜欢

转载自blog.csdn.net/qq_43199575/article/details/134030400
num