40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
|
"""
|
||
|
|
验证 TensorRT Engine 是否可以加载
|
||
|
|
"""
|
||
|
|
import os
|
||
|
|
|
||
|
|
engine_path = r"C:\Users\16337\PycharmProjects\Security_project\yolov8n.engine"
|
||
|
|
|
||
|
|
print(f"检查文件: {engine_path}")
|
||
|
|
print(f"文件存在: {os.path.exists(engine_path)}")
|
||
|
|
|
||
|
|
if os.path.exists(engine_path):
|
||
|
|
size = os.path.getsize(engine_path)
|
||
|
|
print(f"文件大小: {size / (1024*1024):.2f} MB")
|
||
|
|
|
||
|
|
# 检查文件头
|
||
|
|
with open(engine_path, 'rb') as f:
|
||
|
|
header = f.read(32)
|
||
|
|
print(f"文件头 (hex): {header[:16].hex()}")
|
||
|
|
# TensorRT engine 文件应该以特定的 magic number 开头
|
||
|
|
# 不应该是 JSON 格式 (04 07 00 00 7b 22...)
|
||
|
|
|
||
|
|
print("\n尝试加载 Engine...")
|
||
|
|
try:
|
||
|
|
import tensorrt as trt
|
||
|
|
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
|
||
|
|
|
||
|
|
with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
|
||
|
|
engine = runtime.deserialize_cuda_engine(f.read())
|
||
|
|
|
||
|
|
if engine:
|
||
|
|
print("✓ Engine 加载成功!")
|
||
|
|
print(f" 输入数量: {engine.num_io_tensors}")
|
||
|
|
else:
|
||
|
|
print("✗ Engine 加载失败 (返回 None)")
|
||
|
|
except Exception as e:
|
||
|
|
print(f"✗ Engine 加载失败: {e}")
|
||
|
|
else:
|
||
|
|
print("文件不存在，请先运行:")
|
||
|
|
print(" yolo export model=yolov8n.pt format=engine device=0 imgsz=320 batch=1 half")
|