File size: 1,124 Bytes
ef79ec1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 | 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()
|