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:
2026-02-02 14:49:47 +08:00
parent 956bcbbc3e
commit 745cadc8e7
18 changed files with 68258 additions and 130 deletions

53
check_engine_output.py Normal file
View 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()