极智AI | TensorRT Parser 构建模型推理方法
Posted 极智视界
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了极智AI | TensorRT Parser 构建模型推理方法相关的知识,希望对你有一定的参考价值。
欢迎关注我的公众号 [极智视界],获取我的更多笔记分享
大家好,我是极智视界,本文介绍一下 TensorRT Parser 构建模型推理方法。
TensorRT 构建模型推理一般有三种方式:(1) 使用框架自带的 TensorRT 接口,如 TF-TRT、Torch-TRT;(2) 使用 Parser 前端解释器,如 TF / Torch / … -> ONNX -> TensorRT;(3) 使用 TensorRT 原生 API 搭建网络。当然难度和易用性肯定是由低到高的,伴随而来的性能和兼容性也是由低到高的。这里我们介绍第二种方式,使用 Parser 前端解释器来构建 TensorRT 模型推理,会分别用现在最主流的 pytorch 和 tensorflow 来进行示例介绍。
文章目录
1 TensorRT Parser - pytorch
基本流程:
(1) pytorch 中创建网络并保存为 .pt 模型;
(2) 使用 pytorch 内部 API 将 .pt 模型转换为 .onnx 模型;
(3) TensorRT 中读取 .onnx 模型构建 Engine 并做推理;
上代码:
# pytorch 中创建网络并保存为 .pt 模型文件
# ...
t.save(net, ptFile)
print("successed building model in pytorch")
# 将 .pt 模型文件转换为 .onnx 模型文件
t.onnx.export(
net,
t.randn(1, 1, h, w, device="cuda"),
"./model.onnx",
examples_outputs = [t.randn(1, 10, device="cuda"), t.randn(1, device="cuda")],
input_names=['x'];
output_names=['y', 'z'],
do_constant_folding=True,
verbose=True,
keep_initializers_as_inputs=True,
opset_version=12,
dynamic_axes="x": 0: "nBatchSize", "z": 0: "nBatchSize"
)
print("successed converting model into onnx")
# tensorrt 中加载 .onnx 创建 engine
logger = trt.Logger(trt.Logger.ERROR)
# ...
# 用 Parser 加载 .onnx
with open(onnxFile, 'rb') as model:
if not parser.parse(model.read()):
print("filed parsing onnx file")
for error in range(parser.num_errors):
print(parser.get_error(error))
exit()
print("successed paring onnx file")
# ..
# 准备 tensorrt runtime 和 buffer,进行推理
context = engine.create_execution_context()
# ...
print("successed running model in tensorrt")
以上展示了 pytorch -> onnx -> parser -> tensorrt infer 的流程。
2 TensorRT Parser - tensorflow
基本流程:
(1) tensorflow中创建网络并保存为 .pt 模型;
(2) 使用 tf2onnx 将 .pb 模型转换为 .onnx 模型;
(3) TensorRT 中读取 .onnx 模型构建 Engine 并做推理;
上代码:
# tensorflow 中创建网络并保存为 .pb 模型
x = tf.compat.v1.placeholder(tf.float32, [None, 28, 28, 1], name='x')
# ...
# 保存为 .pb 模型
constantGraph = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['z'])
with tf.gfile.FastGFile("./model.pb", mode='wb') as f:
f.write(constantGraph.SerializeToString())
sess.close()
print("successed building model in tensorflow")
# 将 .pb 模型转换为 .onnx 模型
os.system("python -m tf2onnx.convert --input %s --output %s --inputs 'x:0' --outputs 'z:0'" % (pbFile, onnxFile))
print("successed converting model into onnx")
# tensorrt 中加载 .onnx 创建 engine
logger = trt.Logger(trt.Logger.ERROR)
# ...
# 用 Parser 加载 .onnx
with open(onnxFile, 'rb') as model:
if not parser.parse(model.read()):
print("filed parsing onnx file")
for error in range(parser.num_errors):
print(parser.get_error(error))
exit()
print("successed paring onnx file")
# ..
# 准备 tensorrt runtime 和 buffer,进行推理
context = engine.create_execution_context()
# ...
print("successed running model in tensorrt")
以上展示了 tensorflow -> onnx -> parser -> tensorrt infer 的流程,可以看到从 parser 解析 onnx 后面这段和 pytorch 那段是一致的。
好了,以上分享了 TensorRT Parser 构建模型推理的方法,希望我的分享能对你的学习有一点帮助。
【公众号传送】
扫描下方二维码即可关注我的微信公众号【极智视界】,获取我的更多经验分享,让我们用极致+极客的心态来迎接AI !
以上是关于极智AI | TensorRT Parser 构建模型推理方法的主要内容,如果未能解决你的问题,请参考以下文章
我的NVIDIA开发者之旅 - 极智AI | TensorRT API 构建模型推理流程
我的NVIDIA开发者之旅 - 极智AI | TensorRT 中 Layer 和 Tensor 的区别
极智AI | 讲解 TensorRT 显式batch 和 隐式batch
极智AI | 讲解 TensorRT Constant 算子