35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
import tensorrt as trt
|
||
import pycuda.driver as cuda
|
||
import numpy as np
|
||
|
||
'''
|
||
ultralytics导出的engine文件 根本用不了啊。反序列化的engine一直为 None，莫名奇妙
|
||
'''
|
||
print(trt.__version__)
|
||
# 1.构建日志记录器
|
||
logger = trt.Logger(trt.Logger.INFO)
|
||
|
||
# 2.反序列化
|
||
runtime = trt.Runtime(logger)
|
||
|
||
engine_path = '/home/admin-root/haotian/python哈汽安全帽识别/5000_HeadHelmet/train2/weights/best.engine'
|
||
# engine_path = '/home/admin-root/220_project/tensorrtx-master/yolov8/build/220version0826.engine'
|
||
with open(engine_path, 'rb') as f:
|
||
serialized_engine = f.read()
|
||
|
||
engine = runtime.deserialize_cuda_engine(serialized_engine)
|
||
|
||
# 打印模型信息
|
||
for idx in range(engine.num_bindings):
|
||
name = engine.get_tensor_name(idx)
|
||
is_input = engine.get_tensor_mode(name)
|
||
op_type = engine.get_tensor_dtype(name)
|
||
shape = engine.get_tensor_shape(name)
|
||
print('input id: ',idx, '\\tis input: ', is_input, '\\tbinding name: ', name, '\\tshape: ', shape, '\\ttype: ', op_type)
|
||
# 3.执行推理
|
||
# 3.1上下文 、
|
||
context = engine.create_execution_context()
|
||
# 为输入输出指定缓冲区
|
||
|
||
|