tenrorrt加载engine并推理

2022-09-09 09:05:16 阅读：296 来源： 互联网

标签：engine stream img tenrorrt shape input output 加载

tensorRT踩坑日常之engine推理

再进行tensorRT进行推理之前，需要将训练好的模型转onnx再进行序列化生成engine，然后反序列化context推理
此文章是进行序列化生成engine和推理的，不知道如何生成engine和onnx的小伙伴可以参考另一篇博客

https://blog.csdn.net/chaocainiao/article/details/124197430?spm=1001.2014.3001.5502

废话不多说，开始进入正题，上代码

# 初始化(创建引擎，为输入输出开辟&分配显存/内存.)
def init():
    model_path = "flame_sim_8.engine"
    # 加载runtime，记录log
    runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
    # 反序列化模型
    engine = runtime.deserialize_cuda_engine(open(model_path, "rb").read())
    # print("输入",engine.get_binding_shape(0))
    # print("输出",engine.get_binding_shape(1))
    # 1. Allocate some host and device buffers for inputs and outputs:
    h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(trt.float32))
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(trt.float32))
    # Allocate device memory for inputs and outputs.
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    # Create a stream in which to copy inputs/outputs and run inference.
    stream = cuda.Stream()
    # 推理上下文
    context = engine.create_execution_context()
    return context, h_input, h_output, stream, d_input, d_output

# 推理
def inference(data_path):
    global context, h_input, h_output, stream, d_input, d_output
    load_normalized_data(data_path, h_input)
    t1 = time.time()
    # 将图片数据送到cuda显存中
    cuda.memcpy_htod_async(d_input, h_input, stream)
    # 模型预测
    context.execute_async(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle)
    # 将结果送回内存中
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    ## 异步等待结果
    stream.synchronize()
    # Return the host output.
    print("推理时间", time.time() - t1)
    return h_output

加载图片

# 加载数据并将其喂入提供的pagelocked_buffer中.
def load_normalized_data(data_path, pagelocked_buffer, target_size=(224, 224)):
    img = cv2.imread(data_path)
    img = cv2.resize(img, target_size, cv2.INTER_AREA)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = img / 127.5
    img -= 1.
    img = np.transpose(np.array([img], dtype="float32"), (0, 3, 1, 2))
    # 此时img.shape为H * W * C: 224, 224, 3
    # print("图片shape", img.shape)
    # Flatten the image into a 1D array, normalize, and copy to pagelocked memory.
    np.copyto(pagelocked_buffer, img.ravel())

主函数

if __name__ == '__main__':
    context, h_input, h_output, stream, d_input, d_output = init()
    img_path = "images/1427216352-2tkd-northern-lights-senja-norway.jpg"
    # for image in range(10):
    output = inference(data_path=img_path)
    print("type  output:", type(output), "output.shape:", output.shape, "output:", output, "\n")

该模型是一个二分类的模型

from：https://blog.csdn.net/chaocainiao/article/details/124316390

标签：engine,stream,img,tenrorrt,shape,input,output,加载
来源： https://www.cnblogs.com/chentiao/p/16671459.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9

tenrorrt加载engine并推理

tensorRT踩坑日常之engine推理