LeRobotビデオリーダーを最大15倍高速化した方法

ブログに戻る RobotデータがLeRobotに収束しています。 LeRobotはロボット学習データの主要なオープンフォーマットとして台頭しています。しかし、それに対するデータ操作は容易ではありません。フレームのデコードはコストが高くメモリを大量に消費し、GPUの前のすべてのステップ（デコード、変換、フレームのアノテーション）でパイプラインが詰まり、GPUは推論またはトレーニングの実行中であってもアイドル状態になります。私たちはその部分を可能な限り容易にするために取り組んできました。その一環として、最近Daft（Daft #7090）にネイティブなLeRobotリーダーを導入しました。daft.datasets.lerobotは、Hugging Faceからデータセットを直接読み込み、フレームごとに1行のデータフレームに変換します。load_video_framesを使用すると、各カメラを画像列にデコードします。しかし、初期バージョンは非常に遅かったです。 問題：フレームごとのリモートオープン LeRobot v3データセットは、各カメラのビデオをMP4シャードとして保存します。これは、多くのエピソードのフレームを連続してパックしたファイルです。フレームをデコードするために、リーダーはシャードを開き、フレームのタイムスタンプにシークします。MP4を開くということは、シークが可能になる前に、タイムスタンプとファイル内のバイト位置をマッピングするメタデータであるインデックスを読み取ることも意味します。元のリーダーはこれらすべてを各行に対して行っていました。各フレームはシャードを再オープンし、各オープンでネットワーク経由でインデックスを再読み込みしていました。リモートデータセットでは、これはフレームあたり約3秒かかり、総コストはフレーム数に比例して増加しました。たとえ連続する行が同じファイルから隣接するフレームを要求していたとしてもです。 修正：シャードごとのバッチデコード デコードは現在バッチUDF（Daft #7184）になっています。行ごとに1回呼び出されるのではなく、関数は一度に16個の連続する行を受け取るため、それらを横断してデコードを計画できます。各バッチについて、3つのことを行います。 1. シャードごとにターゲットをグループ化します。 バッチの1回のパスで、ターゲットをそれが指すシャードごとにグループ化します。各行のターゲット時間は、シャード内のエピソードの開始オフセットに、エピソード内のフレームのタイムスタンプを加えたものです。結果はシャードごとに1つのリストになり、そのシャードがサービスする必要がある行インデックスとターゲットタイムを保持します。 by_shard = {} 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秒以内であれば現在のクラスタに参加し、そうでなければ新しいクラスタを開始します。 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) このギャップは何のためでしょうか？シークはデコードを先行するキーフレームから再開するため、短いギャップをそのままデコードする方が、再シークするよりも安価です。しかし、シャードは多くのエピソードを連続してパックしており、バッチ内の2つのターゲットは数分離れている可能性があります。それほど長いギャップをデコードすると無駄が生じるため、しきい値を超えるものは独自のクラスタと独自のシークを取得します。 3. クラスタごとに1回のシークと1回のフォワードパスを行います。 各クラスタについて、デコーダーは最も早いターゲットに先行するキーフレームに一度シークし、そのまま読み進めます。デコードされた各フレームはクラスタのターゲットと比較され、これまでに見た最も近いフレームが各ターゲット用に保持され、最後のターゲットを過ぎるとパスは停止します。 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秒に短縮され、コスト曲線は線形からフラットになりました。 多様性（av1/h264/mp4v、5-30 fps、128×128から1280×720、1-3カメラ）のために選択された6つの公開LeRobot v3データセット全体で、バッチ処理されたリーダーは4〜13倍高速です。 実験をスケールアップします。1080pデータセット（pepijn223/egodex-test）では、632フレームすべてをデコードするのにかかる時間は29分から2分未満に短縮されました。これは15倍高速です。なぜなら、バッチ処理のコストはフレームではなくバッチごとに増加するからです。 さらに、このリーダー上に構築されたハンドトラッキングパイプラインでは、フレームに手のポーズをアノテーションする場合、12個のリモートフレームをデコードしてMediaPipeハンドトラッキングを実行するのに、エンドツーエンドで44.8秒かかっていたのが9.8秒に短縮されました（ベンチマーク）。 試してみてください。 import time from daft.datasets import lerobot # フレームごとに1行。カメラは画像列にデコードされます。 df = lerobot.read("pepijn223/egodex-test", load_video_frames="observation.image") # 表示される8行のみをデコードします - シャードは1回開かれます。 t0 = time.perf_counter() df.show() print(f"{time.perf_counter() - t0:.1f}s") # ネットワーク経由で約7秒 完全なアノテーションパイプラインについては、daft-physical-aiを参照してください。ベンチマークハーネスと完全な結果はDaftリポジトリのbenchmarking/lerobotにあります。