如何让我们的LeRobot视频读取器速度提升最高15倍

LeRobot已成为机器人学习数据的主要开放格式，但处理其数据并不容易：解码帧既昂贵又占用内存，而在GPU闲置时，解码、转换和注释帧的步骤常常成为流水线的瓶颈。为此，我们最近在Daft中引入了原生LeRobot读取器（Daft #7090），通过daft.datasets.lerobot可直接从Hugging Face读取数据集为数据帧，每行一帧，并使用load_video_frames将每个摄像头的视频解码为图像列。

然而，初始版本的速度非常慢。问题在于：每个帧都要远程打开一次文件。LeRobot v3数据集将每个摄像头的视频存储为MP4分片——这些文件将多个片段的帧打包在一起。解码一个帧时，读取器需要打开分片并寻址到该帧的时间戳。打开MP4还需要读取其索引（即映射时间戳到字节位置的元数据），然后才能进行寻址。原始读取器逐帧执行此操作：每个帧都重新打开其分片，每次打开都通过网络重新读取索引。对于远程数据集，这大约需要每帧3秒，而且总成本随帧数线性增长——即使连续行需要来自同一文件的相邻帧。

修复方法是批量解码。现在解码是一个批量UDF（Daft #7184）：函数一次接收16个连续行，而不是逐行调用，从而可以跨行规划解码。每个批次执行三个步骤：首先，按分片分组行；其次，排序和聚类目标时间戳（10秒内的目标归为一簇）；最后，每个聚类只进行一次寻址和一次前向解码。寻址到聚类中最早目标的前一个关键帧，然后持续解码，将每个解码帧与聚类中的目标比较，保留最接近的帧，直到超过最后一个目标为止。

结果令人印象深刻：从远程数据集解码8帧从25秒降至3.9秒，成本曲线从线性变为平坦。在六个公共LeRobot v3数据集上，批量读取器速度提升4-13倍。在1080p数据集上，解码全部632帧从29分钟降至不到2分钟——提升了15倍。此外，在一个基于此读取器的手部追踪流水线中，解码12个远程帧并运行MediaPipe手部追踪，端到端时间从44.8秒降至9.8秒，检测结果完全一致。

要尝试此功能，只需几行代码。使用daft.datasets.lerobot读取数据集，并指定加载视频帧的摄像头。该读取器仅需一次分片打开即可解码显示的行。完整的注释流水线可参考daft-physical-ai，基准测试结果可在Daft仓库的benchmarking/lerobot中找到。