hangX/lerobot
1
0
Fork 0
forked from tangger/lerobot
Code Pull Requests Activity

fix issue with saving freshly computed stats

Browse Source
This commit is contained in:
Alexander Soare
2024-04-17 08:49:28 +01:00
parent 3f1c322d56
commit 1331d3b4e4
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
lerobot/common/datasets/factory.py
View File

@@ -62,7 +62,7 @@ def make_dataset(
stats = torch.load(precomputed_stats_path) stats = torch.load(precomputed_stats_path)
else: else:
logging.info(f"compute_stats and save to {precomputed_stats_path}") logging.info(f"compute_stats and save to {precomputed_stats_path}")
# instantiate a one frame dataset with light transform # Create a dataset for stats computation.
stats_dataset = clsfunc( stats_dataset = clsfunc(
dataset_id=cfg.dataset_id, dataset_id=cfg.dataset_id,
split="train", split="train",
@@ -70,7 +70,8 @@ def make_dataset(
transform=Prod(in_keys=clsfunc.image_keys, prod=1 / 255.0), transform=Prod(in_keys=clsfunc.image_keys, prod=1 / 255.0),
) )
stats = compute_stats(stats_dataset) stats = compute_stats(stats_dataset)
torch.save(stats, stats_path) os.makedirs(precomputed_stats_path.parent, exist_ok=True)
torch.save(stats, precomputed_stats_path)
else: else:
stats = torch.load(stats_path) stats = torch.load(stats_path)