PyTorch DataLoader的num_workers过高导致CPU爆满,根本原因是子进程过多且默认占满线程,引发调度与pickle开销;实操应设num_workers=0(调试/Windows)或2–4(Linux/macOS),配pin_memory=True,避免ToTensor等CPU密集操作,并通过环境变量或spawn方式限制线程。

PyTorch DataLoader的num_workers设太高导致CPU爆满
根本原因不是PyTorch本身吃CPU,而是DataLoader启了太多子进程读数据,每个又默认用满所有可用线程——尤其在小数据集或SSD不快的机器上,IO没瓶颈,反而是进程调度和pickle序列化在空转耗CPU。
实操建议:
-
num_workers=0最简单:直接在主进程加载,适合调试、小数据或Windows(fork开销大) - Linux/macOS下推荐
num_workers=2到num_workers=4,再高收益极低,还易触发GIL争抢和内存拷贝 - 务必配
pin_memory=True(仅GPU训练时),否则num_workers>0时数据从worker传回主进程会多一次内存拷贝,拖慢且增CPU负载 - 避免用
torchvision.transforms.ToTensor()这类含PIL解码的操作在DataLoader里——它们是CPU密集型,应提前存为.pt或用torch.compile加速预处理逻辑
torch.set_num_threads()对DataLoader无效?
这个函数只限制PyTorch内部算子(如torch.mm、torch.nn.functional.conv2d)的OpenMP线程数,但DataLoader子进程用的是Python标准库multiprocessing,完全绕过它。
真正起作用的是:
立即学习“Python免费学习笔记(深入)”;
- 启动前设环境变量:
os.environ["OMP_NUM_THREADS"] = "1"(影响OpenMP) - 同时设:
os.environ["MKL_NUM_THREADS"] = "1"(影响Intel MKL) - 对每个
DataLoader子进程生效,需在__getitem__开头加torch.set_num_threads(1)——但更稳妥是在主进程启动前全局设好,再用spawn方式启动worker(mp.set_start_method("spawn"))
Windows上DataLoader卡死+CPU 100%的典型表现
错误现象:DataLoader卡在第一个batch,任务管理器看到一堆python.exe子进程,CPU持续100%,GPU显存却空着。这是Windows对fork支持差,子进程继承了主进程的CUDA上下文,引发死锁。
必须做的三件事:
- 在
if __name__ == "__main__":下启动训练——防止子进程重复导入主模块 - 显式指定
mp.set_start_method("spawn")(PyTorch 1.9+默认,但旧版或某些打包环境可能仍用fork) - 把
DataLoader定义移到__main__块内,别放在模块顶层;或者用num_workers=0保底
用torch.utils.data.Dataset做内存映射替代全量加载
当数据集能放进内存但加载慢(比如上百GB的.npy堆叠文件),DataLoader反复读磁盘就是CPU和IO双高。这时不用改num_workers,而该换数据组织方式。
关键操作:
- 用
numpy.memmap或torch.load(..., map_location="cpu")配合weights_only=True加载大文件,实现按需读页 - 自定义
Dataset.__getitem__中,只用memmap[idx]取单样本,不调np.load或torch.load - 避免在
__init__里做任何os.listdir或pandas.read_csv——这些应提前提取索引并缓存为.pt文件
复杂点在于内存映射需要对齐数据结构,一旦dtype或shape有微小错位,读出来就是乱码——建议先用np.memmap写个最小验证脚本,确认offset和mode="r"能对上再集成进Dataset。


















