极智AI | TensorRT Parser 构建模型推理方法

  一起养成写作习气！这是我参与「日新计划 6 月更文挑战」的第14天，点击检查活动详情。

  欢迎重视我的大众号 [极智视界]，获取我的更多笔记共享

  大家好，我是极智视界，本文介绍一下 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 | TensorRT Parser 构建模型推理办法》