关键词：
由于博主最近实验需要将tensorrt部署到端到端的抓取网络当中，但之前没有使用过tensorrt，查阅了很多资料，踩了很多坑，最后才部署成功。于是想着记录一下，本篇内容主要以Unet和grcnn（antipodal robotic grasping）为例来讲解端到端pytorch模型转换为tensorrt。
文章索引
第一章：前期准备
博主的软件环境：ubuntu18+cuda11.3+cudnn8.6.0+python3.8+torch1.12.0+tensorrt8.5.2.2，GPU为RTX3070。由于cuda+cudnn安装网络上有很多教程，这里就不再赘述。
Step1：安装tensorRT8.5.2.2
可以通过官网下载。
因为去nvidia官网还需要登录，为了方便这里提供百度云网盘下载：提取码ltjy。
解压文件：
tar -xzvf TensorRT-8.5.2.2.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz
添加环境变量：
export PATH="$PATH:*****/TensorRT-8.5.2.2/bin"#自己下载的TensorRT/bin所在地址
export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:*******/TensorRT/lib"#自己下载的TensorRT/lib所在地址
创建并且激活虚拟环境：
conda create -n pt12 python=3.8
conda activate pt12
安装wheel文件：
cd TensorRT-8.5.2.2/python
pip install tensorrt-8.5.2.2-cp38-none-linux_x86_64.whl
Step2：安装onnx-tensorrt工具包
由onnx文件转到trt文件有很多种方法，如果不需要使用INT8进行量化推理，推荐使用该工具包来进行转换。
git clone https://github.com/onnx/onnx-tensorrt.git
cd onnx-tensorrt
git checkout 8.0-GA
git submodule update --init
mkdir build && cd build
cmake .. -DTENSORRT_ROOT=/******/TensorRT-8.5.2.2 #刚才装的位置
与该博客类似，报错1：cmake版本太低
解决：升级cmake版本
pip install cmake --upgrade
或去cmake官网下载
报错2：Could NOT find Protobuf（missing：Protobuf_LIBRARIESProtobuf_INCLUDE_DIR）
解决：安装libprotobuf-dev protobuf-compiler
sudo apt-get install libprotobuf-dev protobuf-compiler
proto --version
开始编译
make -j8
报错：/usr/include/NvInferRuntimeCommon.h:56:10:fatal error:cuda_runtime_api.h:没有那个文件或目录
解决：配置cuda相关环境变量
sudo gedit ~/.bashrc
export PATH=/usr/local/cuda-11.3/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH
export CPATH=/usr/local/cuda-11.3/targets/x86_64-linux/include:$CPATH
export LD_LIBARARY_PATH=/usr/local/cuda-11.3/targets/x86_64-linux/lib:$LD_LIBARARY_PATH
source ~/.basrc
开始安装：
sudo make install
运行测试：
onnx2trt -V
查看版本号，至此所有前期准备完成。
第二章：Pytorch-Unet to TensorRT-Unet
要将pytorch模型转换为tensorrt的engine模型需要经过一下两个步骤。
1、将pt或者模型文件转换为onnx文件；
2、使用转换工具将onnx文件转换为trt文件。
Pytorch-Unet通过输入一张rgb图片通过网络得到其分割后的灰度图，其github项目地址在此，其相关讲解可以参考该博客。这里直接上干货。
Step1：从github拉取代码
git clone https://github.com/milesial/Pytorch-UNet.git
cd Pytorch-UNet
git checkout v1.0
下载数据集：
通过github中代码的项目readme文件，下载其数据集，将train_hq.zip里面的图片放入data/img里，将train_mask.zip里面的图片放入data/mask里。
Step2：训练网络
为方便后续部署，修改utils中的dataset.py中preprocess函数内容，将NewW，NewH修改为960，640。
直接运行train.py文件即可。
Step3：将pt模型文件转换为onnx
test.py：
这里的dummy_input改为我们修改的（1，3，640，960）初始化模型参数保证与train.py中一致。
train.py
Step4：将onnx转换为trt模型
onnx2trt unet_deconv.onnx -o unet_deconv.trt
报错1：onnx模型太复杂，无法转换。
解决：安装onnxsim工具
pip install onnx-simplifier
python -m onnxsim input_onnx_model output_onnx_model
报错2：不能打开目标文件libnvinfer.so.8
解决：链接动态库
sudo gedit /etc/ld.so.conf
添加一行：
/home/lab/xcy/TensorRT-8.5.2.2/lib #自己TensorRT中lib所在路径
sudo ldconfig
报错3：找不到libnvinfer.so.8.4.3
解决：
sudo cp /home/lab/xcy/TensorRT-8.5.2.2/lib/libvinfer_build_resource.so.8.4.3 /usr/lib
Step5：编写推理代码
inference.py
import os
import sys
import time
# from PIL import Image
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
# TensorRT logger singleton
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def allocate_buffers(engine):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
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__()
for binding in engine:
size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if 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 load_engine(trt_path):
# 反序列化引擎
with open(trt_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
class TRTInference(object):
"""Manages TensorRT objects for model inference."""
def __init__(self, trt_engine_path, onnx_model_path, trt_engine_datatype=trt.DataType.FLOAT, batch_size=1):
"""Initializes TensorRT objects needed for model inference.
Args:
trt_engine_path (str): path where TensorRT engine should be stored
uff_model_path (str): path of .uff model
trt_engine_datatype (trt.DataType):
requested precision of TensorRT engine used for inference
batch_size (int): batch size for which engine
should be optimized for
"""
# Initialize runtime needed for loading TensorRT engine from file
# TRT engine placeholder
self.trt_engine = None
# Display requested engine settings to stdout
print("TensorRT inference engine settings:")
print(" * Inference precision - ".format(trt_engine_datatype))
print(" * Max batch size - \\n".format(batch_size))
# If we get here, the file with engine exists, so we can load it
if not self.trt_engine:
print("Loading cached TensorRT engine from ".format(
trt_engine_path))
self.trt_engine = load_engine(
trt_engine_path)
# This allocates memory for network inputs/outputs on both CPU and GPU
self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(self.trt_engine)
# Execution context is needed for inference
self.context = self.trt_engine.create_execution_context()
def infer(self, full_img, output_shapes, new_width, new_height):
"""Infers model on given image.
Args:
image_path (str): image to run object detection model on
"""
assert new_width > 0 and new_height > 0, "Scale is too small"
# resize and transform to array
scale_img = cv2.resize(full_img, (new_width, new_height))
print("scale image shape:".format(scale_img.shape))
# scale_img = np.array(scale_img)
# HWC to CHW
scale_img = scale_img.transpose((2, 0, 1))
# 归一化
if scale_img.max() > 1:
scale_img = scale_img / 255
# 扩增通道数
# scale_img = np.expand_dims(scale_img, axis=0)
# 将数据成块
scale_img = np.array(scale_img, dtype=np.float32, order='C')
# Copy it into appropriate place into memory
# (self.inputs was returned earlier by allocate_buffers())
np.copyto(self.inputs[0].host, scale_img.ravel())
# Output shapes expected by the post-processor
# output_shapes = [(1, 11616, 4), (11616, 21)]
# When infering on single image, we measure inference
# time to output it to the user
inference_start_time = time.time()
# Fetch output from the model
trt_outputs = do_inference(
self.context, bindings=self.bindings, inputs=self.inputs,
outputs=self.outputs, stream=self.stream)
print("network output shape:".format(trt_outputs[0].shape))
# Output inference time
print("TensorRT inference time: ms".format(
int(round((time.time() - inference_start_time) * 1000))))
# Before doing post-processing, we need to reshape the outputs as the common.do_inference will
# give us flat arrays.
outputs = [output.reshape(shape) for output, shape in zip(trt_outputs, output_shapes)]
# And return results
return outputs
# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
ps：这里是完整的推理代码，网上很多资料都不全，所以贴上去了。
predict.py
根据实际情况需要设置的参数：
engine_file_path：engine的文件路径
onnx_file_path：onnx文件路径
new_width, new_height: 输入的宽和高
trt_engine_datatype：engine的精度支持fp32和fp16
image_path：测试图片路径
import tensorrt as trt
import numpy as np
import cv2
import inference as inference_utils # TRT/TF inference wrappers
if __name__ == "__main__":
# 1. 网络构建
# Precision command line argument -> TRT Engine datatype
TRT_PRECISION_TO_DATATYPE =
16: trt.DataType.HALF,
32: trt.DataType.FLOAT
# datatype: float 32
trt_engine_datatype = TRT_PRECISION_TO_DATATYPE[16]
# batch size = 1
max_batch_size = 1
engine_file_path = "unet_deconv_sim.trt"
onnx_file_path = "unet_deconv_sim.onnx"
new_width, new_height = 960, 640
output_shapes = [(1, new_height, new_width)]
trt_inference_wrapper = inference_utils.TRTInference(
engine_file_path, onnx_file_path,
trt_engine_datatype, max_batch_size,
)
# 2. 图像预处理
image_path = "example.jpg"
img = cv2.imread(image_path)
# inference
trt_outputs = trt_inference_wrapper.infer(img, output_shapes, new_width, new_height)[0]
# 输出后处理
out_threshold = 0.5
print("the size of tensorrt output : ".format(trt_outputs.shape))
output = trt_outputs.transpose((1, 2, 0))
# 0/1像素值
output[output > out_threshold] = 255
output[output <= out_threshold] = 0
output = output.astype(np.uint8)
result = cv2.resize(output, (img.shape[1], img.shape[0]))
cv2.imwrite("best_output_deconv.jpg", result)
最后结果：
因为这个网络只是用来测试能否部署tensorrt。为了节省时间，训练模型只使用了一个epoch，故精度不高。
第三章：Pytorch-grcnn to Tensorrt-grcnn
antipodal robotic grasping网络输入一张224x224x(1,3,4)的图片，这张图片可以是rgbd类型，rgb类型和d类型，输出抓取pos，抓取sin，抓取cos和抓取夹爪开度。其github项目在此，这里讲解输入rgb图片网络转换的方法。
Step1：训练rgb输入pytorch网络
这个步骤就不讲解了，其github官网上readme有教程，去看一看。
Step2：将模型文件转换为onnx文件
其onnxtotrt.py代码内容如下：
由于输入图片为rgb是三个通道的，输入类型即为（3，224，224）于是需要修改dummy_input的内容，再初始化其模型参数（输入通道，通道个数）等等。这里需要注意的是由于其训练模型不是按.pt文件保存的，所以这里在加载权重时只需要调用torch.load即可。其余与第二章的转换一致。
Step3：将onnx转换为.trt文件
由于篇幅原因，这里直接贴上代码：
onnx2trt grcnn.onnx -o grcnn.trt
Step4：编写推理代码
得到trt文件后，开始写inference和predict代码。
inferencetest.py中的重点：
由于网络模型输出有4个，所以在do_inference得到output后，需要通过切片操作，分别取4个结果。
predict.py中的重点：
这里的输入图片预处理与原网络中run_offline.py输入图片一致，只是将深度设置成了false
run_offline.py
与pytorch不同的是，要部署tensorrt，需要将里面的predict函数换成tensorrt的推理函数，其余不变。对于其后处理函数也与原网络相同，只是需要将pytorch所支持的tensor变量操作改为numpy操作。
最后得出结果如下图
通过评估代码，将数据集中所有图像用于验证，得到原推理速度为33ms每一张图片：
经过tensorrt加速后的推理速度为：
可以看出还是有加速效果的（博主电脑拉）估计在台式机上会更好。此外，原网络占用显存会比经过tensorrt加速后网络占用显存大很多，博主忘了截图了就口头说明一下。其经过tensorrt开源代码在此。不要白嫖~点个star再克隆吧！
在 ubuntu 或 windows 上运行的 MS Edge 进行端到端测试有啥区别吗？
】在ubuntu或windows上运行的MSEdge进行端到端测试有啥区别吗？【英文标题】：ArethereanydifferencesbetweenMSEdgerunningonubuntuorwindowsforEnd-to-endtesting?在ubuntu或windows上运行的MSEdge进行端到端测试有什么区别吗？【发布时间】：2021-09-0304:56:39... 查看详情
volte端到端业务详解|终端能力类(代码片段)
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客2.6.1S1AP协议1.UE网络能力这个信... 查看详情
volte端到端业务详解|掉话问题
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客10.2.1基本流程传统移动通信网... 查看详情
iperf3 在底层使用啥算法来测量网络基础设施中端到端之间的带宽和延迟？
】iperf3在底层使用啥算法来测量网络基础设施中端到端之间的带宽和延迟？【英文标题】：Whatalgorithmiperf3useunderthehoodtomeasurethebandwidthandlatencybetweenendtoendpointsinanetworkinfrastructure?iperf3在底层使用什么算法来测量网络基础设施中端到... 查看详情
volte端到端业务详解|编解码协商流程
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客由于VoLTE网络中不同网元、不同... 查看详情
中科大利用神经网络端到端训练
...种教育情境感知的认知诊断框架，使用神经网络以及端到端的训练框架，自适应学习不同教育情境信息的量化影响，并结合现有认知诊断 查看详情
漫谈无丢包网络拥塞控制和端到端原则
将一条公网链路用TCP隧道分割成多个段，每一个段根据该段的网络特征使用不同的拥塞控制算法，然后把多条如此分割的长链路合并在一起，这就是一个overlay网络，SDWAN数据面核心不过如此。由于使用TCP搭建隧道&... 查看详情
volte端到端业务详解|基本概念
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客2G/3G网络的掉话一般是在通话过... 查看详情
fasterrcnn近似端到端法详解
有了交替训练法的基础，近似端到端的训练方法就非常好理解了。相比于交替训练法，端到端训练主要有两点不一样：1）网络结构：RPN的输出个数：此处的输出个数不是指损失函数，而是指RPN的输出不... 查看详情
为啥faster-rcnn端到端训练只做近似？
】为啥faster-rcnn端到端训练只做近似？【英文标题】：Whyfaster-rcnnendtoendtrainingonlymakesapproximation?为什么faster-rcnn端到端训练只做近似？【发布时间】：2018-05-0107:07:05【问题描述】：在更快的rcnn(https://arxiv.org/abs/1506.01497)中，有两种... 查看详情
volte端到端业务详解|volte基本概念
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客5.1.1域选由于支持VoLTE的UE可以... 查看详情
端到端图像编码和vvc的结合
...》简介提案提出了一种混合编码方法，对视频I帧使用端到端的图像编码器编码，P和B帧使用VVC编码，如Fig.1，其中P和B帧内的intra块还是使用VVC编码。基于神经网络的端到端图像编码器 查看详情
volte端到端业务详解|volte用户呼叫流程
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客5.3.1主叫流程如图5-37所示，... 查看详情
volte端到端业务详解|volte用户注册流程
书籍来源：艾怀丽《VoLTE端到端业务详解》一边学习一边整理书中的笔记，并与大家分享，侵权即删，谢谢支持！附上汇总贴：VoLTE端到端业务详解|汇总_COCOgsta的博客-CSDN博客注册是用户向签约网络请求授... 查看详情
orange选择mavenir5g独立组网openran用于法国端到端实验云网络
...其5GOpenRAN软件，以用于欧洲首个5G独立组网(StandAlone)端到端云网络。位于法国布列塔尼大区拉尼永(Lannion)的实验网络完全 查看详情
fasterrcnn近似端到端法详解
有了交替训练法的基础，近似端到端的训练方法就非常好理解了。相比于交替训练法，端到端训练主要有两点不一样：1）网络结构：RPN的输出个数：此处的输出个数不是指损失函数，而是指RPN的输出不... 查看详情
（转）一文学会用tensorflow搭建神经网络
一文学会用Tensorflow搭建神经网络 本文转自：http://www.jianshu.com/p/e112012a4b2d字数2259 阅读3168 评论8 喜欢11cs224d-Day6:快速入门Tensorflow本文是学习这个视频课程系列的笔记，课程链接是youtube上的，讲的很好，浅显易... 查看详情
分层测试：端到端测试
...系列文章https://www.cnblogs.com/yuxiuyan/tag/分层测试/1.什么是端到端测试端到端测试（End-To-EndTesting,简称E2E测试）是一种从头到尾测试整个软件产品以确保应用程序流程按预期运行的技术。它定义了产品的系统依赖性，并确保所有集... 查看详情