mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-06 08:29:46 +00:00
feat(Train): enable buckets with streaming dataset (#4312)
This commit is contained in:
@@ -142,6 +142,22 @@ repo_id = "yaak-ai/L2D-v3"
|
|||||||
dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub
|
dataset = StreamingLeRobotDataset(repo_id) # streams directly from the Hub
|
||||||
```
|
```
|
||||||
|
|
||||||
|
Datasets stored in an [HF Storage Bucket](https://huggingface.co/docs/hub/storage-buckets) (`hf://buckets/`) can be streamed the same way by passing `repo_type="bucket"`:
|
||||||
|
|
||||||
|
```python
|
||||||
|
dataset = StreamingLeRobotDataset("my-org/my-bucket", repo_type="bucket")
|
||||||
|
```
|
||||||
|
|
||||||
|
Both options are available in `lerobot-train` through `--dataset.streaming=true`, and `--dataset.repo_type=bucket` to stream from a bucket instead of a Hub dataset repo:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
lerobot-train \
|
||||||
|
--dataset.repo_id=my-org/my-bucket \
|
||||||
|
--dataset.repo_type=bucket \
|
||||||
|
--dataset.streaming=true \
|
||||||
|
...
|
||||||
|
```
|
||||||
|
|
||||||
<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
|
<div style="display:flex; justify-content:center; gap:12px; flex-wrap:wrap;">
|
||||||
<figure style="margin:0; text-align:center;">
|
<figure style="margin:0; text-align:center;">
|
||||||
<img
|
<img
|
||||||
|
|||||||
@@ -29,6 +29,9 @@ class DatasetConfig:
|
|||||||
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
|
# "dataset_index" into the returned item. The index mapping is made according to the order in which the
|
||||||
# datasets are provided.
|
# datasets are provided.
|
||||||
repo_id: str
|
repo_id: str
|
||||||
|
# Hub repository type: "dataset" (default) or "bucket" for an HF Storage Bucket streamed over
|
||||||
|
# hf://buckets/. Buckets are streaming-only, so "bucket" requires streaming=true.
|
||||||
|
repo_type: str = "dataset"
|
||||||
# Root directory for a concrete local dataset tree (e.g. 'dataset/path'). If None, local datasets are
|
# Root directory for a concrete local dataset tree (e.g. 'dataset/path'). If None, local datasets are
|
||||||
# looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
|
# looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
|
||||||
root: str | None = None
|
root: str | None = None
|
||||||
@@ -48,6 +51,16 @@ class DatasetConfig:
|
|||||||
eval_split: float = 0.0
|
eval_split: float = 0.0
|
||||||
|
|
||||||
def __post_init__(self) -> None:
|
def __post_init__(self) -> None:
|
||||||
|
if self.repo_type not in ("dataset", "bucket"):
|
||||||
|
raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {self.repo_type!r}")
|
||||||
|
if self.repo_type == "bucket" and not self.streaming:
|
||||||
|
raise ValueError(
|
||||||
|
"repo_type='bucket' is streaming-only: set streaming=true to train from an HF Storage Bucket."
|
||||||
|
)
|
||||||
|
if self.repo_type == "bucket" and self.eval_split != 0.0:
|
||||||
|
raise ValueError(
|
||||||
|
"eval_split requires map-style datasets and is not supported with repo_type='bucket'."
|
||||||
|
)
|
||||||
if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
|
if self.depth_output_unit not in (DEPTH_METER_UNIT, DEPTH_MILLIMETER_UNIT):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}"
|
f"depth_output_unit must be '{DEPTH_METER_UNIT}' or '{DEPTH_MILLIMETER_UNIT}', got {self.depth_output_unit!r}"
|
||||||
|
|||||||
@@ -84,10 +84,17 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
|||||||
|
|
||||||
if isinstance(cfg.dataset.repo_id, str):
|
if isinstance(cfg.dataset.repo_id, str):
|
||||||
ds_meta = LeRobotDatasetMetadata(
|
ds_meta = LeRobotDatasetMetadata(
|
||||||
cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
|
cfg.dataset.repo_id,
|
||||||
|
root=cfg.dataset.root,
|
||||||
|
revision=cfg.dataset.revision,
|
||||||
|
repo_type=cfg.dataset.repo_type,
|
||||||
)
|
)
|
||||||
delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
|
delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
|
||||||
if not cfg.dataset.streaming:
|
if not cfg.dataset.streaming:
|
||||||
|
if cfg.dataset.repo_type == "bucket":
|
||||||
|
raise ValueError(
|
||||||
|
"repo_type='bucket' is streaming-only: set dataset.streaming=true to train from an HF Storage Bucket."
|
||||||
|
)
|
||||||
dataset = LeRobotDataset(
|
dataset = LeRobotDataset(
|
||||||
cfg.dataset.repo_id,
|
cfg.dataset.repo_id,
|
||||||
root=cfg.dataset.root,
|
root=cfg.dataset.root,
|
||||||
@@ -111,6 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
|||||||
max_num_shards=cfg.num_workers,
|
max_num_shards=cfg.num_workers,
|
||||||
tolerance_s=cfg.tolerance_s,
|
tolerance_s=cfg.tolerance_s,
|
||||||
return_uint8=True,
|
return_uint8=True,
|
||||||
|
repo_type=cfg.dataset.repo_type,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
|
raise NotImplementedError("The MultiLeRobotDataset isn't supported for now.")
|
||||||
|
|||||||
@@ -36,3 +36,22 @@ def test_dataset_config_none_episodes_ok():
|
|||||||
|
|
||||||
def test_dataset_config_empty_episodes_ok():
|
def test_dataset_config_empty_episodes_ok():
|
||||||
DatasetConfig(repo_id="user/repo", episodes=[])
|
DatasetConfig(repo_id="user/repo", episodes=[])
|
||||||
|
|
||||||
|
|
||||||
|
def test_dataset_config_bucket_streaming_ok():
|
||||||
|
DatasetConfig(repo_id="user/repo", repo_type="bucket", streaming=True)
|
||||||
|
|
||||||
|
|
||||||
|
def test_dataset_config_invalid_repo_type():
|
||||||
|
with pytest.raises(ValueError, match="repo_type"):
|
||||||
|
DatasetConfig(repo_id="user/repo", repo_type="model")
|
||||||
|
|
||||||
|
|
||||||
|
def test_dataset_config_bucket_requires_streaming():
|
||||||
|
with pytest.raises(ValueError, match="streaming-only"):
|
||||||
|
DatasetConfig(repo_id="user/repo", repo_type="bucket")
|
||||||
|
|
||||||
|
|
||||||
|
def test_dataset_config_bucket_rejects_eval_split():
|
||||||
|
with pytest.raises(ValueError, match="eval_split"):
|
||||||
|
DatasetConfig(repo_id="user/repo", repo_type="bucket", streaming=True, eval_split=0.1)
|
||||||
|
|||||||
Reference in New Issue
Block a user