TensorRT 加速 Keras 模型在 Jetson 上的推理

将一个训练好的 Keras 模型通过 TensorRT 加速并 Push 到 Jetson TX2 上的流程框图如下：

下面对一些关键代码及步骤进行解释：

Keras model to Tensorflow frozen graph

这一步可以在任意一台机器上完成，不限于 Jetson TX2 或者其 Host PC 上，只要配置了 tensorflow 和 keras 即可。

1234567891011121314151617181920212223242526272829303132333435

import tensorflow as tf from keras.models import model_from_jsonfrom keras import backend as Kimport os def (model_file, weights_file): with open(model_file, 'r') as f: json_string = f.read() K.set_learning_phase(0) model = model_from_json(json_string) model.load_weights(weights_file) # rename output nodes output_node_prefix = 'output_node_' output_node_names = [] output_nodes = [] for i in range(len(model.outputs)): output_node_name = output_node_prefix + str(i) output_nodes.append(tf.identity(model.outputs[i], name=output_node_name)) output_node_names.append(output_node_name) print('output node names are: ', output_node_names) # freeze the graph sess = K.get_session() constant_graph = tf.graph_util.convert_variables_to_constants(sess=sess, input_graph_def=sess.graph.as_graph_def(), output_node_names=output_node_names) # write graph as pb output_dir = 'frozen_graph' if not os.path.isdir(output_dir): os.mkdir(output_dir) filename = 'model.pb' return tf.train.write_graph(constant_graph, output_dir, filename, as_text=False)

代码的关键要点：

1. K.set_learning_phase(0) 将学习阶段设置为 0， 即 test 阶段，像 BatchNormalization 这种 layer 在 test 阶段和 train 阶段的行为是不一样的。
2. tf.graph_util.convert_variables_to_constants 将变量转换为常量，这一步使得我们可以用一个文件完全地（fully）表达一个神经网络，即如果不做这一步的话，tf.train.write_graph 写到文件里的仅仅有网络结构而没有相应的训练得到的权重参数。

Tensorflow frozen graph to UFF

就目前[2018/08/03]来说，NVCaffe 框架搭建的神经网络可以直接被导入并转化为推理引擎，来自其他框架的神经网络模型可以转化为 UFF 或者 ONNX 格式，然后通过相应的解析器解析成推理引擎。

将 Tensorflow frozen graph 转化为 UFF 文件，需要安装 uff Python 包。
这一步可以在 Jetson TX2 或者其 Host PC 上完成，只要安装了 uff 包即可。

uff 包是 TensorRT 3.0.4 for Ubuntu 16.04 and CUDA 9.0 tar package 的一部分，因此需要下载这个 tar 包，同时这个 tar 包也可以用来安装 TensorRT，由于已经通过 JetPack 在 Jetson TX2 上安装了 TensorRT，这一点可以在上一篇文章或者下图中看到，因此我在这里仅取用安装其中的 uff。

12	tar -xzf TensorRT-3.0.4.Ubuntu-16.04.3.x86_64.cuda-9.0.cudnn7.0.tar.gzsudo pip install TensorRT-3.0.4/uff/uff-0.2.0-py2.py3-none-any.whl

12345678

import uff# generate uff from frozen graphuff_model = uff.from_tensorflow_frozen_model( frozen_file=frozen_graph_filename, output_nodes=output_node_names, output_filename=uff_model_filename, text=False,)

UFF to plan (engine)

这一步在安装了 TensorRT 的机器上完成，在此文章中为 Jetson TX2。
这一步可以参考 Nvidia Two Days to a Demo Advanced - TensorFlow to TensorRT Image Classification 的 github 仓库 中的 src/uff_to_plan.cpp。
但是需要注意的是，该文件的 71 行方法调用存在错误：

1	parser->registerInput(inputName.c_str(), DimsCHW(3, inputHeight, inputWidth));

编译时报错：

ȥ /usr/include/aarch64-linux-gnu/NvUffParser.h 查看函数声明，发现需要一个 UffInputOrder 类型的参数:

在这个头文件内搜索可以找到 UffInputOrder 这一个枚举类型:

修改 src/uff_to_plan.cpp 的行 71 为

1	parser->registerInput(inputName.c_str(), DimsCHW(3, inputHeight, inputWidth), UffInputOrder::kNHWC);

即可。

predict.cu

这一步在安装了 TensorRT 的机器上完成，在此文章中为 Jetson TX2，即在 Jetson TX2 完成 .cu 文件的编译。

1234567891011121314151617181920212223242526272829303132333435

#include <fstream>#include <sstream>#include <NvInfer.h>using namespace std;using namespace nvinfer1;class Logger : public ILogger{ void log(Severity severity, const char * msg) override { if (severity != Severity::kINFO) cout << msg << endl; }} gLogger;/* load the engine */cout << "Loading TensorRT engine from plan file..." << endl;ifstream planFile(planFilename); if (!planFile.is_open()){ cout << "Could not open plan file." << endl; return 1;}stringstream planBuffer;planBuffer << planFile.rdbuf();string plan = planBuffer.str();IRuntime *runtime = createInferRuntime(gLogger);ICudaEngine *engine = runtime->deserializeCudaEngine((void*)plan.data(), plan.size(), nullptr);

原文:大专栏  TensorRT 加速 Keras 模型在 Jetson 上的推理

来源：博客园

作者：三鲜豆皮

链接：https://www.cnblogs.com/sanxiandoupi/p/11631912.html