在 python 上使用 TensorRT .engine 文件进行推理

Posted 2023-02-23

【中文标题】在 python 上使用 TensorRT .engine 文件进行推理

【英文标题】：Inference with TensorRT .engine file on python

【发布时间】：2020-04-04 10:30:07

【问题描述】：

我使用 Nvidia 的迁移学习工具包 (TLT) 进行训练，然后使用 tlt-converter 将 .etlt 模型转换为 .engine 文件。

我想在 python 中使用这个 .engine 文件进行推理。但是由于我使用 TLT 进行训练，所以我没有任何冻结的图形或 pb 文件，这是所有 TensorRT 推理教程所需要的。

我想知道 .engine 文件是否可以进行 python 推理。 如果不是，那么支持哪些转换（UFF、ONNX）可以实现这一点？

【问题讨论】：

您的意思是在 Python 中使用 TensorRT 运行 .engine 文件？ this documentation 解释你需要什么吗？

是的，但我无法确定 h_output 中的值。由于我使用了 TLT(detectnetv2)，我不知道模型规格，因此我无法理解它预测的浮点数列表。

@Sharan，嗨，你解决了这个问题吗？我也想仅使用张量 rt 推断detectnetv2，但不知道如何正确准备输入和处理输出...

@bronstein87 不，我无法解决它。我什至在这个问题上悬赏，但仍然没有得到答案。

【参考方案1】：

可以通过 .engine 文件进行 Python 推理。下面的示例从磁盘加载一个 .trt 文件（字面上与 .engine 文件相同）并执行单一推理。

在这个项目中，我在使用之前使用 onnx2trt 可执行文件将 ONNX 模型转换为 TRT 模型。您甚至可以使用 ONNX 作为中间件将 PyTorch 模型转换为 TRT。

import tensorrt as trt
import numpy as np
import os

import pycuda.driver as cuda
import pycuda.autoinit



class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()

class TrtModel:
    
    def __init__(self,engine_path,max_batch_size=1,dtype=np.float32):
        
        self.engine_path = engine_path
        self.dtype = dtype
        self.logger = trt.Logger(trt.Logger.WARNING)
        self.runtime = trt.Runtime(self.logger)
        self.engine = self.load_engine(self.runtime, self.engine_path)
        self.max_batch_size = max_batch_size
        self.inputs, self.outputs, self.bindings, self.stream = self.allocate_buffers()
        self.context = self.engine.create_execution_context()

                
                
    @staticmethod
    def load_engine(trt_runtime, engine_path):
        trt.init_libnvinfer_plugins(None, "")             
        with open(engine_path, 'rb') as f:
            engine_data = f.read()
        engine = trt_runtime.deserialize_cuda_engine(engine_data)
        return engine
    
    def allocate_buffers(self):
        
        inputs = []
        outputs = []
        bindings = []
        stream = cuda.Stream()
        
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding)) * self.max_batch_size
            host_mem = cuda.pagelocked_empty(size, self.dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            
            bindings.append(int(device_mem))

            if self.engine.binding_is_input(binding):
                inputs.append(HostDeviceMem(host_mem, device_mem))
            else:
                outputs.append(HostDeviceMem(host_mem, device_mem))
        
        return inputs, outputs, bindings, stream
       
            
    def __call__(self,x:np.ndarray,batch_size=2):
        
        x = x.astype(self.dtype)
        
        np.copyto(self.inputs[0].host,x.ravel())
        
        for inp in self.inputs:
            cuda.memcpy_htod_async(inp.device, inp.host, self.stream)
        
        self.context.execute_async(batch_size=batch_size, bindings=self.bindings, stream_handle=self.stream.handle)
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out.host, out.device, self.stream) 
            
        
        self.stream.synchronize()
        return [out.host.reshape(batch_size,-1) for out in self.outputs]


        
        
if __name__ == "__main__":
 
    batch_size = 1
    trt_engine_path = os.path.join("..","models","main.trt")
    model = TrtModel(trt_engine_path)
    shape = model.engine.get_binding_shape(0)

    
    data = np.random.randint(0,255,(batch_size,*shape[1:]))/255
    result = model(data,batch_size)
    

大家注意安全！

【参考方案2】：

您可以使用 python 来推断 .engine 文件。有两种方法可以做到这一点，

您需要在系统上安装 Tensorrt 及其兼容的 cuda。在同一环境中，您需要将 .etlt 文件转换为 .engine 文件。稍后您可以使用 python 脚本执行推理。在不同的系统上重现相同的活动是相当乏味的。 另一种简单的方法是从 tlt docker 本身运行推理脚本。这样，docker 提供了所有必要的兼容包。这种方法可以有效地扩展。您可以参考下面的文章了解详细信息。 https://azmoosa.medium.com/deploying-nvidia-tlt-and-tensorrt-applications-using-docker-containers-ecb7abd6366f