Skip to content

Commit

Permalink
[MetaSchedule] Improve the script for TorchBench model tuning & bench…
Browse files Browse the repository at this point in the history
…marking (#13255)

This PR adds features to the `python/tvm/meta_schedule/testing/torchbench/run.py`.

- Integrate with the TVM PyTorch integration to handle boolean tensor and unaligned memory.
- Deduplicate collected tuning tasks to prevent thousands of tasks created by hundreds of subgraphs with similar structure.
- Add option to cast model to float32, which are more stable numerically than float16 and prevents inaccurate result from many models.
- Add option to choose search strategy in MetaSchedule.
- Inspect output error if the actual output doesn't match the expectation. Also save the actual output and expected output for further analysis if needed.
- Save subgraphs and their example input for debug purpose.
- Print MetaSchedule profiling information at the end of execution.
- Detach PyTorch tensor before exporting to dlpack.
- Fix the sys path to avoid conflict with the `benchmarks` package installed by TorchBench dependency.
- Trim all command line args passed in, in order to prevent breaking some TorchBench model that depends on args.
- Empty cuda cache before starting the actual benchmark.
  • Loading branch information
yelite authored Nov 3, 2022
1 parent 215f0e2 commit b98b9f9
Show file tree
Hide file tree
Showing 2 changed files with 218 additions and 62 deletions.
Loading

0 comments on commit b98b9f9

Please sign in to comment.