本文向所有亲们介绍在python当中配置tensorrt环境、使用tensorrt环境进行推理的教程，主要分为两大部分，第一部分环境配置，第二部分前向推理。
第一部分 环境配置
第一步：检查你的系统类型、cuda、cudnn版本，其中cuda查看命令及本人环境显示如下：
nvcc -V
显示如下
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Wed_Jun__2_19:15:15_PDT_2021
Cuda compilation tools, release 11.4, V11.4.48
Build cuda_11.4.r11.4/compiler.30033411_0
查看ubuntu系统信息命令：
sudo lsb_release -a
显示如下：
o LSB modules are available.
Distributor ID: Ubuntu
Description: Ubuntu 18.04.6 LTS
Release: 18.04
Codename: bionic
第二步： 官方下载链接下载tensorrt包。注意，需要亲亲注册哦，不然无法下载的啦。本人下载的是tar格式的
第三步：解压文件：
tar -xvf TensorRT-8.4.3.1.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz
第四部：激活你的当前python环境的情况下，进入你的解压包，安装对应python版本的tensorrt文件。在python文件夹下的内容如下，可以看到有好几个python对应版本的tensorrt，选择亲亲当前环境下的python版本进行安装就可以了，我的python版本是3.8，我就安装tensorrt-8.4.3.1-cp38-none-linux_x86_64.whl。
cd TensorRT-8.4.3.1/python/
pip install tensorrt-8.4.3.1-cp38-none-linux_x86_64.whl
第五步： 这下亲亲是不是觉得已经安装完成了可以使用了，当然不是，此时，一旦你在python中导入tensorrt时，会显示找不到包了的。所以，亲亲们要把解压出来的tensorrt文件当中的lib添加到环境变量当中哦！
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/TensorRT-8.4.3.1/lib
当然了，你还可以把lib和include中的所有文件复制到系统路径下面，命令如下：
sudo cp -r /opt/TensorRT-8.4.3.1/lib/* /usr/lib
sudo cp -r /opt/TensorRT-8.4.3.1//include/* /usr/include
第六步：检查python中tensorrt是否安装成功
>>> import tensorrt as tr
>>> tr.__version__
'8.4.3.1'
>>>
至此，python中配置tensorrt环境就完成了（码字不易，赶紧点赞，别想白嫖）
第二部分 模型转化
第一步：将pytorch中训练出来的pth模型转化为onnx,转化脚本如下：
import torch
import onnx
def pth_to_onnx(input_path,output_path):
'''
1)声明：使用本函数之前，必须保证你手上已经有了.pth模型文件.
2)功能：本函数功能四将pytorch训练得到的.pth文件转化为onnx文件。
'''
torch_model = torch.load(input_path) # pytorch模型加载,此处加载的模型包含图和参数
#torch_model = selfmodel() # 若只需要保存参数，可以换成这一种，其中selfmodel需要自己编写
torch_model.eval()
x = torch.randn(1,1,28,28) # 输入一张28*28的灰度图像并生成张量
export_onnx_file = output_path #输出.onnx文件的文件路径及文件名
torch.onnx.export(torch_model,
x,
export_onnx_file,
opset_version=9, #操作的版本，稳定操作集为9
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={"input": {0: "batch_size"}, # 批处理变量
"output": {0: "batch_size"}}
)
#onnx_model = onnx.load('model_all.onnx') #加载.onnx文件
#onnx.checker.check_model(onnx_model)
#print(onnx.helper.printable_graph(onnx_model.graph)) #打印.onnx文件信息
第二步：将onnx格式的转化为tensorrt支持的格式，转化方法比较多，这里只写一种在终端当中的方法（说实话，其他方法我也没有试哈哈）：
（1）进入到之前安装的Tensorrt安装路径，该路径下有一个trtexec文件
cd /opt/TensorRT-8.4.3.1/bin
（2）使用最简单的命令转化格式：
sudo ./trtexec --onnx=/home/photo-de/mm.onnx --saveEngine=m1.engine --fp16
其中，–onnx后面是你的onnx文件的路径， --saveEngine后面跟的是生成的engine文件的保存路径。–fp16 指生成半精度 模型，加快推理速度。如果只需要简单的转化以下就做到这一步可以了，否则请看（3）
（3）./trtexec 后面有很多参数，可以根据自己的实际情况来设置相应的参数。如下（这部分是引用）：
#静态输入尺度
./trtexec --onnx=<onnx_file> \
--explicitBatch \ #在构建引擎时使用显式批大小(默认=隐式)显示批处理
--saveEngine=<tensorRT_engine_file> \ #输出engine
--workspace=<size_in_megabytes> \ #设置工作空间大小单位是MB(默认为16MB)
--fp16 #除了fp32之外，还启用fp16精度(默认=禁用)
#动态输入尺度
./trtexec --onnx=<onnx_file> \
--minShapes=input:<shape_of_min_batch> \ #最小的batchsize x 通道数 x 输入尺寸x x 输入尺寸y
--optShapes=input:<shape_of_opt_batch> \ #最佳输入维度，跟maxShapes一样就好
--maxShapes=input:<shape_of_max_batch> \ #最大输入维度
--workspace=<size_in_megabytes> \ #设置工作空间大小单位是MB(默认为16MB)
--saveEngine=<engine_file> \ #输出engine
--fp16 #除了fp32之外，还启用fp16精度(默认=禁用)
举个例子：
#小尺寸的图片可以多batchsize即8x3x416x416
/home/zxl/TensorRT-7.2.3.4/bin/trtexec --onnx=mymodel.onnx \
--minShapes=input:1x3x416x416 \
--optShapes=input:8x3x416x416 \
--maxShapes=input:8x3x416x416 \
--workspace=4096 \
--saveEngine=mymodel.engine \
--fp16
#由于内存不够了所以改成4x3x608x608
/home/zxl/TensorRT-7.2.3.4/bin/trtexec --onnx=mymodel.onnx \
--minShapes=input:1x3x608x608 \
--optShapes=input:4x3x608x608 \
--maxShapes=input:4x3x608x608 \
--workspace=4096 \
--saveEngine=mymodel.engine \
--fp16
到此，我们得到了最终的engine文件。
第三部分 使用python+tensorrt进行模型推理
在你的推理环境当中需要再次安装pycuda，numpy。直接上代码：
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]
def predicts(model_path,images):
batch_size = 1
model = TrtModel(model_path)
result = model(images, batch_size)
return result
if __name__ == '__main__':
model_path = 'you model '
images = 'you image'
pred = predicts(model_path,images)
本文含有隐藏内容，请 开通VIP 后查看