admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-19 07:51:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Optimize dataset updates by incrementally concatenating new data instead of reloading from disk, reducing memory usage and improving performance.

Browse Source
This commit is contained in:
Michel Aractingi
2025-09-05 18:37:48 +02:00
parent 992fb177c3
commit 0747afdba7
1 changed files with 14 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/datasets/lerobot_dataset.py
+14 -23
View File
@@ -28,7 +28,7 @@ import pandas as pd
import PIL.Image import PIL.Image
import torch import torch
import torch.utils import torch.utils
from datasets import Dataset from datasets import Dataset, concatenate_datasets
from huggingface_hub import HfApi, snapshot_download from huggingface_hub import HfApi, snapshot_download
from huggingface_hub.constants import REPOCARD_NAME from huggingface_hub.constants import REPOCARD_NAME
from huggingface_hub.errors import RevisionNotFoundError from huggingface_hub.errors import RevisionNotFoundError
@@ -316,17 +316,16 @@ class LeRobotDatasetMetadata:
path.parent.mkdir(parents=True, exist_ok=True) path.parent.mkdir(parents=True, exist_ok=True)
df.to_parquet(path, index=False) df.to_parquet(path, index=False)
# Update the Hugging Face dataset by reloading it. # Update the Hugging Face dataset incrementally instead of reloading from disk
# This process should be fast because only the latest Parquet file has been modified. # This eliminates repeated load_episodes calls that cause cache bloat
# Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache. if self.episodes is None:
self.episodes = load_episodes(self.root)
return
# Explicitly delete old dataset to free memory before reloading # Remove columns from df that start with 'stats/'
if hasattr(self, "episodes") and self.episodes is not None: df = df.drop(columns=[col for col in df.columns if col.startswith("stats/")])
del self.episodes new_episode_dataset = Dataset.from_pandas(df)
self.episodes = None self.episodes = concatenate_datasets([self.episodes, new_episode_dataset])
gc.collect()
self.episodes = load_episodes(self.root)
def save_episode( def save_episode(
self, self,
@@ -1064,17 +1063,8 @@ class LeRobotDataset(torch.utils.data.Dataset):
else: else:
df.to_parquet(path) df.to_parquet(path)
# Update the Hugging Face dataset by reloading it. new_hf_dataset = Dataset.from_pandas(df)
# This process should be fast because only the latest Parquet file has been modified. self.hf_dataset = concatenate_datasets([self.hf_dataset, new_hf_dataset])
# Therefore, only this file needs to be converted to PyArrow; the rest is loaded from the PyArrow memory-mapped cache.
# Explicitly delete old dataset to free memory before reloading
if hasattr(self, "hf_dataset") and self.hf_dataset is not None:
del self.hf_dataset
self.hf_dataset = None
gc.collect()
self.hf_dataset = self.load_hf_dataset()
metadata = { metadata = {
"data/chunk_index": chunk_idx, "data/chunk_index": chunk_idx,
@@ -1093,7 +1083,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
if self.meta.episodes is None: if self.meta.episodes is None:
# Initialize indices for a new dataset made of the first episode data # Initialize indices for a new dataset made of the first episode data
chunk_idx, file_idx = 0, 0 chunk_idx, file_idx = 0, 0
latest_duration_in_s = 0 latest_duration_in_s = 0.0
new_path = self.root / self.meta.video_path.format( new_path = self.root / self.meta.video_path.format(
video_key=video_key, chunk_index=chunk_idx, file_index=file_idx video_key=video_key, chunk_index=chunk_idx, file_index=file_idx
) )
@@ -1119,6 +1109,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
) )
new_path.parent.mkdir(parents=True, exist_ok=True) new_path.parent.mkdir(parents=True, exist_ok=True)
shutil.move(str(ep_path), str(new_path)) shutil.move(str(ep_path), str(new_path))
latest_duration_in_s = 0.0
else: else:
# Update latest video file # Update latest video file
concat_video_files([latest_path, ep_path], self.root, video_key, chunk_idx, file_idx) concat_video_files([latest_path, ep_path], self.root, video_key, chunk_idx, file_idx)