From ced3de4c94f4fa1996246095d973de42c037fa29 Mon Sep 17 00:00:00 2001 From: Remi Date: Sat, 11 May 2024 19:28:22 +0300 Subject: [PATCH] Fix hanging in visualize_dataset.py when num_workers > 0 (#165) --- lerobot/scripts/visualize_dataset.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/lerobot/scripts/visualize_dataset.py b/lerobot/scripts/visualize_dataset.py index 1835e90..d4fafe6 100644 --- a/lerobot/scripts/visualize_dataset.py +++ b/lerobot/scripts/visualize_dataset.py @@ -47,6 +47,7 @@ local$ rerun ws://localhost:9087 """ import argparse +import gc import logging import time from pathlib import Path @@ -115,15 +116,17 @@ def visualize_dataset( spawn_local_viewer = mode == "local" and not save rr.init(f"{repo_id}/episode_{episode_index}", spawn=spawn_local_viewer) + + # Manually call python garbage collector after `rr.init` to avoid hanging in a blocking flush + # when iterating on a dataloader with `num_workers` > 0 + # TODO(rcadene): remove `gc.collect` when rerun version 0.16 is out, which includes a fix + gc.collect() + if mode == "distant": rr.serve(open_browser=False, web_port=web_port, ws_port=ws_port) logging.info("Logging to Rerun") - if num_workers > 0: - # TODO(rcadene): fix data workers hanging when `rr.init` is called - logging.warning("If data loader is hanging, try `--num-workers 0`.") - for batch in tqdm.tqdm(dataloader, total=len(dataloader)): # iterate over the batch for i in range(len(batch["index"])): @@ -196,7 +199,7 @@ def main(): parser.add_argument( "--num-workers", type=int, - default=0, + default=4, help="Number of processes of Dataloader for loading the data.", ) parser.add_argument(