明确加载器必须交付什么
在优化之前先写好输出契约:元素数量、每个字段的类型、维度、目标范围以及不完整输入的处理规则。区分样本的标识符与它在 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 环境时,请单独明确这些需求。