119 lines
3.5 KiB
Python
119 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
测试 TensorRT 引擎加载和推理
|
||
使用随机图像测试
|
||
"""
|
||
|
||
import numpy as np
|
||
import time
|
||
from ultralytics import YOLO
|
||
|
||
|
||
def test_tensorrt_engine():
|
||
"""测试 TensorRT 引擎"""
|
||
print("TensorRT 引擎测试")
|
||
print("=" * 60)
|
||
|
||
engine_path = "C:/Users/16337/PycharmProjects/Security/yolo11n.engine"
|
||
|
||
# 1. 加载引擎
|
||
print(f"🚀 加载 TensorRT 引擎: {engine_path}")
|
||
try:
|
||
model = YOLO(engine_path, task='detect')
|
||
print("✅ 引擎加载成功")
|
||
except Exception as e:
|
||
print(f"❌ 引擎加载失败: {e}")
|
||
return
|
||
|
||
# 2. 测试单帧推理
|
||
print(f"\n🔄 测试单帧推理...")
|
||
test_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
|
||
|
||
try:
|
||
start_time = time.time()
|
||
results = model.predict(
|
||
test_image,
|
||
imgsz=640,
|
||
conf=0.45,
|
||
verbose=False,
|
||
device=0
|
||
)
|
||
elapsed = (time.time() - start_time) * 1000
|
||
print(f"✅ 单帧推理成功，耗时: {elapsed:.1f}ms")
|
||
except Exception as e:
|
||
print(f"❌ 单帧推理失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return
|
||
|
||
# 3. 测试批量推理
|
||
print(f"\n🔄 测试批量推理（batch=4）...")
|
||
test_images = [np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8) for _ in range(4)]
|
||
|
||
try:
|
||
start_time = time.time()
|
||
results = model.predict(
|
||
test_images,
|
||
imgsz=640,
|
||
conf=0.45,
|
||
verbose=False,
|
||
device=0
|
||
)
|
||
elapsed = (time.time() - start_time) * 1000
|
||
per_frame = elapsed / 4
|
||
print(f"✅ 批量推理成功，总耗时: {elapsed:.1f}ms，每帧: {per_frame:.1f}ms")
|
||
except Exception as e:
|
||
print(f"❌ 批量推理失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return
|
||
|
||
# 4. 性能测试
|
||
print(f"\n🔄 性能测试（100帧）...")
|
||
inference_times = []
|
||
|
||
try:
|
||
for i in range(100):
|
||
test_image = np.random.randint(0, 255, (640, 640, 3), dtype=np.uint8)
|
||
|
||
start_time = time.time()
|
||
results = model.predict(
|
||
test_image,
|
||
imgsz=640,
|
||
conf=0.45,
|
||
verbose=False,
|
||
device=0
|
||
)
|
||
elapsed = (time.time() - start_time) * 1000
|
||
inference_times.append(elapsed)
|
||
|
||
if (i + 1) % 20 == 0:
|
||
print(f" 进度: {i+1}/100")
|
||
|
||
print(f"\n{'='*60}")
|
||
print("性能统计")
|
||
print(f"{'='*60}")
|
||
print(f"平均推理延迟: {np.mean(inference_times):.1f}ms")
|
||
print(f"P50推理延迟: {np.percentile(inference_times, 50):.1f}ms")
|
||
print(f"P95推理延迟: {np.percentile(inference_times, 95):.1f}ms")
|
||
print(f"P99推理延迟: {np.percentile(inference_times, 99):.1f}ms")
|
||
print(f"最小延迟: {np.min(inference_times):.1f}ms")
|
||
print(f"最大延迟: {np.max(inference_times):.1f}ms")
|
||
print(f"{'='*60}")
|
||
|
||
print("\n✅ 所有测试通过！TensorRT 引擎工作正常")
|
||
|
||
except Exception as e:
|
||
print(f"❌ 性能测试失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
test_tensorrt_engine()
|
||
except Exception as e:
|
||
print(f"\n❌ 测试异常: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|