From 328b083e66df01c8d586695e25d8ae54df9b09ee Mon Sep 17 00:00:00 2001 From: George Hotz Date: Sat, 11 May 2024 11:02:44 -0700 Subject: [PATCH] lil profiling script --- test/external/external_benchmark_schedule.py | 23 ++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 test/external/external_benchmark_schedule.py diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py new file mode 100644 index 0000000000..991d051d66 --- /dev/null +++ b/test/external/external_benchmark_schedule.py @@ -0,0 +1,23 @@ +from extra.models.resnet import ResNet50 +from tinygrad import Tensor +from tinygrad.helpers import Profiling, Timing, getenv +from tinygrad.engine.realize import lower_schedule + +if __name__ == "__main__": + mdl = ResNet50() + img = Tensor.empty(64, 3, 224, 224) + + PROFILE = getenv("PROFILE", 1) + + with Profiling(PROFILE): + with Timing("***** model forward in "): + out = mdl(img) + + with Profiling(PROFILE): + with Timing("***** model schedule in "): + sched = out.schedule() + + with Profiling(PROFILE): + with Timing("***** model lower in "): + ei = list(lower_schedule(sched)) +