moebiusT7/book-ocr-studio

1"""Small pinned-memory PCIe probe, launched only on idle GPUs by calibrate_gpu.py."""2import json,time,statistics3import torch4n=32*1024*10245host=torch.empty(n,dtype=torch.uint8,pin_memory=True);host.fill_(1)6device=torch.empty(n,dtype=torch.uint8,device='cuda')7for _ in range(3):device.copy_(host,non_blocking=True);torch.cuda.synchronize()8results={}9for name in ['h2d','d2h']:10 values=[]11 for _ in range(8):12  start=time.perf_counter()13  if name=='h2d':device.copy_(host,non_blocking=True)14  else:host.copy_(device,non_blocking=True)15  torch.cuda.synchronize();values.append(n/(time.perf_counter()-start)/1e9)16 results[name+'_gbps']=statistics.median(values)17print(json.dumps(results))18