501 lines
18 KiB
Python
501 lines
18 KiB
Python
"""
|
||
GPU 压力测试模块 - 测试 RTX 3050 极限性能
|
||
目标：
|
||
1. 测试不同分辨率下最大每秒处理帧数
|
||
2. 测试最大接入摄像头路数
|
||
3. 测试不同摄像头数量下单路最大帧数
|
||
4. 测试不同抽帧率下最大摄像头数
|
||
"""
|
||
|
||
import os
|
||
import gc
|
||
import json
|
||
import time
|
||
import signal
|
||
from typing import Dict, List, Optional, Tuple
|
||
from dataclasses import dataclass, asdict
|
||
from datetime import datetime
|
||
|
||
from .utils import setup_logging, ensure_dir
|
||
|
||
logger = setup_logging()
|
||
|
||
|
||
@dataclass
|
||
class StressTestResult:
|
||
"""压力测试结果"""
|
||
test_type: str # max_fps, max_cameras, per_camera_fps, frame_skip
|
||
resolution: int
|
||
batch_size: int
|
||
num_cameras: int
|
||
target_fps: float
|
||
frame_skip: int # 抽帧间隔 (1=不抽帧, 2=每2帧取1帧, etc.)
|
||
|
||
# 结果
|
||
actual_fps: float
|
||
per_camera_fps: float
|
||
gpu_utilization: float
|
||
memory_used_mb: float
|
||
avg_latency_ms: float
|
||
p95_latency_ms: float
|
||
is_stable: bool # 是否稳定运行
|
||
error_msg: Optional[str] = None
|
||
|
||
timestamp: str = ""
|
||
|
||
def __post_init__(self):
|
||
if not self.timestamp:
|
||
self.timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
|
||
|
||
class StressTestRunner:
|
||
"""压力测试运行器"""
|
||
|
||
def __init__(self, model_path: str, output_dir: str = "./stress_results"):
|
||
self.model_path = model_path
|
||
self.output_dir = output_dir
|
||
self.results: List[StressTestResult] = []
|
||
self._interrupted = False
|
||
self._engine_cache: Dict[int, str] = {}
|
||
|
||
# 测试参数
|
||
self.test_duration = 15 # 每次测试秒数
|
||
self.warmup_sec = 2
|
||
self.cooldown_sec = 2
|
||
|
||
ensure_dir(output_dir)
|
||
self._results_file = f"{output_dir}/stress_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
|
||
|
||
signal.signal(signal.SIGINT, self._signal_handler)
|
||
|
||
def _signal_handler(self, signum, frame):
|
||
logger.warning("收到中断信号，保存当前结果...")
|
||
self._interrupted = True
|
||
self._save_results()
|
||
|
||
def _clear_gpu(self):
|
||
"""清理 GPU 显存"""
|
||
gc.collect()
|
||
try:
|
||
import torch
|
||
if torch.cuda.is_available():
|
||
torch.cuda.empty_cache()
|
||
torch.cuda.synchronize()
|
||
except:
|
||
pass
|
||
time.sleep(self.cooldown_sec)
|
||
|
||
def _save_results(self):
|
||
"""保存结果到文件"""
|
||
data = [asdict(r) for r in self.results]
|
||
with open(self._results_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
logger.info(f"结果已保存: {self._results_file}")
|
||
|
||
def _build_engine(self, resolution: int) -> str:
|
||
"""构建或获取缓存的 engine"""
|
||
if resolution in self._engine_cache:
|
||
return self._engine_cache[resolution]
|
||
|
||
from .engine_builder import TRTEngineBuilder
|
||
|
||
builder = TRTEngineBuilder(self.model_path, "./engines")
|
||
engines = builder.build_all_engines([resolution], "fp16")
|
||
|
||
self._engine_cache[resolution] = engines[resolution]
|
||
return engines[resolution]
|
||
|
||
def _run_single_test(
|
||
self,
|
||
resolution: int,
|
||
batch_size: int,
|
||
num_cameras: int,
|
||
target_fps: float,
|
||
frame_skip: int = 1
|
||
) -> Optional[StressTestResult]:
|
||
"""执行单次测试"""
|
||
from .inference_engine import TRTInferenceEngine
|
||
from .decode_thread import FrameQueueManager
|
||
from .batch_assembler import BatchAssembler
|
||
from .metrics_collector import MetricsCollector
|
||
|
||
engine = None
|
||
queue_manager = None
|
||
metrics = None
|
||
|
||
try:
|
||
# 获取 engine
|
||
engine_path = self._build_engine(resolution)
|
||
|
||
logger.info(f"测试: {resolution}x{resolution}, batch={batch_size}, "
|
||
f"cameras={num_cameras}, fps={target_fps}, skip={frame_skip}")
|
||
|
||
# 初始化组件
|
||
engine = TRTInferenceEngine(engine_path, num_streams=1, device_id=0)
|
||
queue_manager = FrameQueueManager(queue_size=2)
|
||
metrics = MetricsCollector(device_id=0, sample_interval_ms=100)
|
||
|
||
# 添加摄像头源 (使用合成数据)
|
||
effective_fps = target_fps / frame_skip # 实际采集帧率
|
||
for i in range(num_cameras):
|
||
queue_manager.add_source(
|
||
source_id=f"cam_{i:02d}",
|
||
source="synthetic",
|
||
target_fps=effective_fps,
|
||
source_type="synthetic",
|
||
resolution=(640, 480)
|
||
)
|
||
|
||
# Batch 组装器
|
||
batch_assembler = BatchAssembler(
|
||
frame_queues=queue_manager.queues,
|
||
batch_size=batch_size,
|
||
imgsz=(resolution, resolution),
|
||
use_gpu_preprocess=True,
|
||
device_id=0
|
||
)
|
||
|
||
# 启动
|
||
queue_manager.start_all()
|
||
|
||
# 预热
|
||
warmup_end = time.time() + self.warmup_sec
|
||
while time.time() < warmup_end and not self._interrupted:
|
||
batch_data = batch_assembler.assemble_batch(timeout=0.1)
|
||
if batch_data:
|
||
batch, _ = batch_data
|
||
engine.infer_sync(batch)
|
||
|
||
# 正式测试
|
||
metrics.reset()
|
||
metrics.start()
|
||
|
||
test_end = time.time() + self.test_duration
|
||
while time.time() < test_end and not self._interrupted:
|
||
batch_data = batch_assembler.assemble_batch(timeout=0.1)
|
||
if batch_data:
|
||
batch, frame_infos = batch_data
|
||
_, latency_ms = engine.infer_sync(batch)
|
||
metrics.record_inference(latency_ms, len(frame_infos))
|
||
|
||
metrics.stop()
|
||
|
||
# 收集结果
|
||
gpu_metrics = metrics.get_gpu_metrics()
|
||
throughput = metrics.get_throughput_metrics()
|
||
|
||
actual_fps = throughput.get("throughput_fps", 0)
|
||
|
||
result = StressTestResult(
|
||
test_type="stress",
|
||
resolution=resolution,
|
||
batch_size=batch_size,
|
||
num_cameras=num_cameras,
|
||
target_fps=target_fps,
|
||
frame_skip=frame_skip,
|
||
actual_fps=actual_fps,
|
||
per_camera_fps=actual_fps / num_cameras if num_cameras > 0 else 0,
|
||
gpu_utilization=gpu_metrics.get("gpu_utilization", {}).get("avg", 0),
|
||
memory_used_mb=gpu_metrics.get("memory_used_mb", {}).get("avg", 0),
|
||
avg_latency_ms=throughput.get("latency", {}).get("avg", 0),
|
||
p95_latency_ms=throughput.get("latency", {}).get("p95", 0),
|
||
is_stable=True
|
||
)
|
||
|
||
logger.info(f" 结果: {actual_fps:.1f} FPS, GPU {result.gpu_utilization:.1f}%, "
|
||
f"延迟 {result.avg_latency_ms:.1f}ms")
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
error_msg = str(e)
|
||
logger.warning(f" 测试失败: {error_msg}")
|
||
|
||
return StressTestResult(
|
||
test_type="stress",
|
||
resolution=resolution,
|
||
batch_size=batch_size,
|
||
num_cameras=num_cameras,
|
||
target_fps=target_fps,
|
||
frame_skip=frame_skip,
|
||
actual_fps=0,
|
||
per_camera_fps=0,
|
||
gpu_utilization=0,
|
||
memory_used_mb=0,
|
||
avg_latency_ms=0,
|
||
p95_latency_ms=0,
|
||
is_stable=False,
|
||
error_msg=error_msg[:200]
|
||
)
|
||
finally:
|
||
# 清理
|
||
if queue_manager:
|
||
queue_manager.stop_all()
|
||
if engine:
|
||
engine.cleanup()
|
||
self._clear_gpu()
|
||
|
||
def test_max_fps(self, resolutions: List[int] = [320, 480]) -> Dict[int, float]:
|
||
"""测试1: 不同分辨率下最大每秒处理帧数"""
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("测试1: 最大处理帧数 (单摄像头, 最大帧率)")
|
||
logger.info("=" * 60)
|
||
|
||
max_fps_results = {}
|
||
|
||
for res in resolutions:
|
||
if self._interrupted:
|
||
break
|
||
|
||
# 使用大 batch 和高帧率测试极限
|
||
for batch_size in [1, 4, 8]:
|
||
result = self._run_single_test(
|
||
resolution=res,
|
||
batch_size=batch_size,
|
||
num_cameras=1,
|
||
target_fps=100, # 高帧率压力测试
|
||
frame_skip=1
|
||
)
|
||
|
||
if result and result.is_stable:
|
||
self.results.append(result)
|
||
if res not in max_fps_results or result.actual_fps > max_fps_results[res]:
|
||
max_fps_results[res] = result.actual_fps
|
||
|
||
self._save_results()
|
||
|
||
logger.info(f"\n最大 FPS 结果: {max_fps_results}")
|
||
return max_fps_results
|
||
|
||
def test_max_cameras(self, resolutions: List[int] = [320, 480]) -> Dict[int, int]:
|
||
"""测试2: 不同分辨率下最大接入摄像头数"""
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("测试2: 最大摄像头接入数 (10 FPS 实时)")
|
||
logger.info("=" * 60)
|
||
|
||
max_cameras = {}
|
||
camera_counts = [1, 3, 5, 10, 15, 20, 25, 30]
|
||
|
||
for res in resolutions:
|
||
if self._interrupted:
|
||
break
|
||
|
||
max_cameras[res] = 0
|
||
|
||
for num_cams in camera_counts:
|
||
if self._interrupted:
|
||
break
|
||
|
||
result = self._run_single_test(
|
||
resolution=res,
|
||
batch_size=min(num_cams, 8),
|
||
num_cameras=num_cams,
|
||
target_fps=10,
|
||
frame_skip=1
|
||
)
|
||
|
||
if result:
|
||
self.results.append(result)
|
||
self._save_results()
|
||
|
||
# 检查是否能稳定处理 (实际 FPS >= 目标 FPS * 90%)
|
||
target_total = num_cams * 10
|
||
if result.is_stable and result.actual_fps >= target_total * 0.9:
|
||
max_cameras[res] = num_cams
|
||
else:
|
||
# 无法处理更多摄像头
|
||
logger.info(f" {res}x{res}: 最大 {max_cameras[res]} 路摄像头")
|
||
break
|
||
|
||
logger.info(f"\n最大摄像头数: {max_cameras}")
|
||
return max_cameras
|
||
|
||
def test_per_camera_fps(
|
||
self,
|
||
resolutions: List[int] = [320, 480],
|
||
camera_counts: List[int] = [1, 3, 5, 10, 15, 30]
|
||
) -> Dict[Tuple[int, int], float]:
|
||
"""测试3: 不同摄像头数量下单路最大帧数"""
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("测试3: 不同摄像头数量下单路最大帧数")
|
||
logger.info("=" * 60)
|
||
|
||
per_camera_results = {}
|
||
|
||
for res in resolutions:
|
||
if self._interrupted:
|
||
break
|
||
|
||
for num_cams in camera_counts:
|
||
if self._interrupted:
|
||
break
|
||
|
||
# 测试高帧率下的实际处理能力
|
||
result = self._run_single_test(
|
||
resolution=res,
|
||
batch_size=min(num_cams, 8),
|
||
num_cameras=num_cams,
|
||
target_fps=30, # 目标 30 FPS
|
||
frame_skip=1
|
||
)
|
||
|
||
if result:
|
||
self.results.append(result)
|
||
per_camera_results[(res, num_cams)] = result.per_camera_fps
|
||
self._save_results()
|
||
|
||
logger.info(f"\n单路最大帧数结果:")
|
||
for (res, cams), fps in per_camera_results.items():
|
||
logger.info(f" {res}x{res}, {cams}路: {fps:.1f} FPS/路")
|
||
|
||
return per_camera_results
|
||
|
||
def test_frame_skip_capacity(
|
||
self,
|
||
resolutions: List[int] = [320, 480],
|
||
frame_skips: List[int] = [1, 2, 3, 5, 10]
|
||
) -> Dict[Tuple[int, int], int]:
|
||
"""测试4: 不同抽帧率下最大摄像头数"""
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("测试4: 不同抽帧率下最大摄像头数")
|
||
logger.info("=" * 60)
|
||
|
||
frame_skip_results = {}
|
||
camera_counts = [5, 10, 15, 20, 25, 30, 40, 50]
|
||
|
||
for res in resolutions:
|
||
if self._interrupted:
|
||
break
|
||
|
||
for skip in frame_skips:
|
||
if self._interrupted:
|
||
break
|
||
|
||
max_cams = 0
|
||
|
||
for num_cams in camera_counts:
|
||
if self._interrupted:
|
||
break
|
||
|
||
# 原始帧率 30 FPS，抽帧后实际处理帧率 = 30/skip
|
||
effective_fps = 30.0 / skip
|
||
|
||
result = self._run_single_test(
|
||
resolution=res,
|
||
batch_size=min(num_cams, 8),
|
||
num_cameras=num_cams,
|
||
target_fps=30, # 源帧率
|
||
frame_skip=skip
|
||
)
|
||
|
||
if result:
|
||
self.results.append(result)
|
||
self._save_results()
|
||
|
||
# 检查是否稳定
|
||
target_total = num_cams * effective_fps
|
||
if result.is_stable and result.actual_fps >= target_total * 0.85:
|
||
max_cams = num_cams
|
||
else:
|
||
break
|
||
|
||
frame_skip_results[(res, skip)] = max_cams
|
||
logger.info(f" {res}x{res}, 抽帧{skip}: 最大 {max_cams} 路")
|
||
|
||
return frame_skip_results
|
||
|
||
def run_all_tests(self):
|
||
"""运行所有压力测试"""
|
||
logger.info("=" * 60)
|
||
logger.info("RTX 3050 GPU 压力测试")
|
||
logger.info("=" * 60)
|
||
|
||
resolutions = [320, 480]
|
||
|
||
# 测试1: 最大 FPS
|
||
max_fps = self.test_max_fps(resolutions)
|
||
|
||
# 测试2: 最大摄像头数
|
||
max_cameras = self.test_max_cameras(resolutions)
|
||
|
||
# 测试3: 单路最大帧数
|
||
per_camera = self.test_per_camera_fps(resolutions, [1, 3, 5, 10, 15, 30])
|
||
|
||
# 测试4: 抽帧容量
|
||
frame_skip = self.test_frame_skip_capacity(resolutions, [1, 2, 3, 5, 10])
|
||
|
||
# 生成报告
|
||
self._generate_report(max_fps, max_cameras, per_camera, frame_skip)
|
||
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("压力测试完成!")
|
||
logger.info(f"结果保存在: {self.output_dir}")
|
||
logger.info("=" * 60)
|
||
|
||
def _generate_report(self, max_fps, max_cameras, per_camera, frame_skip):
|
||
"""生成测试报告"""
|
||
report_path = f"{self.output_dir}/stress_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
|
||
|
||
lines = [
|
||
"# RTX 3050 GPU 压力测试报告",
|
||
f"\n生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
|
||
"\n## 1. 最大处理帧数 (单摄像头)",
|
||
"\n| 分辨率 | 最大 FPS |",
|
||
"|--------|----------|",
|
||
]
|
||
|
||
for res, fps in max_fps.items():
|
||
lines.append(f"| {res}×{res} | {fps:.1f} |")
|
||
|
||
lines.extend([
|
||
"\n## 2. 最大摄像头接入数 (10 FPS 实时)",
|
||
"\n| 分辨率 | 最大摄像头数 |",
|
||
"|--------|--------------|",
|
||
])
|
||
|
||
for res, cams in max_cameras.items():
|
||
lines.append(f"| {res}×{res} | {cams} |")
|
||
|
||
lines.extend([
|
||
"\n## 3. 不同摄像头数量下单路最大帧数",
|
||
"\n| 分辨率 | 摄像头数 | 单路 FPS |",
|
||
"|--------|----------|----------|",
|
||
])
|
||
|
||
for (res, cams), fps in per_camera.items():
|
||
lines.append(f"| {res}×{res} | {cams} | {fps:.1f} |")
|
||
|
||
lines.extend([
|
||
"\n## 4. 不同抽帧率下最大摄像头数",
|
||
"\n| 分辨率 | 抽帧间隔 | 实际帧率 | 最大摄像头数 |",
|
||
"|--------|----------|----------|--------------|",
|
||
])
|
||
|
||
for (res, skip), cams in frame_skip.items():
|
||
effective_fps = 30.0 / skip
|
||
lines.append(f"| {res}×{res} | 每{skip}帧取1帧 | {effective_fps:.1f} FPS | {cams} |")
|
||
|
||
lines.extend([
|
||
"\n## 5. 部署建议",
|
||
"\n根据测试结果，推荐配置：",
|
||
"\n| 场景 | 分辨率 | 摄像头数 | 抽帧 | 说明 |",
|
||
"|------|--------|----------|------|------|",
|
||
])
|
||
|
||
# 根据结果生成建议
|
||
if 320 in max_cameras:
|
||
lines.append(f"| 高并发 | 320×320 | {max_cameras.get(320, 10)} | 1 | 最大并发 |")
|
||
if 480 in max_cameras:
|
||
lines.append(f"| 高精度 | 480×480 | {max_cameras.get(480, 5)} | 1 | 精度优先 |")
|
||
|
||
with open(report_path, 'w', encoding='utf-8') as f:
|
||
f.write('\n'.join(lines))
|
||
|
||
logger.info(f"报告已生成: {report_path}")
|
||
|
||
|
||
def run_stress_test(model_path: str, output_dir: str = "./stress_results"):
|
||
"""运行压力测试的入口函数"""
|
||
runner = StressTestRunner(model_path, output_dir)
|
||
runner.run_all_tests()
|