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()
|