Pytorch DataLoader中的num，PyTorch DataLoader中的num工作机制和优化策略

温馨提示：这篇文章已超过412天没有更新，请注意相关的内容是否还可用！

在PyTorch中，DataLoader是一个重要的组件，用于数据加载和预处理。其中的参数num指的是每个批次加载的数据数量。这个参数决定了在训练过程中每次从数据集中取出的样本数量。通过调整这个参数，可以控制每次迭代时模型接收的数据量，从而影响训练的速度和效率。在使用DataLoader时，根据具体需求和计算资源调整num参数是很重要的。

num_workers是Dataloader的一个参数，默认值为0，它告诉DataLoader实例要使用多少个子进程进行数据加载，这个参数与CPU相关，与GPU无关。

Pytorch DataLoader中的num，PyTorch DataLoader中的num工作机制和优化策略第1张

当num_worker设为0时，DataLoader在每轮迭代时不会自主加载数据到RAM中，这意味着在RAM中查找batch时，如果找不到，则需要加载相应的batch，因此速度较慢。

当num_worker不为0时，DataLoader在每轮加载数据时，会一次性创建指定数量的worker，这些worker由batch_sampler分配特定的batch，并负责将其加载到RAM中，这种方式的好处是，由于下一轮迭代的batch可能已经在上一轮或更早的轮次中加载，因此寻找batch的速度更快，这也增加了内存开销和CPU负担，建议将num_workers设置为电脑或服务器的CPU核心数，如果CPU性能强、RAM充足，可以设置为更大的值。

值得注意的是，num_workers的值与模型训练的快速性有关，但并不影响训练出的模型的性能，Detectron2中num_workers的默认值为4。

选择最合适的num_workers值

最合适的num_workers值与使用的数据集有关，在选择num_workers的值之前，可以使用以下脚本进行性能测试：

Pytorch DataLoader中的num，PyTorch DataLoader中的num工作机制和优化策略第2张

import time
import multiprocessing as mp
import torch
import torchvision
from torchvision import transforms
transform = transforms.Compose([
    torchvision.transforms.ToTensor(),
    torchvision.transforms.Normalize((0.1307,), (0.3081,))
])
trainset = torchvision.datasets.MNIST(
    root='dataset/',
    train=True,
    download=True,
    transform=transform
)
print(f"Number of CPU cores: {mp.cpu_count()}")
for num_workers in range(2, mp.cpu_count(), 2):  
    train_loader = torch.utils.data.DataLoader(trainset, shuffle=True, num_workers=num_workers, batch_size=64, pin_memory=True)
    start = time()
    for epoch in range(1, 3):
        for i, data in enumerate(train_loader, 0):
            pass
    end = time()
    print(f"Finished with {end - start} seconds, num_workers={num_workers}")

通过运行此脚本并观察结果，可以选择最适合的num_workers值，在一个拥有24个CPU的服务器上，最合适的num_workers值可能是14。