梳理TensorFlow模型在Jetson TX2上进行inference的主要流程

1 内容概述

本文档主要内容：梳理TensorFlow模型在Jetson TX2上进行inference的主要流程，涉及到相关软件的安装、依赖库的编译配置以及PC端的深度网络模型在Jetson TX2的移植步骤。

2 开发环境

注意：PC端和Jetson TX2所使用的TensorRT安装包是不一样的，前者使用的是Tesla GPUs版本，而后者使用的是Jetson Platforms版本

2.1 PC端

1) Pycharm：Python程序的开发环境；

2) Python 3.5：

3) TensorFlow：数值计算的开源软件库，主要用于深度学习模型的搭建；

4) TensorRT Python API：在PC端使用TensorRT所需的Python库，以whl文件形式存在，存放在TensorRT安装包的python目录下；

5) UFF Python API：转换成uff文件的Python库，存放在TensorRT安装包的uff目录下。

2.2 Jetson TX2端

1) TensorRT 3.0：NVIDIA提供的API，包括c++ API和Python API，用于加速基于NVIDIA gpu硬件平台的深度网络模型inference程序，目前主要支持caffe、TensorFlow模型的转换；

2) CUDA 8.0，cuDNN 7.0;

3) OpenCV 3.3：图像处理开源库，在Jetson TX2平台上需使用cmake方式进行编译安装；

3 深度网络模型移植

本文移植主要针对利用TensorFlow搭建的网络模型，下面用图例简单介绍深度网络模型从training到inference的整个流程，图中NVIDIA DIGITS为NVIDIA开发的方便模型训练的web应用：

 

从PC端训练模型到TX2进行inference，主要有三个步骤：

1) 训练模型并保存为pb文件；

2) Pb文件转uff文件；

3) 利用TensorRT加载并运行模型；

前两个步骤是在PC端执行，而最后一个步骤是在Jetson TX2上进行。本教程引用了TensorRT-3-User-Guide.pdf文档中的部分内容。

3.1 训练模型保存为pb文件（PC端）

Pb文件包含网络结构和模型参数，将TensorFlow训练得到的模型保存为pb的方法参考网址。

保存pb文件最核心的代码是：

constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ["output"])

with tf.gfile.FastGFile(pb_file_path, mode='wb') as f:

f.write(constant_graph.SerializeToString())

3.2 pb文件转uff文件（PC端）

为了将生成的pb格式的模型文件转为TensorRT 3.0可以解析的uff文件，需在命令行执行以下命令：

convert-to-uff tensorflow -o name_of_output_uff_file --input_file  name_of_input_pb_file  -O  name_of_output_tensor

加粗部分是需要更改的，name_of_output_uff_file是即将输出的uff文件名字，

name_of_input_pb_file是输入的pb模型文件名字，name_of_output_tensor是pb模型图中输出的节点名字。

可通过以下命令得到网络各个节点的名称：

convert-to-uff tensorflow --input_file name_of_input_pb_file –l

注：在由pb文件转换至uff文件时，必须指定输出节点名称。

 

3.3 加载并运行模型（Jetson端）

在运用TensorRT进行inference主要包括两个阶段：build（构建网络）和deployment（运行），编程语言为c++，技术文档请参考doc路径下TensorRT-3-User-Guide.pdf中章节3.3： SampleUffMNIST UFF Usage。示例程序请参考sample/ sampleUffMNIST路径下的程序文件sampleUffMNIST.cpp。

3.3.1 构建阶段（build）

1) 创建builder

// 创建network解析变量parser

auto parser = createUffParser();     

// 为parser指定网络模型的输入和输出

parser->registerInput("Input_0", DimsCHW(1, 28, 28));     

parser->registerOutput("Binary_3");

// 创建builder变量

IBuilder* builder = createInferBuilder(gLogger);

INetworkDefinition* network = builder->createNetwork();

2) 解析uff model到network

parser->parse(uffFile, *network, nvinfer1::DataType::kFLOAT)；

uffFile是uff文件的路径

3) 创建engine

builder->setMaxBatchSize(maxBatchSize);     // 设置输入数据的batch

builder->setMaxWorkspaceSize(MAX_WORKSPACE);  

// 设置engine运行时空间（即内存）大小，这里MAX_WORKSPACE可设置为1<<30。

ICudaEngine* engine = builder->buildCudaEngine(*network);  // engine返回值

3.3.2 执行阶段（deployment）

1) 创建execution context

IExecutionContext* context = engine.createExecutionContext();

2) 获取engine的binding个数以及索引

int nbBindings = engine.getNbBindings();

3) 计算输入、输出节点维度大小和数据类型

auto buffersSizes = calculateBindingBufferSizes(engine, nbBindings, batchSize);

4) 为输入、输出节点分配CUDA内存

参考sampleUffMNIST.cpp中的safeCudaMalloc函数，得到buffers数组，包括输入节点和输出节点的内存块；

5) 将输入数据导入输入节点的内存中

参考createMnistCudaBuffer（）函数，可根据需要进行改写；

6) 执行inference

context->execute(batchSize, buffers);  // buffers为存储输入、输出节点内存的数组；

7) 输出结果

根据上一步得到输出节点数据，即buffers[1]，一般为分类结果，或者语义分割的输出数据。根据需要，进行后续处理。