LeRobotはロボット学習データの主要なオープンフォーマットとして台頭していますが、そのデータ処理は容易ではありません。フレームのデコードはコストが高くメモリを消費し、GPUがアイドル状態である間、デコード、変換、アノテーションの各ステップがパイプラインのボトルネックになりがちです。そこで、私たちはDaftにネイティブなLeRobotリーダー(Daft #7090)を最近導入しました。daft.datasets.lerobotを使用すると、Hugging Faceから直接データセットをデータフレームとして読み取り、1行が1フレームに対応し、load_video_framesで各カメラのビデオを画像列にデコードします。
しかし、初期バージョンは非常に遅かった。問題は、フレームごとにリモートファイルを開くことでした。LeRobot v3データセットは、各カメラのビデオをMP4シャード(複数のエピソードのフレームを連続してパックしたファイル)として保存します。フレームをデコードするには、リーダーがシャードを開き、そのフレームのタイムスタンプにシークします。MP4を開くには、シークの前にインデックス(タイムスタンプをファイル内のバイト位置にマッピングするメタデータ)も読み取る必要があります。元のリーダーはこれをフレームごとに行っていました。つまり、各フレームがシャードを再オープンし、毎回ネットワーク経由でインデックスを再読み取りしていました。リモートデータセットでは、これが1フレームあたり約3秒かかり、フレーム数に比例して総コストが増大しました。連続する行が同じファイルの隣接フレームを要求する場合でも同様です。
修正方法は、バッチデコードです。デコードはバッチUDF(Daft #7184)になりました。関数は1行ごとではなく、16行ずつまとめて受け取るため、それらを横断してデコードを計画できます。各バッチでは3つのことを行います。まず、行をシャードごとにグループ化します。次に、ターゲットタイムスタンプをソートし、10秒以内のものをクラスタリングします。最後に、各クラスタに対して1回だけシークし、前方デコードを行います。クラスタ内の最も早いターゲットの前のキーフレームにシークし、その後デコードを続け、各ターゲットに最も近いフレームを保持し、最後のターゲットを超えるまで続けます。
結果は印象的です。リモートデータセットからの8フレームのデコードが25秒から3.9秒に短縮され、コスト曲線は線形からフラットになりました。6つの多様な公開LeRobot v3データセットでは、バッチリーダーが4〜13倍高速化しました。1080pデータセットで632フレームすべてをデコードする場合、29分から2分未満に短縮され、15倍の高速化を達成しました。さらに、このリーダー上に構築したハンドトラッキングパイプラインでは、12のリモートフレームのデコードとMediaPipeによるハンドトラッキングの実行が、エンドツーエンドで44.8秒から9.8秒に短縮され、検出結果は同一でした。
試すには、数行のコードが必要です。daft.datasets.lerobotを使用してデータセットを読み取り、ロードするカメラを指定します。このリーダーは、表示された行のデコードに1回のシャードオープンしか必要としません。完全なアノテーションパイプラインについてはdaft-physical-aiを、ベンチマーク結果についてはDaftリポジトリのbenchmarking/lerobotを参照してください。