Skip to content

Conversation

@zhangnju
Copy link

What this does

Explain what this PR does. Feel free to tag your PR with the appropriate label(s).
This PR adds torch profiler in training/eval tools, which can help find the performance bottleneck of lerobot training/eval tool

How it was tested

Explain/show how you tested your changes.
add the option of "profiling" to enable this feature, and the trace json file will be generated in the current working folder

Training example
TOKENIZERS_PARALLELISM=false lerobot-train --policy.type=smolvla --policy.repo_id=xxxx/smolvla-test --dataset.repo_id=lerobot/svla_so100_stacking/ --batch_size=64 --steps=20 --policy.push_to_hub False --profiling True

Eval example:
TOKENIZERS_PARALLELISM=false lerobot-eval --policy.path=HuggingFaceVLA/smolvla_libero --env.type=libero --env.task=libero_object --eval.batch_size=1 --eval.n_episodes=1 --policy.device=cuda --profiling True

How to checkout & try? (for the reviewer)

Provide a simple way for the reviewer to try out your changes.

based on this PR, run "pip install .[all]" to install lerobot

run the below command to generate tracing files for SmolVLA model

Training example
TOKENIZERS_PARALLELISM=false lerobot-train --policy.type=smolvla --policy.repo_id=xxxx/smolvla-test --dataset.repo_id=lerobot/svla_so100_stacking/ --batch_size=64 --steps=20 --policy.push_to_hub False --profiling True

Eval example:
TOKENIZERS_PARALLELISM=false lerobot-eval --policy.path=HuggingFaceVLA/smolvla_libero --env.type=libero --env.task=libero_object --eval.batch_size=1 --eval.n_episodes=1 --policy.device=cuda --profiling True


## SECTION TO REMOVE BEFORE SUBMITTING YOUR PR

**Note**: Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR. Try to avoid tagging more than 3 people.

**Note**: Before submitting this PR, please read the [contributor guideline](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md#submitting-a-pull-request-pr).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant