Jetson学习笔记(二):TensorRT 查看模型的输入输出

通过下面代码运行即可

import tensorrt as trt

def get_engine(engine_path):
    # If a serialized engine exists, use it instead of building an engine.
    print("Reading engine from file {}".format(engine_path))
    with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

TRT_LOGGER = trt.Logger()
# engine = get_engine("yolov4_1.trt")

engine = get_engine("./models/face_reg_mnet.engine")
for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * 1
        dims = engine.get_binding_shape(binding)
        print('***',size)
        print('*****',dims)
        print('***********',binding)
        print("input =", engine.binding_is_input(binding))
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # print("dtype = ", dtype)

输出结果: