python3 TensorRT 查看模型的输入输出

最新推荐文章于 2024-08-12 08:43:15 发布

文若书生

最新推荐文章于 2024-08-12 08:43:15 发布

阅读量2.5k

点赞数 2

分类专栏：人工智能

本文链接：https://blog.csdn.net/ZHOUYONGXYZ/article/details/111901546

版权

人工智能专栏收录该内容

6 篇文章 0 订阅

订阅专栏

1. 查看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)