为您的项目提供 GPU · 无 KYC 的加密货币支付 如何租用
简体中文
打开控制台
实用指南 / KERNODECK

DataLoader 阻塞:先检查数据,再检查 workers

先用 num_workers=0 和固定顺序开始,检查一个样本,再检查一个完整 batch,并把读取、变换、拼接和 GPU 传输分开。然后再逐步引入 workers。GPU 在等待并不证明存储很慢:数据错误、序列化或代价高昂的拼接都可能在计算前就阻塞整条链路。

2 分钟阅读 · 开发者指南

明确加载器必须交付什么

在优化之前先写好输出契约:元素数量、每个字段的类型、维度、目标范围以及不完整输入的处理规则。区分样本的标识符与它在 batch 中的位置。一个变换可能改变形状或过滤掉某个输入;训练程序必须知道这是否被允许。

取一个具有代表性的样本,其中包含一个普通文件、一个边界情况以及数据集的最后一个元素。用与 Dataset 完全相同的预处理方式打开每一个元素。然后检查它们的拼接。单个访问成功并不证明多个结果可以被堆叠。对于文本,要记录 padding 和 mask;对于图像,要记录通道数、维度以及轴的顺序。

固定范围:数据在本地还是远程、是否包含解码、变换是固定还是随机。保持它在两种设置之间;表面上的收益可能来自被删掉的工作。

退回单进程以读取错误

先用 num_workers=0、shuffle=False 和小 batch 复现。此时加载在主进程中执行,错误堆栈通常更易读。DataLoader 文档建议用这种方式调试。在解码之前记录失败元素的标识符,但不要把其敏感内容抄进日志。

逐步分离:原始访问、变换、collate_fn,然后是传输。如果流程在传输之前就失败,改 CUDA 并不是首要方向。如果只在多个 workers 下阻塞,就检查传给这些进程的对象和资源。比较第一次迭代和后续迭代:worker 启动可以解释初始等待,但不能据此断定问题会持续存在。

timeout 可以让等待变得可见,但既不能修复不可用的数据源,也不能修复阻塞的 worker。保留最后已知的步骤,并减少输入数量,而不是无限增大这个超时值。

完整示例:期望三个通道,却出现一张不同的图像

考虑四条教学记录。前三条给出形状为 [3, 16, 16] 的张量,第四条为 [1, 16, 16]。在要求三个通道的契约下,必须在堆叠之前识别出第四个元素。此场景并未在此处实际运行;它描述的是基于所选形状的预期结果。

下面的函数假定每条记录都有字段 id、x 和 y,x 是 CPU 张量,y 是整数索引。它会拒绝不一致,而不是悄悄丢弃该图像。对于你的项目,要明确决定单色图像应当转换为三个通道,还是在导入时被拒绝。这个决定取决于数据的含义以及模型期望的预处理。

修正之后,四个标识符都应保留,拼接后的张量形状应为 [4, 3, 16, 16]。再为目标添加一道防护:一张尺寸正确的图像仍可能带有无效标注。

建议的教学拼接示例,未执行
import torch
from torch.utils.data import DataLoader

def assemble(records):
    for item in records:
        if tuple(item["x"].shape) != (3, 16, 16):
            raise ValueError(f"Forme inattendue pour {item['id']}")
    return {
        "ids": [item["id"] for item in records],
        "x": torch.stack([item["x"] for item in records]),
        "y": torch.tensor([item["y"] for item in records],
                          dtype=torch.long),
    }

# dataset 是生成所描述记录的 Dataset。
# 在多进程脚本中,请在 main 保护下创建 loader。
if __name__ == "__main__":
    loader = DataLoader(dataset, batch_size=4, num_workers=0,
                        shuffle=False, collate_fn=assemble)
    iterator = iter(loader)
    batch = next(iterator)

在不改变数据的前提下重新引入 workers

从零逐步增加少量 workers,同时保持 batch、顺序和变换不变。先完整跑一个 epoch,再跑第二个:某些错误只在迭代器重启或资源已被消耗时才出现。只有准备工作能够真正并行推进时,提高并行度才有意义。

启动方法取决于操作系统和 Python 版本。使用 spawn 时,用 if __name__ == '__main__' 保护程序入口,并将 Dataset、collate_fn 和 worker 函数定义在模块级别,而不是局部的 lambda 中。进程相关文档也解释了为什么继承的锁或线程可能导致死锁。当库有要求时,保持访问的初始化对每个进程独立。

对于 IterableDataset,通过标识符检查 worker 之间的分区:多个 worker 不应各自消费完全相同的流。不要只评估 batch 数量,还要查找重复和缺失的元素。

用清晰单位衡量等待时间与吞吐量

使用两种互补的观测方式。单独遍历数据加载器,统计在指定时间间隔内准备好的样本数。集成遍历则考察模型消费这些数据时发生的情况。前者有助于隔离准备工作,但它并不自动代表训练吞吐量。

在你的测试方案中,统计实际交付的样本数,然后除以经过的秒数。说明哪些轮次被排除(用于启动)、数据缓存、变换以及重复次数。保留每一轮的值,而不是只挑选最好的一次。下表是一张记录表:没有任何性能数值被预先填入。

如果形状会变化,每秒样本数可能掩盖负载的变化。加入相关单位,例如解码的像素数或实际准备好的 token 数,同时保留样本数。为了在完整循环中定位等待,把读取下一个 batch 与计算分别命名。

左右滚动表格即可查看所有列。
用你自己的负载填写记录表,不预设任何性能数值
配置检查项观测时长预期结论
workers=0标识符、形状、目标以秒为单位测量正确的基准
少量 workers相同的输入集合以秒为单位测量实际收益或额外开销
相同配置,第二个 epoch无丢失或重复以秒为单位测量启动与缓存的影响

分别处理内存、预取和传输

workers 和待处理的 batch 会消耗主机内存。在试验期间持续监控内存,不要轻易认定只有 VRAM 才重要。更深的预取可能把等待转移到别处,同时增加内存占用;它并不能保证每秒产出更多结果。先缩小可疑变量,并在相同范围内进行比较。

pin_memory 和非阻塞传输涉及数据向加速器的传输。PyTorch 优化指南将其视为需结合硬件和负载考察的手段。它们无法修正错误的解码。先在 workers 中使用 CPU 数据,然后在驱动计算的主进程中组织传输。实际收益与重叠效果必须经过观测,而不能凭空假设。

若使用 persistent_workers,请注意两个 epoch 之间保留的资源和状态。仅在单个 batch 上表现良好的设置,不足以验证文件是否正确关闭或数据源是否已更新。

只有在数据仍然正确时才接受某个设置

预期结果是一个循环,在既定范围内接收所有预定输入,且没有静默错误。比较优化前后的标识符与目标值。若舍弃最后一个不完整的 batch,请说明 drop_last。若变换是随机的,应检查其策略,而不是要求像素完全相等,因为那与该策略相矛盾。

保留能满足实测需求的最简设置。如果存储、解码或模型本身已经是瓶颈,增加 worker 数量可能毫无改善。服务器的 CPU、内存和存储资源不能仅凭其 GPU 名称推断:在准备 Kernodeck 环境时,请单独明确这些需求。

您的疑问

num_workers=0 会禁用 GPU 训练吗?

不会。它只是把数据加载放到主进程中。模型仍然可以在 GPU 上计算。该设置的主要作用是更直接地读取加载、变换和组装过程中的错误。

worker 数量应该和 CPU 核心数一样吗?

并不自动如此。合适的设置取决于预处理工作、内存、数据访问以及模型的消费速度。在保持相同负载并检查所交付输入的前提下,比较几个数值。

更小的 batch 能修复停止的 worker 吗?

它可能改变内存压力,但无法修复无效的标注、无法序列化的资源或无法读取的文件。请先用零个 worker 复现,以确定出问题的环节。

next(iterator) 所花的时间能衡量磁盘吗?

不能。在 iterator=iter(loader) 之后,next(iterator) 只是在等待一个 batch。解码、变换、组装、进程间通信和预取都可能参与其中。单独的读取测量与完整循环的追踪回答的是不同的问题。