本文详细介绍了如何将基于TensorFlow训练的模型转换为ONNX格式，然后利用TensorRT进行优化，最终部署在Jetson Nano上进行边缘计算。过程中涉及模型保存、ONNX转换、ONNX模型验证、TRT引擎构建以及推理代码实现，并解决了动态输入的问题，实现了TensorFlow模型在嵌入式设备上的高效推理。
项目背景：
把tensorflow训练好的模型部署到嵌入式设备jetson nano上，做边缘计算。
TensorRT简介
方案
tensorflow -> onnx -> trt
-
tensorflow：训练好的模型要保存为saved_model格式（在【04】keras高级API 文章中有说明）
-
onnx：ONNX | Home
-
trt：部署在jetson nano上推理的引擎
步骤
保存训练好的模型
tf.saved_model.save(model, 'path')
path目录结构如下
.
├── assets
├── saved_model.pb
└── variables
├── variables.data-00000-of-00001
└── variables.index
saved_model 转 onnx
点击链接，跳转到官方文档
tensorflow-onnx/README.md at master · onnx/tensorflow-onnx · GitHub
安装必要的库：
pip install tensorflow # tensorflow这里使用2.3.0版本，具体安装方法参考其他资料，这里不赘述
pip install onnxruntime # 这是运行onnx文件的库，用来验证转换后的模型是否正确（可选）
pip install -U tf2onnx # 转换
执行命令
python -m tf2onnx.convert --saved-model tensorflow-model-path --opset 11 --output model.onnx
'''
;tensorflow-model-path 这个是saved_model的路径
;opset 这个值默认是9，这里选11
;model.onnx 输出onnx的文件名
'''
如果顺利的话，一般在当前文件夹下可以看到生成了一个新文件model.onnx。如果出了问题可以查看issue或者google。
借助工具可以查看onnx的内容
验证转换后的onnx的模型的正确性
import onnx
import onnxruntime as ort
import numpy as np
import cv2
# "加载load"
model=onnx.load('model.onnx')
# 检查模型格式是否完整及正确
try:
onnx.checker.check_model(model)
except onnx.checker.ValidationError as e:
print('The model is invalid: %s' % e)
else:
print('The model is valid!')
ort_session = ort.InferenceSession('model.onnx')
# 图片预处理
cv_img = cv2.imread("image.jpg") # 读取图片
cv_img = cv2.resize(cv_img,(512, 512),fx=0,fy=0,interpolation=cv2.INTER_LINEAR) # 维度变换
cv_images = cv_img.astype('float32') # 数据类型转换
cv_images = cv_images / 127.5 - 1 # 归一化到 -1 到 1 之间
cv_images = np.expand_dims(cv_images, 0) # 增加一个维度
# 推理
ort_inputs = {ort_session.get_inputs()[0].name: cv_images}
ort_outs = ort_session.run(None, ort_inputs)
print(ort_outs[0])
# 图片后处理
一般验证ort_outs[0]，与tensorflow预测的值一样即可。
onnx 转 trt
命令形式
首先安装好TensorRT，这里的版本是7.1.3.0。
在TensorRT的目录下面有一个bin文件夹，里面有一个trtexec的文件，使用这个文件来进行转换。（如果是ubuntu18.4，TensorRT的路径一般在/usr/src/tersorrt）
进入到上面的路径，执行命令
./trtexec --onnx=/home/nano/model.onnx --saveEngine=/home/nano/model.trt --workspace=6000
'''
;--onnx 这个是model.onnx的路径
;--saveEngine 这个是trt引擎保存的路径，文件是trt格式的
;--workspace 是内存分配空间 单位MiB
'''
如果顺利的话，你将可以得到一个model.trt文件，这个文件就可以部署在jetson nano上，提升模型的推理速度。但是，大概率会报错：
- Your ONNX model has been generated with INT64 weights, while TensorRT does not natively support INT64. Attempting to cast down to INT32.
这只是个警告，可以忽略，不影响文件输出，以及后续的推理。
- Network has dynamic or shape inputs, but no optimization profile has been defined.
网络具有动态或形状输入，但尚未定义优化配置文件。
用上面打开onnx模型的网站，看到我们的输入输出的形状是：输入(None, 512, 512, 3)、输出(None, 512, 512, 1)，这个None就是他说的动态输入。
（打开onnx2trt的github地址https://github.com/onnx/onnx-tensorrt ，在他的issue中可以找到解决方案：https://github.com/onnx/onnx-tensorrt/issues/518）
您必须使用优化配置文件：https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work_dynamic_shapes
没错，我们需要开始啃官方文档了。
用代码转换
首先需要在TensorRT的官方示例中找到一个common.py文件，这个将在后续转换和推理中需要用到。
打开TensorRT的开发者指南，找到第四章Python API。其中4.1小节是onnx转trt的教程。动态输入解决办法在第8.2小节。经过不断的阅读理解官方文档，于是有了下面的代码onnx2trt.py。可以对照的官方文档一句一句的理解。
import tensorrt as trt
import common
'''
通过加载onnx文件，构建engine
'''
onnx_file_path = "model.onnx"
G_LOGGER = trt.Logger(trt.Logger.WARNING)
# 1、动态输入第一点必须要写的
explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
batch_size = 1 # trt推理时最大支持的batchsize
with trt.Builder(G_LOGGER) as builder, builder.create_network(explicit_batch) as network, \
trt.OnnxParser(network, G_LOGGER) as parser:
builder.max_batch_size = batch_size
config = builder.create_builder_config()
config.max_workspace_size = common.GiB(1) # common文件可以自己去tensorrt官方例程下面找
config.set_flag(trt.BuilderFlag.TF32)
print('Loading ONNX file from path {}...'.format(onnx_file_path))
with open(onnx_file_path, 'rb') as model:
print('Beginning ONNX file parsing')
parser.parse(model.read())
print('Completed parsing of ONNX file')
print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
# 动态输入问题解决方案
profile = builder.create_optimization_profile()
profile.set_shape("input_1", (1, 512, 512, 3), (1, 512, 512, 3), (1, 512, 512, 3))
config.add_optimization_profile(profile)
engine = builder.build_engine(network, config)
print("Completed creating Engine")
# 保存engine文件
engine_file_path = 'model_fp32.trt'
with open(engine_file_path, "wb") as f:
f.write(engine.serialize())
然后运行这个文件python onnx2trt.py。然后你就有了model_fp32.trt文件。如果遇到其他问题，找issue或者google。
参考链接：
【1】官方手册
https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#python_topics
【2】参考示例代码
https://github.com/NVIDIA/TensorRT/blob/main/samples/python/introductory_parser_samples/onnx_resnet50.py
【3】博客
pytorch模型-转-onnx-转-tensorrt实际操作详解_pangxing6491的博客-优快云博客
TensorRT 7 动态输入和输出_我是一个菜鸟，虚心学习的菜鸟。-优快云博客
input只允许输入日期_tensorrt动态输入（Dynamic shapes）_hateful harmful的博客-优快云博客
推理
有了trt文件后，就可以调用它进行推理了。参考官方文档和示例，python代码如下：
import cv2
import numpy as np
import os
import tensorrt as trt
import time
import common
TRT_LOGGER = trt.Logger()
engine_file_path = "model_fp32.trt"
with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime, \
runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:
inputs, outputs, bindings, stream = common.allocate_buffers(engine)
print(inputs, outputs, bindings, stream)
# 前处理部分
t1 = time.clock()
cv_img = cv2.imread("image030.jpg")
cv_img = cv2.resize(cv_img, (512, 512), fx=0, fy=0, interpolation=cv2.INTER_LINEAR)
cv_images = cv_img.astype('float32')
cv_images = cv_images / 127.5 - 1 # 归一化到 -1 到 1 之间
cv_images = np.expand_dims(cv_images, 0)
# 前处理结束
# 开始推理
inputs[0].host = cv_images
trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
print(trt_outputs.shape)
print("Time:", time.clock() - t1)
# 由于trt_outputs是一个一维的list，而我们需要的输出是(1, 512, 512, 2)的张量，所以要整形
trt_output = [output.reshape(shape) for output, shape in zip(trt_outputs, [(1, 512, 512, 2)])]
这里比较坑的地方就是输出是list，要经过整形转换成我们需要的维度的张量。
经过测试，在TensorRT推理的结果和在电脑上用tensorflow预测的结果差不多。
参考链接：
【1】官方手册
https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#python_topics
【2】参考示例代码
https://github.com/NVIDIA/TensorRT/blob/main/samples/python/yolov3_onnx/onnx_to_tensorrt.py
【3】博客
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
