这里探究了如何将基于TensorFlow2.x框架训练的模型转换成通用性较强的onnx静态图模型，以及推理效率较高的TensorRT模型。

每当碰到台风天气，或者平时下暴雨，现在的城市排水系统都有滞后性，如果碰上短时间强降雨就很容易导致路面积水，危害人民生命财产安全。因此需要设计算法可自动识别地面积水情况，从而通知相关部门快速处理。

对于上述应用场景，精度、性能以及模型部署的可迁移性都是在算法设计实现过程中需要重点考虑的因素，为了提高算法效率，同时增强模型的可移植性，本文探究了如何将基于TensorFlow2.x框架训练的模型转换成通用性较强的onnx静态图模型，以及推理效率较高的TensorRT模型。

TensorFlow是主流深度学习框架之一，在全球有着数十万名用户。由于TensorFlow1.x代码构建及接口问题饱含诟病，Google在2019年推出Tensorflow2避免TensorFlow 1.x 版本的诸多缺陷，Tensorflow2.x与Keras的代码构建方式类似，且包含了大量Keras接口，同时TensorFlow 2.x 也支持通过 tf.function 将动态图优先模式的代码转化为静态图模式，实现开发和运行效率双赢。

NVIDIA®TensorRT™是一个深度学习平台，用于优化神经网络模型，其核心库是使用C++去加速NVIDIA生产的GPU。一般来说从tensorflow2.x到tensorrt部署主要有两种方式：

  1. model.h5->model.pb->model.uff
  2. model.h5->model.pb-> model.onnx->model.engine

由于极市平台提供的demo通常支持通过.onnx转engine（demo地址：src/SampleDetector.cpp · cvmart/ev_sdk_demo4.0_vehicle_plate_detection - 码云 - 开源中国 (gitee.com)），因此本文主要介绍第二种部署方式。

运行环境：

TensorFlow 2.4
CUDA 11.1
CUDNN 8
TensorRT 8.2.1.8
tf2onnx 1.13.0
onnx 1.12.0

部署步骤：

通常来说，采用model.fit()接口生成的模型文件后缀为.h5，首先需要将该模型转换为静态图.pb格式，python代码如下：

from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
def h5_to_pb(model):
    model.load_weights(h5_save_path, by_name=True, skip_mismatch=True)
    print(tf.__version__)
    model.summary()
    full_model = tf.function(lambda Input: model(Input))
    full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
    # Get frozen ConcreteFunction
    frozen_func = convert_variables_to_constants_v2(full_model)
    frozen_func.graph.as_graph_def()

    layers = [op.name for op in frozen_func.graph.get_operations()]
    print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
    print(layer)
    print("-" * 50)
    print("Frozen model inputs: ")
    print(frozen_func.inputs)
    print("Frozen model outputs: ")
    print(frozen_func.outputs)
    # Save frozen graph from frozen ConcreteFunction to hard drive
    tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                      logdir=h5_save_paths,
                      name="model.pb",
                      as_text=False)
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.
  • 20.
  • 21.
  • 22.
  • 23.
  • 24.
  • 25.
  • 26.

其中model为加载到内存中的动态图模型。.pb模型转换成功后，通过tf2onnx库将其转换为onnx通用模型文件，代码如下：

python -m tf2onnx.convert --graphdef /project/train/models/model.pb --output /project/train/models/seg_model.onnx --inputs Input:0 --outputs Identity:0
  • 1.

需要注意输入和输出层名称（该名称在上一步转pb的过程中会打印出来）

极市部署demo中已经继承了从onnx到trt的过程，如果需要自行转换trt文件则在配置好TensorRT环境后，通过以下代码输出：

/project/train/TensorRT-8.2.1.8/bin/trtexec 
--onnx=/project/train/models/seg_model.onnx 
--saveEngine=/project/train/models/seg_model.trt --workspace=2048 --fp16
  • 1.
  • 2.
  • 3.

在trt部署过程中需要注意一点，由于pytorch默认输入为C×H×W，而TensorFlow默认输入维度顺序为H×W×C，所以在进行数据预处理时存在一些差异，在C++部署过程中采用以下代码进行数据预处理：

void normalize_(Mat &img)
{
   img.convertTo(img, CV_32F);
    int row = img.rows;
    int col = img.cols;
    std::vector<float> mean {0.406, 0.456, 0.485};
    std::vector<float> stds{ 0.225, 0.224, 0.229}; 
    for (int c = 0; c < 3; c++)
    {
        for (int i = 0; i < row; i++)
        {
            for (int j = 0; j < col; j++)
            {
                float pix =  img.ptr<float>(i)[j * 3 + c];
                img.ptr<float>(i)[j * 3 + c] = (pix / 255.0 - mean[c]) / stds[c];
            }
        }
    }
}
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.
  • 7.
  • 8.
  • 9.
  • 10.
  • 11.
  • 12.
  • 13.
  • 14.
  • 15.
  • 16.
  • 17.
  • 18.
  • 19.

通过上述步骤即完成了TensorFlow2.x从.h5到TensorRT的.engine转换。使用TensorRT进行FP16推理通常能够比TensorFlow推理.pb速度提升三倍所有，比推理.h5速度提升8倍左右。

好简单就完事了哦