如何将 LeRobot 视频读取器的速度提升 15 倍

# 我们如何将 LeRobot 视频读取器速度提升至 15 倍 来源：https://www.eventual.ai/blog/how-we-made-our-lerobot-video-reader-up-to-15x-faster ## 机器人数据正在向 LeRobot 收敛 LeRobot 已成为机器人学习数据的主流开放格式。但对其执行数据操作并不容易：解码帧成本高昂且消耗大量内存，而且在将数据送入 GPU 之前的每一步——解码、变换、注释帧——都是流水线瓶颈所在，导致 GPU 闲置，无论您是在进行推理还是训练。 我们一直在努力让这一环节尽可能简便。作为其中的一部分，我们最近在 Daft 中引入了一个原生 LeRobot 读取器（Daft #7090 (https://github.com/Eventual-Inc/Daft/pull/7090)）。`daft.datasets.lerobot` 直接从 Hugging Face 读取数据集到 DataFrame 中，每一行对应一帧。通过 `load_video_frames`，它将每个摄像头解码为一个图像列。 然而，初始版本的运行速度非常慢。 ## 问题：每帧一次远程打开 LeRobot v3 数据集将每个摄像头的视频存储为 MP4 分片——这些文件将多个 episodes 的帧连续打包在一起。要解码一帧，读取器需要打开分片并寻址到该帧的时间戳。打开 MP4 还需要读取其索引——即映射时间戳到文件中字节位置的元数据——然后才能进行任何寻址操作。 原始读取器对每一行都执行所有这些操作：每一帧都重新打开它的分片，每次打开都会通过网络重新读取索引。对于远程数据集，这大约需要每帧 3 秒，而且总成本随帧数线性增长——即使连续的行请求来自同一文件的相邻帧。 ## 修复：按分片批量解码 现在解码是一个批量 UDF（https://docs.daft.ai/en/stable/custom-code/func/#batch-udfs-with-daftfuncbatch）（Daft #7184 (https://github.com/Eventual-Inc/Daft/pull/7184)）：该函数不再是每行调用一次，而是每次接收 16 个连续行，从而可以跨行规划解码。对于每个批次，它执行三个步骤。 **1. 按分片对行进行分组。** 对批次进行一次遍历，将行按其指向的分片分组。每一行的目标时间就是其 episode 在分片内的起始偏移加上该帧在 episode 内的时间戳。结果是每个分片对应一个列表，包含该分片需要服务的行索引和目标时间： ```python by_shard = {} # 遍历每一行；每个 `file` 是对该行分片的引用。 for i, file in enumerate(files): abs_ts = from_timestamp[i] + frame_timestamp[i] by_shard.setdefault(file.path, []).append((i, abs_ts)) ``` 然后每个分片只打开一次，服务于其所有目标，而非每帧打开一次。 **2. 对目标进行排序和聚类。** 在一个分片内，目标按时间戳升序排序，然后进行一次遍历：如果目标与前一个目标的时间差在 10 秒以内，则加入当前聚类；否则开始一个新的聚类： ```python targets = sorted(abs_timestamps, key=lambda t: t[1]) # 按时间戳升序 clusters = [[targets[0]]] for t in targets[1:]: if t[1] - clusters[-1][-1][1] > 10.0: # 秒 clusters.append([t]) else: clusters[-1].append(t) ``` 为什么设置这个时间差？一次寻址操作会从前面的关键帧开始重新解码，因此直接解码一个小的间隙比分次寻址更便宜。但一个分片会连续打包许多 episodes，批次中的两个目标可能相隔几分钟——解码如此大的间隙会浪费工作，因此超过阈值的部分会建立自己的聚类并执行自己的寻址。 **3. 每个聚类一次寻址和一次正向遍历。** 对于每个聚类，解码器只寻址一次到最早目标之前的关键帧，然后持续读取：将每个解码帧与聚类的目标进行比较，为每个目标保留当前看到的最接近的帧，并在超过最后一个目标后停止： ```python container.seek(earliest_target_pts, backward=True) # 前一个关键帧 for frame in container.decode(stream): ts = float(frame.pts * stream.time_base) for row, target in cluster: # 如果这是目前看到的最接近 `target` 的帧，则保留 ... if ts >= latest_target + tail: break ``` 该更改仅涉及 Python 代码，输出与旧的逐行解码完全相同（字节级一致）。 ## 结果 从远程数据集中解码 8 帧，时间从 25 秒降至 3.9 秒，成本曲线从线性变为平坦： 原始 vs 批量 在六个具有代表性的公共 LeRobot v3 数据集上（av1/h264/mp4v，5-30 fps，128×128 到 1280×720，1-3 个摄像头），批量读取器速度提升 4-13 倍： 原始 vs 批量（公共数据集） 扩大实验规模：在一个 1080p 数据集（pepijn223/egodex-test (https://huggingface.co/datasets/pepijn223/egodex-test)）上，解码全部 632 帧从 **29 分钟降至不到 2 分钟——提升 15 倍**——原因在于批量成本随批次数量增长而非帧数增长： 原始 vs 批量（大规模） 此外，在一个基于此读取器构建的手部追踪流水线（用于注释带有手部姿态的帧）中，解码 12 帧远程数据并运行 MediaPipe 手部追踪，端到端时间从 44.8 秒降至 9.8 秒，且检测结果完全相同（基准 (https://github.com/Eventual-Inc/Daft/pull/7267)）： 原始 vs 批量（手部追踪工作负载） ## 尝试使用 ```python 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") # 通过网络约 7s ``` 如需在此读取器之上构建完整的注释流水线，请参阅 daft-physical-ai (https://github.com/Eventual-Inc/daft-physical-ai)。基准测试框架和完整结果位于 Daft 仓库中：`benchmarking/lerobot` (https://github.com/Eventual-Inc/Daft/tree/main/benchmarking/lerobot)。