TensorRT部署YOLOv5(04)—构建TensorRT引擎

在上一篇文章中，已经对什么是TensorRT，使用TensorRT进行深度学习模型部署推理的完整流程进行了初步的介绍。本文将详细介绍TensorRT引擎的构建，包括tf2onnx的使用、trtexec命令行的使用和重要参数介绍、如何使用Python API进行TensorRT引擎构建等

Tensorflow模型转换为ONNX

使用以下命令安装tf2onnx

pip install tf2onnx -i https://pypi.tuna.tsinghua.edu.cn/simple

tf2onnx是以Python命令行模块的形式使用的，可在终端中输入以下内容，查看tf2onnx的帮助信息

python -m tf2onnx -h
optional arguments:
  -h, --help            show this help message and exit
  --input INPUT         input from graphdef
  --graphdef GRAPHDEF   input from graphdef
  --saved-model SAVED_MODEL
                        input from saved model
  --tag TAG             tag to use for saved_model
  --signature_def SIGNATURE_DEF
                        signature_def from saved_model to use
  --concrete_function CONCRETE_FUNCTION
                        For TF2.x saved_model, index of func signature in __call__ (--signature_def is ignored)
  --checkpoint CHECKPOINT
                        input from checkpoint
  --keras KERAS         input from keras model
  --tflite TFLITE       input from tflite model
  --tfjs TFJS           input from tfjs model
  --large_model         use the large model format (for models > 2GB)
  --output OUTPUT       output model file
  --inputs INPUTS       model input_names (optional for saved_model, keras, and tflite)
  --outputs OUTPUTS     model output_names (optional for saved_model, keras, and tflite)
  --ignore_default IGNORE_DEFAULT
                        comma-separated list of names of PlaceholderWithDefault ops to change into Placeholder ops
  --use_default USE_DEFAULT
                        comma-separated list of names of PlaceholderWithDefault ops to change into Identity ops using their default value
  --rename-inputs RENAME_INPUTS
                        input names to use in final model (optional)
  --rename-outputs RENAME_OUTPUTS
                        output names to use in final model (optional)
  --use-graph-names     (saved model only) skip renaming io using signature names
  --opset OPSET         opset version to use for onnx domain
  --dequantize          remove quantization from model. Only supported for tflite currently.
  --custom-ops CUSTOM_OPS
                        comma-separated map of custom ops to domains in format OpName:domain. Domain 'ai.onnx.converters.tensorflow' is used by default.
  --extra_opset EXTRA_OPSET
                        extra opset with format like domain:version, e.g. com.microsoft:1
  --load_op_libraries LOAD_OP_LIBRARIES
                        comma-separated list of tf op library paths to register before loading model
  --target {rs4,rs5,rs6,caffe2,tensorrt,nhwc}
                        target platform
  --continue_on_error   continue_on_error
  --verbose, -v         verbose output, option is additive
  --debug               debug mode
  --output_frozen_graph OUTPUT_FROZEN_GRAPH
                        output frozen tf graph to file
  --inputs-as-nchw INPUTS_AS_NCHW
                        transpose inputs as from nhwc to nchw
  --outputs-as-nchw OUTPUTS_AS_NCHW
                        transpose outputs as from nhwc to nchw

Usage Examples:

python -m tf2onnx.convert --saved-model saved_model_dir --output model.onnx
python -m tf2onnx.convert --input frozen_graph.pb  --inputs X:0 --outputs output:0 --output model.onnx
python -m tf2onnx.convert --checkpoint checkpoint.meta  --inputs X:0 --outputs output:0 --output model.onnx

For help and additional information see:
    https://github.com/onnx/tensorflow-onnx

tf2onnx可以接收很多种类型的深度学习模型格式，例如Tensorflow的.pb、checkpoint、.h5、tflite、tfjs等，也可以接收keras模型，在本实验中，我最开始使用.h5和checkpoint并未成功，原因未找到，.pb是可以正确转换的，后续没有再对该问题进行研究，一直使用.pb格式

将Tensorflow导出的model.pb模型转换为ONNX模型model.onnx的命令如下

python -m tf2onnx.convert --saved-model model.pb --output model.onnx \
--inputs-as-nchw input

使用--inputs-as-nchw input的原因是Tensorflow与ONNX使用的内存排布方式不同，Tensorflow中使用的是NHWC，而ONNX使用的是NCHW，因此添加该转换参数，将输入的内存排布修改未NCHW

tf2onnx的--opset可以手动指定opset version，因为ONNX和TensorRT的版本存在一些对应关系，如果ONNX模型转换TensorRT过程中存在问题，可能是opset version的问题，有可能可以通过手动指定opset version解决。本实验不存在这种问题，因此未手动指定opset，使用默认opset进行转换

Nvidia为适配ONNX推出了一个[ONNX GraphSurgeon](ONNX GraphSurgeon — ONNX GraphSurgeon 0.3.25 documentation (nvidia.com))工具，该工具可以对ONNX的模型计算图进行一些节点修改，可以通过该工具解决一些算子不匹配问题或者进行模型计算图优化。本实验未使用ONNX GraphSurgeon对模型进行修改，这部分的详细内容本专栏不涉及

BatchSize

这里的BatchSize指的是模型推理过程输入数据的BatchSize。TensorRT既支持固定的BatchSize，也支持动态BatchSize(Dynamic Shape)，固定的BatchSize能够让TensorRT在构建期更好的优化模型。Dynamic Shape能够让开发者在构建期不指定InputShape，而在运行期再通过API来指定。本实验由于是在JestonNano上进行YOLOv5模型推理，JestonNano本身性能算力有限，因此本实验中在ONNX模型中手动固定了BatchSize为1，以下为使用onnx修改BatchSize的代码

import onnx

onnx_model = onnx.load_model('model.onnx')

batch_size = 1
inputs = onnx_model.graph.input
for input in inputs:
    dim1 = input.type.tensor_type.shape.dim[0]
    dim1.dim_value = batch_size
model_name = 'model-bs1.onnx'
onnx.save_model(onnx_model, model_name)

ONNX模型转换为TensorRT引擎

ONNX模型转换为TensorRT引擎有两种方式

• trtexec命令行工具

• TensorRT API

我比较推荐使用trtexec命令行工具来进行转换，该工具不仅可以进行推理引擎的转换，还可以进行转换后的模型推理性能测试，而且转换过程只需要选取正确的参数即可，不需要额外编写代码

使用trtexec转换模型为TensorRT引擎

默认的trtexec命令并没有添加到环境变量，而是位于/usr/src/tensorrt/bin目录下，将model-bs1.onnx转换为TensorRT引擎的命令如下

/usr/src/tensorrt/bin/trtexec --onnx=model-bs1.onnx --saveEngine=model.trt --workspace=3200

以上命令执行后，trtexec开始工作，这需要一段时间，trtexec首先会将ONNX格式进行解析，并进行适用于在GPU上推理的计算图优化过程，然后是各层具体实现算子及参数的自动调优，整个过程都是自动进行的，用户所需要做的仅仅是输入以上命令，然后等待即可

有几点需要说明，首先是trtexec并不仅仅可以接收ONNX格式，也可以接收一些其他格式的模型，本文仅使用了ONNX格式，因此对于其他格式不在此介绍。--saveEngine后跟要输出的引擎文件名称，这里输入的"model.trt"只是一个文件名字，可以输入任何合法的名称。--workspace是为trtexec指定的转换过程可用内存空间，单位是MB，由于JestonNano开发板内存大小为4G，以上命令是按照板子实际可用内存配置的，用户需按照自己的实际情况进行设置，太小的数值可能导致构建过程失败

转换精度

转换过程非常重要的一点是转换精度，不同精度的引擎，引擎大小及内部参数占用内内存大小差异很大，由于TensorRT使用GPU进行推理，整个过程涉及到CPU和GPU之间的内存拷贝，较低精度的模型能够在很大程度上的降低内存拷贝和访存时间，对推理速度的影响很大，代价是牺牲一些精度。默认未指定转换精度的情况下，trtexec生成的推理引擎是FP32的，JestonNano上最低支持的精度是FP16，INT8精度在Nvidia NX等更高端的计算平台上才支持，因此INT8精度转换及后续处理不在本系列讨论范围内，需要注意的是INT8精度转换过程需要进行校准操作，需要提供一组校准数据集(unlabel)以及指定校准模式等操作，详细内容可以在TensorRT官方文档中找到。要将模型转换为FP16精度，可按照以下命令操作

/usr/src/tensorrt/bin/trtexec --onnx=model-bs1.onnx --saveEngine=xxx-fp16.trt --fp16 --workspace=3200

只需要增加--fp16选项即可。需要清楚一点，这里使用--fp16选项，trtexec转换后的模型，输入层和输出层依然是fp32，只有中间层是fp16，这一点在后续推理章节会详细介绍。如果要让输入层和输出层也变为fp16，需要增加--inputIOFormats和--outputIOFormats选项，如下

/usr/src/tensorrt/bin/trtexec --onnx=model-bs1.onnx --saveEngine=xxx-fp16.trt --explicitBatch 
--inputIOFormats=fp16:chw --outputIOFormats=fp16:chw --fp16 --workspace=3200

FP16比FP32转换所需要的时间更长

转换输出信息

trtexec转换过程末尾的输出信息非常重要，对引擎的推理性能进行了多轮推理测试，这对实际模型的部署过程性能评估具有指导意义

重点关注Throughput和Latency两个参数，这两个参数其实是互为倒数，例如图中所示，平均Latency为108.75ms，Throughput为9.18516，意味着不考虑前处理和后处理耗时，模型推理的极限帧率大概是9.18fps，实际进行推理时的推理延迟和这里的测试数值基本一致

使用TensorRT API转换模型为TensorRT引擎

使用TensorRT API的转换脚本我这里直接给出，具体内容直接看代码吧，就不详细介绍了

import os
import sys
import logging
import argparse
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit



logging.basicConfig(level=logging.INFO)
logging.getLogger("Engine").setLevel(logging.INFO)
logger = logging.getLogger('Engine')



class EngineCalibrator(trt.IInt8MinMaxCalibrator):

    def __init__(self, cache_file):
        super().__init__()
        self.cache_file = cache_file
        self.calib_batcher = None
        self.batch_allocation = None
        self.batch_generator = None
    
    def set_calib_batcher(self, calib_batcher: CalibBatcher):
        self.calib_batcher = calib_batcher
        size = int(np.dtype(self.image_batcher.dtype).itemsize * np.prod(self.image_batcher.shape))
        self.batch_allocation = cuda.mem_alloc(size)
        self.batch_generator = self.image_batcher.get_batch()
    
    def get_batch_size(self):
        if self.calib_batcher:
            return self.calib_batcher.batch_size
        return 1
    
    def get_batch(self, names):
        if not self.calib_batcher:
            return None
        try:
            batch, _, _ = next(self.batch_generator)
            logger.info("Calibrating image {} / {}".format(self.calib_batcher.image_index, 
                self.calib_batcher.num_images))
            cuda.memcpy_htod(self.batch_allocation, np.ascontiguousarray(batch))
            return [int(self.batch_allocation)]
        except StopIteration:
            logger.info("Finished calibration batches")
            return None
    
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                logger.info("Using calibration cache file: {}".format(self.cache_file))
                return f.read()
    
    def write_calibration_cache(self, cache):
        if self.cache_file is None:
            return
        with open(self.cache_file, "wb") as f:
             logger.info("Writing calibration cache data to: {}".format(self.cache_file))
             f.write(cache)

def main(args):

    logger = trt.Logger(trt.Logger.WARNING)
    trt.init_libnvinfer_plugins(logger, namespace="")
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreateFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    success = parser.parse_from_file(args.onnx)
    for idx in range(parser.num_errors):
        print(parser.get_error(idx))

    if not success:
        return

    inputs = [network.get_input(i) for i in range(network.num_inputs)]
    outputs = [network.get_output(i) for i in range(network.num_outputs)]

    logger.info('Network Description')
    for input in inputs:
        logger.info("Input '{}' with shape {} and dtype {}".format(input.name, input.shape, input.dtype))
    for output in outputs:
        logger.info("Output '{}' with shape {} and dtype {}".format(output.name, output.shape, output.dtype))

    config = builder.create_builder_config()
    memsize = args.mem * 1024 * 1024
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, memsize)

    if args.precision == 'fp16':
        if not builder.platform_has_fast_fp16:
            logger.warning("FP16 is not supported natively on this platform/device")
        config.set_flag(trt.BuilderFlag.FP16)
    elif args.precision == 'int8':
        if not builder.platform_has_fast_int8:
            logger.warning("INT8 is not supported natively on this platform/device")
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = EngineCalibrator(args.calib_cache)
        if args.calib_cache is None or not os.path.exists(args.calib_cache):
            calib_shape = [args.calib_batch_size] + list(inputs[0].shape[1:])
            calib_dtype = trt.nptype(inputs[0].dtype)
            config.int8_calibrator.set_image_batcher(
                CalibBatcher(args.calib_path, calib_shape, calib_dtype, 
                    args.calib_num, exact_batch=True)
            )

    serialized_engine = builder.builder_serialized_network(network, config)
    with open(args.output, 'wb') as f:
        logger.info("Serializing engine to file: {:}".format(args.engine))
        f.write(serialized_engine)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--onnx', required=True, help='input onnx file')
    parser.add_argument('--engine', default='engine.trt', help='output engine filename')
    parser.add_argument('--precision', default='fp32', choices=['fp32', 'fp16', 'int8'], 
        help='the precision of mode, fp32/fp16/int8')
    parser.add_argument('--mem', default=32, type=int, help='memory pool limit(MiB)')
    parser.add_argument('--calib_path', help='the int8 calibration images directory')
    parser.add_argument('--calib_cache', default='./calibration.cache', 
        help='the int8 calibration cache save path')
    parser.add_argument('--calib_num', default=5000, type=int, 
        help='the int calibration images number, default 5000')
    parser.add_argument('--calib_batch_size', default=8, type=int,
        help='the int calibration batch size, default 8')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    if args.precision == 'int8' and not args.calib_path:
        logger.error("int8 --calib_path is required")
        sys.exit(1)
    try:
        main(args)
    except SystemError:
        pass