PyTorch中next(iter(train_dataloader))执行耗时过长的原因与解决咨询

问题描述 #

我尝试用以下代码加载本地约225张图片的数据集，发现执行next(iter(train_dataloader))时耗时异常长，明明数据集规模很小，想知道原因和可行的优化方法。

加载数据集的代码：

# Set the batch size
BATCH_SIZE = 32 

# Create data loaders
train_dataloader, test_dataloader, class_names = data_setup.create_dataloaders(
  train_dir=train_dir,
  test_dir=test_dir,
  transform=manual_transforms, # use manually created transforms
  batch_size=BATCH_SIZE
)

# Get a batch of images
image_batch, label_batch = next(iter(train_dataloader)) # 耗时过长的操作

创建DataLoader的实现代码：

import os

from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import pdb

NUM_WORKERS = os.cpu_count()

def create_dataloaders(
  train_dir: str, 
  test_dir: str, 
  transform: transforms.Compose, 
  batch_size: int, 
  num_workers: int=NUM_WORKERS
):
# Use ImageFolder to create dataset(s)
train_data = datasets.ImageFolder(train_dir, transform=transform)
test_data = datasets.ImageFolder(test_dir, transform=transform)

# Get class names
class_names = train_data.classes

# Turn images into data loaders
train_dataloader = DataLoader(
  train_data,
  batch_size=batch_size,
  shuffle=True,
  num_workers=num_workers,
  pin_memory=True,
)
test_dataloader = DataLoader(
  test_data,
  batch_size=batch_size,
  shuffle=False, # don't need to shuffle test data
  num_workers=num_workers,
  pin_memory=True,
)

return train_dataloader, test_dataloader, class_names

耗时过长的原因 #

• 多进程启动开销：当前num_workers设为CPU核心数，对于仅225张图的小数据集，启动多个worker进程的初始化开销远大于实际加载、处理数据的时间，反而拖慢整体速度。
• 数据变换开销：如果manual_transforms包含复杂操作（比如大尺寸图片的随机裁剪、多轮像素增强、频繁的PIL与Tensor格式转换），每张图片的处理时间累加后，会被DataLoader的调度逻辑放大。
• 首次加载的磁盘IO与缓存缺失：第一次执行时，所有图片都需要从磁盘读取，没有内存缓存，小数据集的单文件IO延迟会被多进程的调度等待放大。
• pin_memory的额外开销：pin_memory=True会将数据锁页到内存以加速GPU传输，但小数据集下，这个操作的开销远大于带来的收益，反而增加了内存操作耗时。

可行的优化方法 #

• 调整num_workers参数：将num_workers设为0（单进程加载）或者1，避免多进程启动的额外开销，这是小数据集下最有效的优化手段。
• 简化数据变换逻辑：检查manual_transforms，移除不必要的数据增强操作；或者提前完成预处理，将图片转为Tensor格式保存到本地，后续直接加载预处理后的文件。
• 关闭pin_memory：将pin_memory改为False，减少内存锁页的额外操作开销。
• 预加载整个数据集到内存：直接把所有训练数据加载到内存中，避免反复磁盘IO，示例代码如下：

  train_data = datasets.ImageFolder(train_dir, transform=manual_transforms)
  # 预加载所有数据到内存
  train_data = [(img, label) for img, label in train_data]
  train_dataloader = DataLoader(
      train_data,
      batch_size=BATCH_SIZE,
      shuffle=True,
      num_workers=0,
      pin_memory=False
  )
  

• 优化shuffle策略：如果不需要每次迭代都重新打乱数据集，可以提前手动打乱一次，再传入DataLoader，减少每次迭代的打乱开销。

内容的提问来源于stack exchange，提问作者Jose Ramon