在Python中使用TensorRT .engine文件进行推理

Question

我使用了Nvidia的Transfer Learning Toolkit（TLT）进行训练，然后使用tlt-converter将.etlt模型转换成.engine文件。

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

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

- Sharan

---

你的意思是在Python中使用TensorRT运行.engine文件吗？这份文档是否解释了你需要的内容？ - jdehesa

是的，但我在弄清楚h_output中的值时遇到了麻烦。由于我使用了TLT(detectnetv2)，我不知道模型规格，因此无法理解它预测的浮点数列表。 - Sharan

@Sharan，你好，你解决了这个问题吗？我也想使用TensorRT来进行推理检测，但不知道如何正确准备输入和处理输出... - bronstein87

@bronstein87 不，我无法解决它。我甚至在这个问题上设置了赏金，但仍然没有得到答案。 - Sharan

Answer 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)
    

大家要注意安全!

- O Vuruşkaner

Answer 2

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

1. 您需要在系统上安装Tensorrt及其兼容的CUDA。在相同的环境中，您需要将.etlt文件转换为.engine文件。稍后，您可以使用Python脚本执行推断。在不同的系统上复制相同的活动非常繁琐。
2. 另一种简单的方法是直接从tlt docker运行推断脚本。这种方式，docker提供了所有必要的兼容包。而且这种方法很有效地扩展。您可以参考下面的文章获取详细信息。 https://azmoosa.medium.com/deploying-nvidia-tlt-and-tensorrt-applications-using-docker-containers-ecb7abd6366f

- Shaik Ahmad