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_framesobservation.image) # 仅解码显示的 8 行数据只需打开一次分片文件。 t0 time.perf_counter() df.show() print(f{time.perf_counter() - t0:.1f}s) # 通过网络读取约需 7 秒关于基于此读取器构建的完整标注管道可参阅相关内容。基准测试工具和完整结果可在 Daft 仓库的对应目录中找到。推荐文章使用 daft - physical - ai 将机器人视频转换为可用于训练的数据这是一个基于 Daft 构建的 Python 库用于将机器人视频转换为可用于训练的数据最初支持将手部跟踪和奖励评分作为 UDF未来还会有更多功能。Daft v0.7.17 支持 LeRobot 数据集、HDF5 文件和本地大语言模型推理该版本提供了 daft.datasets.lerobot 用于处理 LeRobot v3 机器人数据支持原生 HDF5 文件并可通过 Transformers 提供程序进行本地大语言模型推理。