(f, *args, **kwargs)
| 59 | |
| 60 | @torch.no_grad() |
| 61 | def 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 | |
| 73 | def get_sam_model(only_one_block=False, batchsize=1): |
| 74 | sam = sam_model_registry[model_type](checkpoint=checkpoint_path).cuda() |
no test coverage detected