Shaun ZhouandClaude Opus 5
d6adea5b05
88 个测试覆盖配置继承、注册表、GPU 守卫、YOLO 格式解析、数据集、模型形状、 标签模式、聚类原语与 SimCLR 损失，以及 train.py / eval.py / 聚类脚本的入口接线。 含两条回归测试： - PCA 白化放大噪声方向导致簇塌陷（默认 whiten=False 的依据） - classes 写成 data.yaml 路径时 num_classes 必须是类别数而非路径字符串长度 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
71 lines
2.1 KiB
Python
71 lines
2.1 KiB
Python
"""推理 profiler 脚本：CPU 上也要能跑出耗时结果（有 GPU 时额外给出显存峰值）。"""
|
|
|
|
import argparse
|
|
import importlib.util
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
from models import build_model # noqa: E402
|
|
|
|
SCRIPT = Path(__file__).resolve().parents[1] / "scripts" / "profile_inference.py"
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def script():
|
|
spec = importlib.util.spec_from_file_location("profile_inference", SCRIPT)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.fixture
|
|
def model():
|
|
return build_model(
|
|
{
|
|
"name": "defect_net",
|
|
"num_classes": 2,
|
|
"backbone": {"name": "tiny_conv", "widths": [8, 16], "depths": [1, 1]},
|
|
"neck": {"name": "upsample", "channels": 8, "num_upsamples": 1},
|
|
"head": {"hidden_channels": 8},
|
|
}
|
|
).eval()
|
|
|
|
|
|
@pytest.fixture
|
|
def args():
|
|
return argparse.Namespace(iters=2, warmup=1, amp=False, postprocess=True, rows=3, trace=None)
|
|
|
|
|
|
def test_measure_reports_latency_and_memory(script, model, args):
|
|
images = torch.randn(2, 3, 32, 32)
|
|
result = script.measure(model, images, args)
|
|
|
|
assert result["batch_size"] == 2
|
|
assert result["latency_ms"] > 0
|
|
assert result["throughput_fps"] > 0
|
|
if torch.cuda.is_available():
|
|
assert result["peak_allocated_mb"] > 0
|
|
assert result["peak_reserved_mb"] >= result["peak_allocated_mb"]
|
|
else:
|
|
assert "peak_allocated_mb" not in result
|
|
|
|
|
|
def test_profile_exports_trace(script, model, args, tmp_path, capsys):
|
|
args.trace = str(tmp_path / "trace.json")
|
|
script.profile_once(model, torch.randn(1, 3, 32, 32), args)
|
|
|
|
assert Path(args.trace).exists()
|
|
out = capsys.readouterr().out
|
|
assert "显存占用 Top" in out and "耗时 Top" in out
|
|
|
|
|
|
def test_postprocess_toggle_runs_both_paths(script, model, args):
|
|
images = torch.randn(1, 3, 32, 32)
|
|
args.postprocess = False
|
|
script.run_once(model, images, args)
|
|
args.postprocess = True
|
|
script.run_once(model, images, args)
|