wangzixiang
/
lerobot
mirror of https://github.com/huggingface/lerobot.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

nit + num-image-writers 8

This commit is contained in:
Remi Cadene 2024-07-12 20:00:41 +02:00
parent d525e1b0f8
commit 30c7a9e6fc
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
lerobot/scripts/control_robot.py
View File

@ -209,7 +209,7 @@ def record_dataset(
num_episodes=50, num_episodes=50,
video=True, video=True,
run_compute_stats=True, run_compute_stats=True,
num_image_writters=4, num_image_writers=8,
force_override=False, force_override=False,
): ):
# TODO(rcadene): Add option to record logs # TODO(rcadene): Add option to record logs
@ -292,7 +292,7 @@ def record_dataset(
# Using `with` to exist smoothly if an execption is raised. # Using `with` to exist smoothly if an execption is raised.
# Using only 4 worker threads to avoid blocking the main thread. # Using only 4 worker threads to avoid blocking the main thread.
futures = [] futures = []
with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writters) as executor: with concurrent.futures.ThreadPoolExecutor(max_workers=num_image_writers) as executor:
# Start recording all episodes # Start recording all episodes
while episode_index < num_episodes: while episode_index < num_episodes:
logging.info(f"Recording episode {episode_index}") logging.info(f"Recording episode {episode_index}")
@ -406,7 +406,7 @@ def record_dataset(
if stop_recording or episode_index == num_episodes: if stop_recording or episode_index == num_episodes:
logging.info("Done recording") logging.info("Done recording")
os.system('say "Done recording"') os.system('say "Done recording"')
logging.info("Waiting for threads writting the images on disk to terminate...") logging.info("Waiting for threads writing the images on disk to terminate...")
listener.stop() listener.stop()
for _ in tqdm.tqdm( for _ in tqdm.tqdm(
concurrent.futures.as_completed(futures), total=len(futures), desc="Writting images" concurrent.futures.as_completed(futures), total=len(futures), desc="Writting images"
@ -615,10 +615,10 @@ if __name__ == "__main__":
) )
parser_record.add_argument( parser_record.add_argument(
"--num-image-writters", "--num-image-writers",
type=int, type=int,
default=4, default=8,
help="Number of threads writting the frames as png images on disk. Don't set too much as you might get unstable fps due to main thread being blocked.", help="Number of threads writing the frames as png images on disk. Don't set too much as you might get unstable fps due to main thread being blocked.",
) )
parser_record.add_argument( parser_record.add_argument(
"--force-override", "--force-override",