TensorRT加速tensorflow实例

准备工作

  1. 生成.pb的模型文件

首先我们需要从保存模型的chekpoint文件中，生成.pb的模型文件。这一步叫做模型的持久化，具体的做法可以参考之前写的这篇文章：
春天不是读书天：[深度学习] TensorFlow中模型的freeze_graph
2. 导入必要的库

  1. import tensorflow as tf
  2. import uff
  3. import tensorrt as trt
  4. import pycuda.driver as cuda 
  5. import pycuda.autoinit
  6. from tensorrt.parsers import uffparser
  • uff：是将刚才的pb转化为TensorRT引擎支持的uff文件，该文件可以序列化，也可以直接当作流传过去。
  • pycyda：用于显卡cuda编程的，如果要使用TensorRT的python API，这是一个必须的库
  • uffparser ：用于解析uff模型
  1. 参数设置
    1. MODEL_DIR = './model_seg/model.pb'
    2. CHANNEL = 3
    3. HEIGHT = 299
    4. WIDTH = 299
    5. ENGINE_PATH = './model_seg/model_.pb.plan'
    6. INPUT_NODE = 'input'
    7. OUTPUT_NODE = ['InceptionV3/Logits/SpatialSqueeze']
    8. INPUT_SIZE = [CHANNEL, HEIGHT ,WIDTH] 
    9. MAX_BATCH_SIZE = 1 
    10. MAX_WORKSPACE = 1<<30
  • MODEL_DIR：第一步中生成的pb模型地址
  • CHANNEL、HEIGHT、WIDTH：图片的通道、高和宽，根据模型的输入大小确定
  • ENGINE_PATH：等会保存TensorRT引擎的地址
  • INPUT_NODE：模型的输入节点
  • OUTPUT_NODE：模型的输出节点，是一个列表，如果有许多个输出节点，就将节点名都列入这个列表中
  • INPUT_SIZE：输入图片的大小，注意通道在前还是后，这里输入的是 CHANNEL, HEIGHT ,WIDTH
  • MAX_BATCH_SIZE：在推理的时候，每次输入几张图片
  • MAX_WORKSPACE：显存的大小1<<30也就是1GB的大小。有的时候，程序运行是会报内存溢出的错，这个时候就可以调小MAX_WORKSPACE，比如2 << 10

    将tensorflow模型转换成TensorRT

  1. pb转uff 并解析模型
    1. G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO)
    2. uff_model = uff.from_tensorflow_frozen_model(FROZEN_GDEF_PATH, OUTPUT_NODE)
    3. parser = uffparser.create_uff_parser()
    4. parser.register_input(INPUT_NODE, INPUT_SIZE, 0)
    5. parser.register_output(OUTPUT_NODE)
    这里做的事情是将pb的文件格式转成了uff文件格式。你需要知道的一个概念是，UFF(Universal Framework Format)是一种描述DNN执行图的数据格式。绑定执行图的是输入与输出，所以parser.register_input和parser.register_output做的事情是将tensorflow模型的输入输出在UFF文件中记录。
    注意，对于多个输出，因为OUTPUT_NODE是一个列表，所以将多个输出节点依次放入列表就可以了。
    如果是多个输入的话，则需要将输入节点名一个个的记录在uff中。register_input()需要3个参数：
  • name – Input name.
  • shape – Input shape.
  • order – Input order on which the framework input was originally.

假设你的模型在输入层同时输入了三张图片，那么你需要定义3个输入节点，并且指定order分别为0、1、2。这里的order指的是模型的输入在uff结构中的顺序，这种order在接下来的binding会得到体现。

  1. parser.register_input(INPUT_NODE1, INPUT_SIZE, 0)
  2. parser.register_input(INPUT_NODE2, INPUT_SIZE, 1)
  3. parser.register_input(INPUT_NODE3, INPUT_SIZE, 2)
  1. 保存模型
    1. engine = trt.utils.uff_to_trt_engine(
    2.      G_LOGGER,
    3.      uff_model,
    4.      parser,
    5.      MAX_BATCH_SIZE,
    6.      MAX_WORKSPACE,
    7.      datatype=trt.infer.DataType.FLOAT)
    以上代码创建了TensorRT中的engine，即引擎，这个engine将负责模型的前向运算。TensorRT是一个用于推理的加速工具，所以前向计算就够了。
    在engine创建成功之后，就可以使用了。不过，一个建议是将结果保存下来。毕竟到目前为止，虽然代码很少，但是将pb文件成功转换成uff文件是不容易的（谁用谁知道！）
    使用以下语句，我们就保存了一个.plan文件。PLAN文件是运行引擎用于执行网络的序列化数据。包含权重，网络中执行步骤以及用来决定如何绑定输入与输出缓存的网络信息。
    trt**.**utils**.**cwrite_engine_to_file('./model_.pb.plan',engine**.**serialize())

    使用TensorRT实现推理

    现在，让我们调用之前保存的plan文件，启用引擎，开始使用TensorRT实现推理
    engine **=** trt**.**utils**.**load_engine(G_LOGGER, './model_.pb.plan')
    引擎叫做engine，而引擎运行的上下文叫做context。engine和context在推理过程中都是必须的，这两者的关系如下：
    1. context = engine.create_execution_context()
    2. engine = context.get_engine()
    在运行前向运算前，我们还需要做一次确认。get_nb_bindings()是为了获取与这个engine相关的输入输出tensor的数量。对于本例中单输入输出的模型，tensor的数量是2。如果有多个输入输出，这个确认值就要相应的变化，比如3个输入，1个输出的模型，tensor的数量就是4。我们需要知道这个数量，是为了之后的显存分配做准备
    1. print(engine.get_nb_bindings())
    2. assert(engine.get_nb_bindings() == 2)
    现在准备好一张可以输入给模型的图像 img.jpg，并且转换成fp32
    1. img = cv2.imread(img.jpg)
    2. img = img.astype(np.float32)
    同时，创建一个array来“接住”输出数据。为什么说“接住”呢，因为之后你就会看到，引擎做前向推理计算的时候，是生成了一个数据流，这个数据流会写入output array中
    1. #create output array to receive data
    2. OUTPUT_SIZE = 10
    3. output = np.zeros(OUTPUT_SIZE , dtype = np.float32)
    我们需要为输入输出分配显存，并且绑定。 ```python

    使用PyCUDA申请GPU显存并在引擎中注册

    申请的大小是整个batchsize大小的输入以及期望的输出指针大小。

    d_input = cuda.mem_alloc(1 img.size img.dtype.itemsize) d_output = cuda.mem_alloc(1 output.size output.dtype.itemsize)

引擎需要绑定GPU显存的指针。PyCUDA通过分配成ints实现内存申请。

bindings = [int(d_input), int(d_output)]

  1. 现在，我们可以开始TensorRT上的推理计算了！
  2. ```python
  3. # 建立数据流
  4. stream = cuda.Stream()
  5. # 将输入传给cuda
  6. cuda.memcpy_htod_async(d_input, img, stream)
  7. # 执行前向推理计算
  8. context.enqueue(1, bindings, stream.handle, None)
  9. # 将预测结果传回
  10. cuda.memcpy_dtoh_async(output, d_output, stream)
  11. # 同步
  12. stream.synchronize()

这个时候，如果你将output打印出来，就会发现output数组中已经有值了，这就是TensorRT计算的结果。
如过你使用tensorflow的方法，对同一组输入数据做预测，看看计算的结果是否一致 ，因为精度的差异会有一些差异，但是大体上来说，使用tensorflow和TensorRT，会得到一致的结果。

特别注意！

TensorRT和Tensorflow的数据格式不一样，Tensorflow是NHWC格式，即channel_last，而TensorRT中是NCHW格式，即channel_first，比如一张RGB图像，在Tensorflow中表示为（224， 224， 3），在TensorRT中就是（3，224， 224）。所以使用TensorRT时，请一定确认图像的格式。

参考资料：https://zhuanlan.zhihu.com/p/64114667

第二种方式

https://blog.csdn.net/zong596568821xp/article/details/86077553
要将tensorflow的pb文件转化为uff格式的文件
python3.5 /usr/lib/python3.5/dist-packages/uff/bin/convert_to_uff.py --input_file models/lenet5.pb
需要注意的是：该转化过程只支持在x86平台上实现，如果想在TX2等嵌入式平台上使用tensorrt，需要先在x86平台上将pb模型转化为uff文件，然后再拷贝到TX2上使用

end