92 lines
3.6 KiB
Python
92 lines
3.6 KiB
Python
"""
|
||
注意力模块 A/B 基准测试
|
||
|
||
目的：在相同骨干与输入下，对比注意力开/关（none/se/cbam）在单尺度与 FPN 前向的耗时差异；可选指定插入位置。
|
||
|
||
示例：
|
||
PYTHONPATH=. uv run python tests/benchmark_attention.py --device cpu --image-size 512 --runs 10 --backbone resnet34 --places backbone_high desc_head
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import time
|
||
from typing import Dict, List
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
from models.rord import RoRD
|
||
|
||
|
||
def bench_once(model: torch.nn.Module, x: torch.Tensor, fpn: bool = False) -> float:
|
||
if torch.cuda.is_available() and x.is_cuda:
|
||
torch.cuda.synchronize()
|
||
t0 = time.time()
|
||
with torch.inference_mode():
|
||
_ = model(x, return_pyramid=fpn)
|
||
if torch.cuda.is_available() and x.is_cuda:
|
||
torch.cuda.synchronize()
|
||
return (time.time() - t0) * 1000.0
|
||
|
||
|
||
def build_model(backbone: str, attention_type: str, places: List[str], device: torch.device) -> RoRD:
|
||
cfg = type("cfg", (), {
|
||
"model": type("m", (), {
|
||
"backbone": type("b", (), {"name": backbone, "pretrained": False})(),
|
||
"attention": type("a", (), {"enabled": attention_type != "none", "type": attention_type, "places": places})(),
|
||
})()
|
||
})()
|
||
model = RoRD(cfg=cfg).to(device)
|
||
model.eval()
|
||
return model
|
||
|
||
|
||
def run_suite(backbone: str, places: List[str], device: torch.device, image_size: int, runs: int) -> List[Dict[str, float]]:
|
||
x = torch.randn(1, 3, image_size, image_size, device=device)
|
||
results: List[Dict[str, float]] = []
|
||
for attn in ["none", "se", "cbam"]:
|
||
model = build_model(backbone, attn, places, device)
|
||
# warmup
|
||
for _ in range(3):
|
||
_ = model(x, return_pyramid=False)
|
||
_ = model(x, return_pyramid=True)
|
||
# single
|
||
t_list_single = [bench_once(model, x, fpn=False) for _ in range(runs)]
|
||
# fpn
|
||
t_list_fpn = [bench_once(model, x, fpn=True) for _ in range(runs)]
|
||
results.append({
|
||
"backbone": backbone,
|
||
"attention": attn,
|
||
"places": ",".join(places) if places else "-",
|
||
"single_ms_mean": float(np.mean(t_list_single)),
|
||
"single_ms_std": float(np.std(t_list_single)),
|
||
"fpn_ms_mean": float(np.mean(t_list_fpn)),
|
||
"fpn_ms_std": float(np.std(t_list_fpn)),
|
||
"runs": int(runs),
|
||
})
|
||
return results
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="RoRD 注意力模块 A/B 基准")
|
||
parser.add_argument("--backbone", type=str, default="resnet34", choices=["vgg16","resnet34","efficientnet_b0"], help="骨干")
|
||
parser.add_argument("--places", nargs="*", default=["backbone_high"], help="插入位置：backbone_high det_head desc_head")
|
||
parser.add_argument("--image-size", type=int, default=512, help="输入尺寸")
|
||
parser.add_argument("--runs", type=int, default=10, help="重复次数")
|
||
parser.add_argument("--device", type=str, default="cpu", help="cuda 或 cpu")
|
||
args = parser.parse_args()
|
||
|
||
device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
|
||
results = run_suite(args.backbone, args.places, device, args.image_size, args.runs)
|
||
|
||
# 简要打印
|
||
print("\n===== Attention A/B Summary =====")
|
||
for r in results:
|
||
print(f"{r['backbone']:<14} attn={r['attention']:<5} places={r['places']:<24} "
|
||
f"single {r['single_ms_mean']:.2f}±{r['single_ms_std']:.2f} | "
|
||
f"fpn {r['fpn_ms_mean']:.2f}±{r['fpn_ms_std']:.2f} ms")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|