MCPcopy Create free account
hub / github.com/pytorch/tutorials / benchmark

Function benchmark

unstable_source/gpu_quantization_torchao_tutorial.py:61–71  ·  view source on GitHub ↗
(f, *args, **kwargs)

Source from the content-addressed store, hash-verified

59
60@torch.no_grad()
61def benchmark(f, *args, **kwargs):
62 for _ in range(3):
63 f(*args, **kwargs)
64 torch.cuda.synchronize()
65
66 torch.cuda.reset_peak_memory_stats()
67 t0 = Timer(
68 stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
69 )
70 res = t0.adaptive_autorange(.03, min_run_time=.2, max_run_time=20)
71 return {'time':res.median * 1e3, 'memory': torch.cuda.max_memory_allocated()/1e9}
72
73def get_sam_model(only_one_block=False, batchsize=1):
74 sam = sam_model_registry[model_type](checkpoint=checkpoint_path).cuda()

Calls 1

fFunction · 0.70

Tested by

no test coverage detected