LeRobot 视频读取器提速 15 倍！解码成本大幅降低，效率显著提升-拓冰建站

机器人数据趋向采用 LeRobot 格式

LeRobot 已成为机器人学习数据的主流开放格式。不过，对其进行数据操作并非易事，解码帧成本高且内存占用大。在将数据送入 GPU 之前的每一步，如解码、转换、标注帧，都会使数据处理管道变慢，此时 GPU 只能闲置，无论进行推理还是训练。为让这一过程尽可能简单，最近在 Daft 中引入了原生 LeRobot 读取器。`daft.datasets.lerobot` 可直接从 Hugging Face 读取数据集，并将其转换为一个数据框，每帧占一行。借助 `load_video_frames`，它能将每个摄像头的视频解码为图像列。然而，初始版本的速度慢得令人痛苦。

问题所在：每帧都进行一次远程打开操作

LeRobot v3 数据集将每个摄像头的视频存储为 MP4 分片文件，这些文件会连续打包多段视频片段的帧。为了解码一帧，读取器需要打开分片文件，并定位到该帧的时间戳位置。在进行定位之前，打开 MP4 文件还意味着要读取其索引，即用于将时间戳映射到文件中字节位置的元数据。最初的读取器对每一行数据都进行上述操作，每一帧都重新打开其分片文件，每次打开都要通过网络重新读取索引。对于远程数据集，解码一帧大约需要 3 秒，而且总解码成本会随帧数线性增加，即使连续行需要的是同一文件中的相邻帧也是如此。

解决方案：按分片进行批量解码

现在，解码操作采用了批量 UDF。该函数不再逐行调用，而是每次接收 16 行连续的数据，这样就能对这些数据的解码操作进行规划。对于每个批次，它会执行以下三个步骤。

1. 按分片对行进行分组

对批次数据进行一次遍历，按行所指向的分片对其进行分组。每行的目标时间是该片段在分片中的起始偏移量加上该帧在片段内的时间戳。最终结果是每个分片对应一个列表，列表中包含该行的索引和该分片需要处理的目标时间。这样，每个分片只需打开一次，就能处理所有目标，而不是每帧都打开一次。

2. 对目标进行排序和聚类

在一个分片中，目标按时间戳升序排序，然后进行一次遍历。如果一个目标与前一个目标的时间间隔在 10 秒以内，则将其加入当前聚类；否则，创建一个新的聚类。设置这个时间间隔是因为定位操作会从前面的关键帧重新开始解码，所以连续解码一个小间隔内的帧比重新定位更高效。但一个分片中会连续打包多个片段，一个批次中的两个目标可能相隔数分钟，如果解码这么长的间隔会浪费计算资源，所以超过这个阈值的目标会被分到不同的聚类中，并进行单独的定位操作。

3. 每个聚类进行一次定位和一次正向遍历

对于每个聚类，解码器先定位到最早目标之前的关键帧，然后持续读取。将每个解码后的帧与该聚类的目标进行比较，为每个目标保留目前为止最接近的帧，一旦超过最后一个目标，遍历就停止。这个改进仅涉及 Python 代码，输出结果与旧的逐行解码方式在字节层面完全相同。

实验结果

从远程数据集解码 8 帧的时间从 25 秒降至 3.9 秒，成本曲线从线性增长变为趋于平稳。在六个具有多样性的公开 LeRobot v3 数据集上，批量读取器的速度提升了 4 - 13 倍。扩大实验规模，在一个 1080p 数据集上，解码所有 632 帧的时间从 29 分钟降至不到 2 分钟，提速达 15 倍，这是因为批量解码的成本随批次数量增长，而非帧数。此外，在基于该读取器构建的手部跟踪管道中，即对帧进行手部姿势标注，解码 12 个远程帧并运行 MediaPipe 手部跟踪的端到端时间从 44.8 秒降至 9.8 秒，检测结果相同。

立即尝试

可使用以下代码尝试：

import time from daft.datasets import lerobot # 每帧占一行；摄像头视频被解码为图像列。 df = lerobot.read("pepijn223/egodex-test", load_video_frames="observation.image") # 仅解码显示的 8 行数据，只需打开一次分片文件。 t0 = time.perf_counter() df.show() print(f"{time.perf_counter() - t0:.1f}s") # 通过网络读取约需 7 秒

关于基于此读取器构建的完整标注管道，可参阅相关内容。基准测试工具和完整结果可在 Daft 仓库的对应目录中找到。