pytorch

1
import torch
2
from pyarkbench import Benchmark, default_args, Timer
3

4
use_new = True
5

6

7
class Basic(Benchmark):
8
    def benchmark(self):
9
        x = [torch.ones(200, 200) for i in range(30)]
10
        with Timer() as big1:
11
            torch.save(x, "big_tensor.zip", _use_new_zipfile_serialization=use_new)
12

13
        with Timer() as big2:
14
            v = torch.load("big_tensor.zip")
15

16
        x = [torch.ones(10, 10) for i in range(200)]
17
        with Timer() as small1:
18
            torch.save(x, "small_tensor.zip", _use_new_zipfile_serialization=use_new)
19

20
        with Timer() as small2:
21
            v = torch.load("small_tensor.zip")
22

23
        return {
24
            "Big Tensors Save": big1.ms_duration,
25
            "Big Tensors Load": big2.ms_duration,
26
            "Small Tensors Save": small1.ms_duration,
27
            "Small Tensors Load": small2.ms_duration,
28
        }
29

30

31
if __name__ == "__main__":
32
    bench = Basic(*default_args.bench())
33
    print("Use zipfile serialization:", use_new)
34
    results = bench.run()
35
    bench.print_stats(results, stats=["mean", "median"])
36