深度学习 linux TensorRT学习笔记--Ubuntu20.04安装TensorRT 8.2.5

目录

前言

1--查看本机环境配置

2--下载并安装Tensor RT

3--实例测试

3-1--验证Onnx模型的可用性

3-2--将Onnx模型转换为推理引擎engine

3-3--基于Tensor RT使用engine模型进行推理

4--参考

前言

        推荐结合官方文档 3.2.3节中的Tar File Installation安装教程进行安装；

1--查看本机环境配置

import torch

print(torch.__version__)

print(torch.version.cuda)

print(torch.backends.cudnn.version())

博主的配置环境：

        ① Ubuntu 20.04

        ② Cuda 11.3

        ③ NVIDIA GeForce RTX 3060

        ④ Pytorch 1.12.0

        ⑤ Python 3.9

2--下载并安装Tensor RT

① 下载地址：Tensor RT官方下载地址

        博主下载的版本为Tensor RT 8.2.5.1。

② 安装依赖：

pip install 'pycuda<2021.1'

# 注意onnxruntime-gpu的版本需与实际环境进行匹配，这里博主选择1.11版本

pip install onnxruntime-gpu==1.11

 ③ 解压并配置环境变量： 

# 解压

tar -zxvf TensorRT-8.2.5.1.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz

# 配置环境变量

sudo gedit ~/.bashrc # 也可以使用vim

# 末尾添加以下两条路径，需根据解压的实际路径

export LD_LIBRARY_PATH=$PATH:/home/liujinfu/Downloads/TensorRT-8.2.5.1/lib:$LD_LIBRARY_PATH

export LIBRARY_PATH=$PATH:/home/liujinfu/Downloads/TensorRT-8.2.5.1/lib::$LIBRARY_PATH

# 保存后刷新环境变量

source ~/.bashrc

④ 安装Tensor RT库：

cd TensorRT-8.2.5.1/python

# 根据Python版本安装，博主为python3.9

pip install tensorrt-8.2.5.1-cp39-none-linux_x86_64.whl

# 安装依赖

cd TensorRT-8.2.5.1/graphsurgeon

pip install graphsurgeon-0.4.5-py2.py3-none-any.whl

⑤ 查看安装版本：

import onnxruntime as ort

import tensorrt

print(ort.get_device())

print(ort.get_available_providers())

print(tensorrt.__version__)

3--实例测试

        深度学习模型部署流程一般为：Pytorch → Onnx → TensorRT；这里博主选取一个姿态估计的Onnx模型（ThreeDPose）作为实例测试：Onnx模型下载地址

3-1--验证Onnx模型的可用性

# 导入第三方库

import onnx

import numpy as np

import onnxruntime as ort

import cv2

# 导入下载的Onnx模型

model_path = './Resnet34_3inputs_448x448_20200609.onnx'

onnx_model = onnx.load(model_path)

onnx.checker.check_model(onnx_model)

# 前处理：读入图像并调整为输入维度

img = cv2.imread("data/InitImg.png")

img = cv2.resize(img,(448,448)).transpose(2,0,1)

img = np.array(img)[np.newaxis, :, :, :].astype(np.float32)

# 设置模型session以及输入信息

sess = ort.InferenceSession(model_path,providers= ort.get_available_providers()) # 这一步可能会报错，一般与onnxruntime的版本有关，需根据实际情况进行调整

input_name1 = sess.get_inputs()[0].name

input_name2 = sess.get_inputs()[1].name

input_name3 = sess.get_inputs()[2].name

# 使用Onnx模型进行推理

output = sess.run(None, {input_name1: img, input_name2: img, input_name3: img})

print(output)

# 后处理

# 代码。。。

        正常运行，没有报错就表明下载的Onnx模型没有问题；

3-2--将Onnx模型转换为推理引擎engine

        借助Tensor RT自带的可执行文件trtexec，将Onnx模型转换为推理引擎：

cd TensorRT-8.2.5.1/bin

./trtexec --onnx=/home/liujinfu/Desktop/Tensor_Test/Resnet34_3inputs_448x448_20200609.onnx --saveEngine=/home/liujinfu/Desktop/Tensor_Test/Resnet34_3inputs_448x448_20200609.trt --workspace=6000

# --onnx=path1 表示待转换的onnx模型

# --saveEngine=path2 表示保存的推理engine模型

3-3--基于Tensor RT使用engine模型进行推理

# 导入第三方库

import torch

import cv2

import tensorrt as trt

import numpy as np

def trt_version():

return trt.__version__

def torch_device_from_trt(device):

if device == trt.TensorLocation.DEVICE:

return torch.device("cuda")

elif device == trt.TensorLocation.HOST:

return torch.device("cpu")

else:

return TypeError("%s is not supported by torch" % device)

def torch_dtype_from_trt(dtype):

if dtype == trt.int8:

return torch.int8

elif trt_version() >= '7.0' and dtype == trt.bool:

return torch.bool

elif dtype == trt.int32:

return torch.int32

elif dtype == trt.float16:

return torch.float16

elif dtype == trt.float32:

return torch.float32

else:

raise TypeError("%s is not supported by torch" % dtype)

class TRTModule(torch.nn.Module):

def __init__(self, engine=None, input_names=None, output_names=None):

super(TRTModule, self).__init__()

self.engine = engine

if self.engine is not None:

# engine创建执行context

self.context = self.engine.create_execution_context()

self.input_names = input_names

self.output_names = output_names

def forward(self, *inputs):

batch_size = inputs[0].shape[0]

bindings = [None] * (len(self.input_names) + len(self.output_names))

# 创建输出tensor，并分配内存

outputs = [None] * len(self.output_names)

for i, output_name in enumerate(self.output_names):

idx = self.engine.get_binding_index(output_name) # 通过binding_name找到对应的input_id

dtype = torch_dtype_from_trt(self.engine.get_binding_dtype(idx)) # 找到对应的数据类型

shape = (batch_size,) + tuple(self.engine.get_binding_shape(idx)) # 找到对应的形状大小

device = torch_device_from_trt(self.engine.get_location(idx))

output = torch.empty(size=shape, dtype=dtype, device=device)

outputs[i] = output

bindings[idx] = output.data_ptr() # 绑定输出数据指针

for i, input_name in enumerate(self.input_names):

idx = self.engine.get_binding_index(input_name)

bindings[idx] = inputs[0].contiguous().data_ptr() # 应当为inputs[i]，对应3个输入。但由于我们使用的是单张图片，所以将3个输入全设置为相同的图片。

self.context.execute_async(

batch_size, bindings, torch.cuda.current_stream().cuda_stream

) # 执行推理

outputs = tuple(outputs)

if len(outputs) == 1:

outputs = outputs[0]

return outputs

if __name__ == "__main__":

logger = trt.Logger(trt.Logger.INFO)

# 加载推理引擎，返回ICudaEngine对象

with open("./Resnet34_3inputs_448x448_20200609.trt", "rb") as f, trt.Runtime(logger) as runtime:

engine = runtime.deserialize_cuda_engine(f.read())

# 查看输入输出的名字，类型，大小

for idx in range(engine.num_bindings):

is_input = engine.binding_is_input(idx)

name = engine.get_binding_name(idx)

op_type = engine.get_binding_dtype(idx)

shape = engine.get_binding_shape(idx)

print('input id:', idx, ' is input: ', is_input, ' binding name:', name, ' shape:', shape, 'type: ', op_type)

trt_model = TRTModule(engine, ["input.1", "input.4", "input.7"], ["499", "504", "516", "530"])

# 加载测试图片

image = cv2.imread("./test1.jpg")

# 前处理

image = cv2.resize(image, (200,64))

image = image.transpose(2,0,1)

img_input = image.astype(np.float32)

img_input = torch.from_numpy(img_input)

img_input = img_input.unsqueeze(0)

img_input = img_input.to('cuda:0')

# 运行模型进行推理

result_trt = trt_model(img_input)

4--参考

TensorRT8.2最新版入门教程

文章链接

 您阅读本篇文章共花了：