pytorch

1
"""Trivial use of Timer API:
2

3
$ python -m examples.simple_timeit
4
"""
5

6
import torch
7

8
import torch.utils.benchmark as benchmark_utils
9

10

11
def main():
12
    timer = benchmark_utils.Timer(
13
        stmt="x + y",
14
        globals={"x": torch.ones((4, 8)), "y": torch.ones((1, 8))},
15
        label="Broadcasting add (4x8)",
16
    )
17

18
    for i in range(3):
19
        print(f"Run: {i}\n{'-' * 40}")
20
        print(f"timeit:\n{timer.timeit(10000)}\n")
21
        print(f"autorange:\n{timer.blocked_autorange()}\n\n")
22

23

24
if __name__ == "__main__":
25
    main()
26