训练代码中num_worker的设置

想要躺平的一枚

已于 2023-10-25 11:51:07 修改

阅读量6.9k

点赞数 8

分类专栏： AI图像算法文章标签： pytorch 深度学习人工智能

于 2023-10-25 11:32:21 首次发布

本文链接：https://blog.csdn.net/qq_43199575/article/details/134030400

版权

AI图像算法专栏收录该内容

23 篇文章 ¥9.90 ¥99.00

订阅专栏

超级会员免费看

num_workers在深度学习训练中影响数据加载速度，理论上越大越好，但实际效果受CPU核心数、batchsize和图像尺寸等因素影响。实验表明，最佳num_workers值可能随batchsize变化。例如，batchsize为64时最佳值为12，而batchsize为128时为8。选择合适值能提高训练效率。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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