feat: TensorRT 固定 batch=4 重构
- tensorrt_engine.py 工业级 Buffer Pool - preprocessor.py 添加 pad_to_batch4() - postprocessor.py 支持批量输出 - settings.py 固定 batch_size=4
This commit is contained in:
53
check_engine_output.py
Normal file
53
check_engine_output.py
Normal file
@@ -0,0 +1,53 @@
|
||||
"""检查 TensorRT Engine 输出的实际 shape"""
|
||||
import sys
|
||||
import os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import numpy as np
|
||||
|
||||
try:
|
||||
import tensorrt as trt
|
||||
import pycuda.driver as cuda
|
||||
import pycuda.autoinit
|
||||
|
||||
engine_path = "./models/yolo11n.engine"
|
||||
|
||||
with open(engine_path, "rb") as f:
|
||||
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
|
||||
engine = runtime.deserialize_cuda_engine(f.read())
|
||||
|
||||
context = engine.create_execution_context()
|
||||
|
||||
print("=" * 60)
|
||||
print("Engine Binding Information")
|
||||
print("=" * 60)
|
||||
|
||||
for i in range(engine.num_bindings):
|
||||
name = engine.get_binding_name(i)
|
||||
shape = engine.get_binding_shape(i)
|
||||
dtype = trt.nptype(engine.get_binding_dtype(i))
|
||||
is_input = engine.binding_is_input(i)
|
||||
|
||||
print(f"\nBinding {i}: {name}")
|
||||
print(f" Shape: {shape}")
|
||||
print(f" Dtype: {dtype}")
|
||||
print(f" Is Input: {is_input}")
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
|
||||
input_shape = engine.get_binding_shape(0)
|
||||
output_shape = engine.get_binding_shape(1)
|
||||
|
||||
print(f"Input shape: {input_shape}")
|
||||
print(f"Output shape: {output_shape}")
|
||||
|
||||
input_size = np.prod([max(1, s) for s in input_shape])
|
||||
output_size = np.prod([max(1, s) for s in output_shape])
|
||||
|
||||
print(f"Input size: {input_size}")
|
||||
print(f"Output size: {output_size}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
Reference in New Issue
Block a user