tangger/lerobot
1
0
Fork 1
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Speedup data loading

Browse Source
This commit is contained in:
Remi Cadene
2025-05-06 15:13:50 +00:00
parent 588bf96559
commit 0309a9fcbc
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
lerobot/common/datasets/lerobot_dataset.py
View File

@@ -736,7 +736,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
query_timestamps = {} query_timestamps = {}
for key in self.meta.video_keys: for key in self.meta.video_keys:
if query_indices is not None and key in query_indices: if query_indices is not None and key in query_indices:
timestamps = self.hf_dataset.select(query_indices[key])["timestamp"] timestamps = self.hf_dataset[query_indices[key]]["timestamp"]
query_timestamps[key] = torch.stack(timestamps).tolist() query_timestamps[key] = torch.stack(timestamps).tolist()
else: else:
query_timestamps[key] = [current_ts] query_timestamps[key] = [current_ts]
@@ -745,7 +745,7 @@ class LeRobotDataset(torch.utils.data.Dataset):
def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict: def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
return { return {
key: torch.stack(self.hf_dataset.select(q_idx)[key]) key: torch.stack(self.hf_dataset[q_idx][key])
for key, q_idx in query_indices.items() for key, q_idx in query_indices.items()
if key not in self.meta.video_keys if key not in self.meta.video_keys
} }