u52f85420
04bf90c09b
- 新增 irstd 核心包、scripts、tests、configs - 新增研究文档：literature_review / mathematical_analysis / novelty_matrix / experiment_plan / research_decision / research_log / manuscript_outline - 新增 evidence、results 数据与审计产物 - 新增 pyproject.toml、requirements、.gitignore
38 lines
2.0 KiB
Python
38 lines
2.0 KiB
Python
"""Hardware measurement only; random input is NOT accuracy evaluation."""
|
|
import argparse,json,time,statistics
|
|
from pathlib import Path
|
|
import torch
|
|
from irstd.models import TinyUNet
|
|
from irstd.core import candidate_pool
|
|
|
|
def main():
|
|
p=argparse.ArgumentParser();p.add_argument('--out',required=True);p.add_argument('--size',type=int,default=256)
|
|
p.add_argument('--device',default='cuda');p.add_argument('--iterations',type=int,default=100);a=p.parse_args()
|
|
if a.iterations<1 or a.size<2:p.error('Invalid benchmark dimensions')
|
|
model=TinyUNet().to(a.device).eval();x=torch.randn(1,1,a.size,a.size,device=a.device)
|
|
cuda=str(a.device).startswith('cuda')
|
|
def sync():
|
|
if cuda:torch.cuda.synchronize()
|
|
if cuda:torch.cuda.reset_peak_memory_stats()
|
|
times=[];full=[]
|
|
with torch.no_grad():
|
|
for i in range(20+a.iterations):
|
|
sync();start=time.perf_counter();y=model(x).sigmoid();sync();middle=time.perf_counter()
|
|
candidate_pool(y[0,0].cpu().numpy());end=time.perf_counter()
|
|
if i>=20:times.append((middle-start)*1000);full.append((end-start)*1000)
|
|
# Exact conv MAC estimate via hooks, measured for declared resolution.
|
|
macs=[0]
|
|
def hook(m,inp,out):
|
|
macs[0]+=out.numel()*(m.in_channels//m.groups)*m.kernel_size[0]*m.kernel_size[1]
|
|
handles=[m.register_forward_hook(hook) for m in model.modules() if isinstance(m,torch.nn.Conv2d)]
|
|
with torch.no_grad():model(x)
|
|
for h in handles:h.remove()
|
|
Path(a.out).write_text(json.dumps(dict(device=a.device,gpu=torch.cuda.get_device_name() if cuda else None,
|
|
input=[1,1,a.size,a.size],dtype='float32',warmup=20,iterations=a.iterations,
|
|
parameters=sum(p.numel() for p in model.parameters()),conv_MACs=macs[0],conv_FLOPs_2_per_MAC=2*macs[0],
|
|
model_latency_median_ms=statistics.median(times),end_to_end_median_ms=statistics.median(full),
|
|
peak_allocated_bytes=torch.cuda.max_memory_allocated() if cuda else None,
|
|
excludes='Disk IO, candidate MLP, graph solver; FLOPs exclude interpolation/activations/pooling',
|
|
accuracy='not measured; random input'),indent=2))
|
|
if __name__=='__main__':main()
|