PyTorch使用DataLoader训练HAM10000数据集速度过慢求助

问题场景 #

使用HAM10000数据集训练小型神经网络，基于PyTorch的DataLoader实现数据加载，但训练过程本身很快，95%的时间消耗在Epoch之间，推测是数据加载环节效率低下。

原数据加载代码 #

class CocoDetectionWithFilenames(CocoDetection):
    def __init__(self, root: str, ann_file: str, transform=None):
        super().__init__(root, ann_file, transform)

    def get_filename(self, idx: int) -> str:
        return self.coco.loadImgs(self.ids[idx])[0]["file_name"]


def get_loaders(root: str, ann_file: str) -> tuple[CocoDetection, DataLoader, DataLoader, DataLoader]:
    transform = transforms.Compose([
        transforms.ToTensor()
    ])
    dataset = CocoDetectionWithFilenames(
        root=root,
        ann_file=ann_file,
        transform=transform
    )
    train_size = int(0.7 * len(dataset))
    valid_size = int(0.15 * len(dataset))
    test_size = len(dataset) - train_size - valid_size
    train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, valid_size, test_size])
    num_workers = os.cpu_count()
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        prefetch_factor=1024
    )
    valid_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
        prefetch_factor=1024
    )
    test_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )

    return dataset, train_loader, valid_loader, test_loader

原训练循环代码 #

def extract_bboxes(targets: list[dict]) -> list[torch.Tensor]:
    bboxes = []

    for target in targets:
        xs, ys, widths, heights = target["bbox"]

        for idx, _ in enumerate(xs):
            x1, y1, width, height = xs[idx], ys[idx], widths[idx], heights[idx]
            # Convert COCO format (x, y, width, height) to (x1, y1, x2, y2)
            x2, y2 = x1 + width, y1 + height

            bboxes.append(torch.IntTensor([x1, y1, x2, y2]))

    return bboxes

num_epochs = 25
train_losses = []
val_losses = []

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for images, targets in train_loader_tqdm:
        images = images.to(device)
        bboxes = extract_bboxes(targets)
        bboxes = torch.stack(bboxes).to(device)

        optimizer.zero_grad(set_to_none=True)

        outputs = model(images)
        loss = criterion(outputs, bboxes)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_train_loss = running_loss / len(train_loader)

    train_losses.append(epoch_train_loss)
    print(f"Epoch {epoch + 1}, Loss: {epoch_train_loss}")
    model.eval()

核心问题分析 #

1. 数据集分配错误：valid_loader和test_loader都错误使用了train_dataset，不仅逻辑错误，还会导致重复加载相同数据，浪费资源。
2. bbox转换在主线程执行：extract_bboxes在训练循环中处理，没有利用DataLoader的多进程预处理能力，导致数据加载与模型训练串行执行，拖慢Epoch间的衔接速度。
3. prefetch_factor设置过大：1024的预取因子会占用大量内存，引发内存竞争，反而降低加载效率。
4. 多进程数量不合理：直接使用os.cpu_count()会导致过多线程竞争CPU资源，反而降低处理速度。

优化方案 #

1. 修改Dataset，将预处理移至__getitem__ #

把bbox格式转换和张量生成放到Dataset的__getitem__方法中，让DataLoader的多进程提前完成预处理：

class CocoDetectionWithFilenames(CocoDetection):
    def __init__(self, root: str, ann_file: str, transform=None):
        super().__init__(root, ann_file, transform)

    def get_filename(self, idx: int) -> str:
        return self.coco.loadImgs(self.ids[idx])[0]["file_name"]
    
    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        img, target = super().__getitem__(idx)
        # 转换COCO格式bbox为(x1,y1,x2,y2)并转为张量
        bboxes = []
        for ann in target:
            x1, y1, w, h = ann["bbox"]
            x2 = x1 + w
            y2 = y1 + h
            bboxes.append([x1, y1, x2, y2])
        bboxes = torch.tensor(bboxes, dtype=torch.float32)
        return img, bboxes

2. 修正DataLoader的数据集分配与参数 #

• 正确分配valid/test数据集
• 调整num_workers为CPU核心数的一半，避免线程竞争
• 降低prefetch_factor至合理值（默认2即可）

def get_loaders(root: str, ann_file: str) -> tuple[CocoDetection, DataLoader, DataLoader, DataLoader]:
    transform = transforms.Compose([
        transforms.ToTensor()
    ])
    dataset = CocoDetectionWithFilenames(
        root=root,
        ann_file=ann_file,
        transform=transform
    )
    train_size = int(0.7 * len(dataset))
    valid_size = int(0.15 * len(dataset))
    test_size = len(dataset) - train_size - valid_size
    train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, valid_size, test_size])
    # 合理设置多进程数
    num_workers = max(1, os.cpu_count() // 2)
    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        prefetch_factor=2
    )
    valid_loader = torch.utils.data.DataLoader(
        valid_dataset,  # 修正为验证集
        batch_size=32,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
        prefetch_factor=2
    )
    test_loader = torch.utils.data.DataLoader(
        test_dataset,  # 修正为测试集
        batch_size=32,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True
    )

    return dataset, train_loader, valid_loader, test_loader

3. 简化训练循环 #

去掉主线程的extract_bboxes，直接使用Dataset预处理好的bbox张量：

num_epochs = 25
train_losses = []
val_losses = []

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0

    for images, bboxes in train_loader_tqdm:
        images = images.to(device)
        bboxes = bboxes.to(device)

        optimizer.zero_grad(set_to_none=True)

        outputs = model(images)
        loss = criterion(outputs, bboxes)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    epoch_train_loss = running_loss / len(train_loader)

    train_losses.append(epoch_train_loss)
    print(f"Epoch {epoch + 1}, Loss: {epoch_train_loss}")
    model.eval()

额外优化建议 #

• 启用CuDNN基准测试：在训练前添加torch.backends.cudnn.benchmark = True，让CuDNN自动选择最优卷积算法，提升模型计算速度。
• 磁盘IO优化：如果使用机械硬盘，更换为SSD可大幅提升图像读取速度。
• 数据缓存：对于重复使用的数据集，可将预处理后的图像和bbox缓存到磁盘，避免重复处理，比如使用joblib或自定义缓存逻辑。

内容的提问来源于stack exchange，提问作者Marek M.