| from __future__ import annotations |
| import argparse,json,time |
| from pathlib import Path |
| import torch |
| from runtime.trt_torch import TorchTensorRTEngine |
|
|
| def bench(engine,feeds,warm=3,iters=20): |
| for _ in range(warm): engine.run(feeds) |
| torch.cuda.synchronize(); times=[] |
| for _ in range(iters): |
| t=time.perf_counter(); engine.run(feeds); torch.cuda.synchronize(); times.append((time.perf_counter()-t)*1000) |
| times.sort(); return {'mean_ms':sum(times)/len(times),'p50_ms':times[len(times)//2],'p95_ms':times[int(len(times)*.95)-1]} |
| def main(): |
| ap=argparse.ArgumentParser(); ap.add_argument('--root',default='.'); a=ap.parse_args(); r=Path(a.root); out={} |
| d=TorchTensorRTEngine(r/'engines/l4-sm89/dinov2.plan'); out['dinov2']=bench(d,{'pixel_values':torch.randn(1,3,518,518,device='cuda')}) |
| s=TorchTensorRTEngine(r/'engines/l4-sm89/dsine.plan'); out['dsine']=bench(s,{'image':torch.randn(1,3,544,544,device='cuda'),'intrins':torch.eye(3,device='cuda').unsqueeze(0)}) |
| out['peak_vram_gib']=torch.cuda.max_memory_allocated()/1024**3; print(json.dumps(out,indent=2)) |
| if __name__=='__main__':main() |
|
|
